crypto: more type annotations

This commit is contained in:
SomberNight 2018-11-10 13:30:34 +01:00
parent a6a003a345
commit aceb022f9d
No known key found for this signature in database
GPG key ID: B33B5F232C6271E9

View file

@ -44,13 +44,13 @@ class InvalidPadding(Exception):
pass pass
def append_PKCS7_padding(data): def append_PKCS7_padding(data: bytes) -> bytes:
assert_bytes(data) assert_bytes(data)
padlen = 16 - (len(data) % 16) padlen = 16 - (len(data) % 16)
return data + bytes([padlen]) * padlen return data + bytes([padlen]) * padlen
def strip_PKCS7_padding(data): def strip_PKCS7_padding(data: bytes) -> bytes:
assert_bytes(data) assert_bytes(data)
if len(data) % 16 != 0 or len(data) == 0: if len(data) % 16 != 0 or len(data) == 0:
raise InvalidPadding("invalid length") raise InvalidPadding("invalid length")
@ -63,7 +63,7 @@ def strip_PKCS7_padding(data):
return data[0:-padlen] return data[0:-padlen]
def aes_encrypt_with_iv(key, iv, data): def aes_encrypt_with_iv(key: bytes, iv: bytes, data: bytes) -> bytes:
assert_bytes(key, iv, data) assert_bytes(key, iv, data)
data = append_PKCS7_padding(data) data = append_PKCS7_padding(data)
if AES: if AES:
@ -75,7 +75,7 @@ def aes_encrypt_with_iv(key, iv, data):
return e return e
def aes_decrypt_with_iv(key, iv, data): def aes_decrypt_with_iv(key: bytes, iv: bytes, data: bytes) -> bytes:
assert_bytes(key, iv, data) assert_bytes(key, iv, data)
if AES: if AES:
cipher = AES.new(key, AES.MODE_CBC, iv) cipher = AES.new(key, AES.MODE_CBC, iv)
@ -90,36 +90,38 @@ def aes_decrypt_with_iv(key, iv, data):
raise InvalidPassword() raise InvalidPassword()
def EncodeAES(secret, s): def EncodeAES(secret: bytes, msg: bytes) -> bytes:
assert_bytes(s) """Returns base64 encoded ciphertext."""
assert_bytes(msg)
iv = bytes(os.urandom(16)) iv = bytes(os.urandom(16))
ct = aes_encrypt_with_iv(secret, iv, s) ct = aes_encrypt_with_iv(secret, iv, msg)
e = iv + ct e = iv + ct
return base64.b64encode(e) return base64.b64encode(e)
def DecodeAES(secret, e):
e = bytes(base64.b64decode(e)) def DecodeAES(secret: bytes, ciphertext_b64: Union[bytes, str]) -> bytes:
e = bytes(base64.b64decode(ciphertext_b64))
iv, e = e[:16], e[16:] iv, e = e[:16], e[16:]
s = aes_decrypt_with_iv(secret, iv, e) s = aes_decrypt_with_iv(secret, iv, e)
return s return s
def pw_encode(s, password):
if password:
secret = sha256d(password)
return EncodeAES(secret, to_bytes(s, "utf8")).decode('utf8')
else:
return s
def pw_decode(s, password): def pw_encode(data: str, password: Union[bytes, str]) -> str:
if password is not None: if not password:
secret = sha256d(password) return data
try: secret = sha256d(password)
d = to_string(DecodeAES(secret, s), "utf8") return EncodeAES(secret, to_bytes(data, "utf8")).decode('utf8')
except Exception:
raise InvalidPassword()
return d def pw_decode(data: str, password: Union[bytes, str]) -> str:
else: if password is None:
return s return data
secret = sha256d(password)
try:
d = to_string(DecodeAES(secret, data), "utf8")
except Exception:
raise InvalidPassword()
return d
def sha256(x: Union[bytes, str]) -> bytes: def sha256(x: Union[bytes, str]) -> bytes: