better exception handling

Author: JarbasAl

hivemind_bus_client/encryption.py

@@ -1,16 +1,30 @@
 import enum
 import json
 from binascii import hexlify, unhexlify
-from typing import Union, Optional, Dict, Any
+from typing import Union, Optional, Dict, Any, Tuple
 
 import pybase64
 from Cryptodome.Cipher import AES, ChaCha20_Poly1305
-
 from cpuinfo import get_cpu_info
-from hivemind_bus_client.exceptions import EncryptionKeyError, DecryptionKeyError, InvalidCipher
+
+from hivemind_bus_client.exceptions import EncryptionKeyError, DecryptionKeyError, InvalidCipher, InvalidKeySize
+
+# Cipher-specific constants
+AES_GCM_KEY_SIZE = 16
+AES_GCM_NONCE_SIZE = 16
+AES_GCM_TAG_SIZE = 16
+CHACHA20_KEY_SIZE = 32
+CHACHA20_NONCE_SIZE = 12
+CHACHA20_TAG_SIZE = 16
 
 
 def cpu_supports_AES() -> bool:
+    """
+    Check if the CPU supports AES encryption.
+
+    Returns:
+        bool: True if AES is supported by the CPU, False otherwise.
+    """
     return "aes" in get_cpu_info()["flags"]
 
 
@@ -66,13 +80,22 @@ def encrypt_as_json(key: Union[str, bytes], data: Union[str, Dict[str, Any]],
 
     bcipher = BinaryCiphers.BINARY_AES_GCM_128 if cipher in aes_ciphers else BinaryCiphers.BINARY_CHACHA20_POLY1305
 
-    ciphertext = encrypt_bin(key, data, cipher=bcipher)
+    try:
+        ciphertext = encrypt_bin(key, data, cipher=bcipher)
+    except InvalidKeySize as e:
+        raise e
+    except Exception as e:
+        raise EncryptionKeyError from e
 
     # extract nonce/tag depending on cipher, sizes are different
     if cipher in aes_ciphers:
-        nonce, ciphertext, tag = ciphertext[:16], ciphertext[16:-16], ciphertext[-16:]
+        nonce, ciphertext, tag = (ciphertext[:AES_GCM_NONCE_SIZE],
+                                  ciphertext[AES_GCM_NONCE_SIZE:-AES_GCM_TAG_SIZE],
+                                  ciphertext[-AES_GCM_TAG_SIZE:])
     else:
-        nonce, ciphertext, tag = ciphertext[:12], ciphertext[12:-16], ciphertext[-16:]
+        nonce, ciphertext, tag = (ciphertext[:CHACHA20_NONCE_SIZE],
+                                  ciphertext[CHACHA20_NONCE_SIZE:-CHACHA20_TAG_SIZE],
+                                  ciphertext[-CHACHA20_TAG_SIZE:])
 
     encoder = pybase64.b64encode if cipher in b64_ciphers else hexlify
 
@@ -90,7 +113,7 @@ def decrypt_from_json(key: Union[str, bytes], data: Union[str, bytes], cipher: J
     Args:
         key (Union[str, bytes]): The decryption key, up to 16 bytes. Longer keys will be truncated.
         data (Union[str, bytes]): The encrypted data as a JSON string or bytes.
-        cipher (Optional[JsonCiphers]): The cipher used for encryption. If None, it is auto-detected.
+        cipher (JsonCiphers): The cipher used for encryption.
 
     Returns:
         str: The decrypted plaintext data.
@@ -99,7 +122,6 @@ def decrypt_from_json(key: Union[str, bytes], data: Union[str, bytes], cipher: J
         InvalidCipher: If an unsupported cipher is provided.
         DecryptionKeyError: If decryption fails due to an invalid key or corrupted data.
     """
-
     aes_ciphers = {JsonCiphers.JSON_B64_AES_GCM_128, JsonCiphers.JSON_HEX_AES_GCM_128}
     b64_ciphers = {JsonCiphers.JSON_B64_AES_GCM_128, JsonCiphers.JSON_B64_CHACHA20_POLY1305}
 
@@ -111,7 +133,10 @@ def decrypt_from_json(key: Union[str, bytes], data: Union[str, bytes], cipher: J
 
     ciphertext = decoder(data["ciphertext"])
     if "tag" not in data:  # web crypto compatibility
-        ciphertext, tag = ciphertext[:-16], ciphertext[-16:]
+        if bcipher == BinaryCiphers.BINARY_AES_GCM_128:
+            ciphertext, tag = ciphertext[:-AES_GCM_TAG_SIZE], ciphertext[-AES_GCM_TAG_SIZE:]
+        else:
+            ciphertext, tag = ciphertext[:-CHACHA20_TAG_SIZE], ciphertext[-CHACHA20_TAG_SIZE:]
     else:
         tag = decoder(data["tag"])
     nonce = decoder(data["nonce"])
@@ -120,7 +145,9 @@ def decrypt_from_json(key: Union[str, bytes], data: Union[str, bytes], cipher: J
     try:
         plaintext = decryptor(key, ciphertext, tag, nonce)
         return plaintext.decode("utf-8")
-    except ValueError as e:
+    except InvalidKeySize as e:
+        raise e
+    except Exception as e:
         raise DecryptionKeyError from e
 
 
@@ -148,6 +175,8 @@ def encrypt_bin(key: Union[str, bytes], data: Union[str, bytes], cipher: BinaryC
     encryptor = encrypt_AES_GCM_128 if cipher == BinaryCiphers.BINARY_AES_GCM_128 else encrypt_ChaCha20_Poly1305
     try:
         ciphertext, tag, nonce = encryptor(key, data)
+    except InvalidKeySize as e:
+        raise e
     except Exception as e:
         raise EncryptionKeyError from e
 
@@ -175,14 +204,20 @@ def decrypt_bin(key: Union[str, bytes], ciphertext: bytes, cipher: BinaryCiphers
 
     # extract nonce/tag depending on cipher, sizes are different
     if cipher == BinaryCiphers.BINARY_AES_GCM_128:
-        nonce, ciphertext, tag = ciphertext[:16], ciphertext[16:-16], ciphertext[-16:]
+        nonce, ciphertext, tag = (ciphertext[:AES_GCM_NONCE_SIZE],
+                                  ciphertext[AES_GCM_NONCE_SIZE:-AES_GCM_TAG_SIZE],
+                                  ciphertext[-AES_GCM_TAG_SIZE:])
     else:
-        nonce, ciphertext, tag = ciphertext[:12], ciphertext[12:-16], ciphertext[-16:]
+        nonce, ciphertext, tag = (ciphertext[:CHACHA20_NONCE_SIZE],
+                                  ciphertext[CHACHA20_NONCE_SIZE:-CHACHA20_TAG_SIZE],
+                                  ciphertext[-CHACHA20_TAG_SIZE:])
 
     decryptor = decrypt_AES_GCM_128 if cipher == BinaryCiphers.BINARY_AES_GCM_128 else decrypt_ChaCha20_Poly1305
     try:
         return decryptor(key, ciphertext, tag, nonce)
-    except ValueError as e:
+    except InvalidKeySize as e:
+        raise e
+    except Exception as e:
         raise DecryptionKeyError from e
 
 
@@ -205,8 +240,8 @@ def encrypt_AES_GCM_128(key: Union[str, bytes], text: Union[str, bytes],
         text = bytes(text, encoding="utf-8")
     if not isinstance(key, bytes):
         key = bytes(key, encoding="utf-8")
-    if len(key) != 16:  # AES-128 uses 128 bit/16 byte keys
-        raise ValueError("AES-GCM-128 requires a 16-byte key")
+    if len(key) != AES_GCM_KEY_SIZE:  # AES-128 uses 128 bit/16 byte keys
+        raise InvalidKeySize("AES-GCM-128 requires a 16-byte key")
     cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
     ciphertext, tag = cipher.encrypt_and_digest(text)
     return ciphertext, tag, cipher.nonce
@@ -226,10 +261,13 @@ def decrypt_AES_GCM_128(key: Union[str, bytes], ciphertext: bytes, tag: bytes, n
         str: The decrypted plaintext.
 
     Raises:
+        InvalidKeySize: If key size is not valid
         ValueError: If decryption or authentication fails.
     """
-    if not isinstance(key, bytes):
-        key = bytes(key, encoding="utf-8")
+    if isinstance(key, str):
+        key = key.encode("utf-8")
+    if len(key) != AES_GCM_KEY_SIZE:  # AES-128 uses 128 bit/16 byte keys
+        raise InvalidKeySize("AES-GCM-128 requires a 16-byte key")
     cipher = AES.new(key, AES.MODE_GCM, nonce)
     return cipher.decrypt_and_verify(ciphertext, tag)
 
@@ -248,13 +286,16 @@ def encrypt_ChaCha20_Poly1305(key: Union[str, bytes],
     Returns:
         tuple[bytes, bytes, bytes]: The ciphertext, authentication tag, and nonce.
     """
-    if not isinstance(text, bytes):
-        text = bytes(text, encoding="utf-8")
-    if not isinstance(key, bytes):
-        key = bytes(key, encoding="utf-8")
-    assert len(key) == 32  # ChaCha20 uses 256 bit/32 byte keys
+    if isinstance(text, str):
+        text = text.encode("utf-8")
+    if isinstance(key, str):
+        key = key.encode("utf-8")
+
+    if len(key) != CHACHA20_KEY_SIZE:  # ChaCha20 uses 256 bit/32 byte keys
+        raise InvalidKeySize("CHACHA20-POLY1305 requires a 32-byte key")
     if nonce:
-        assert len(nonce) == 12  # 92bits/12bytes bytes per RFC7539
+        if len(nonce) != CHACHA20_NONCE_SIZE:  # 92bits/12bytes per RFC7539
+            raise InvalidKeySize("CHACHA20-POLY1305 requires a 12-byte nonce per RFC7539")
     cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce)
     ciphertext, tag = cipher.encrypt_and_digest(text)
     return ciphertext, tag, cipher.nonce
@@ -277,14 +318,17 @@ def decrypt_ChaCha20_Poly1305(key: Union[str, bytes],
         str: The decrypted plaintext.
 
     Raises:
+        InvalidKeySize:
         ValueError: If decryption or authentication fails.
     """
-    if not isinstance(key, bytes):
-        key = bytes(key, encoding="utf-8")
+    if isinstance(key, str):
+        key = key.encode("utf-8")
 
-    assert len(key) == 32  # ChaCha20 uses 256 bit/32 byte keys
+    if len(key) != CHACHA20_KEY_SIZE:  # ChaCha20 uses 256 bit/32 byte keys
+        raise InvalidKeySize("CHACHA20-POLY1305 requires a 32-byte key")
     if nonce:
-        assert len(nonce) == 12  # 92bits/12bytes per RFC7539
+        if len(nonce) != CHACHA20_NONCE_SIZE:  # 92bits/12bytes per RFC7539
+            raise InvalidKeySize("CHACHA20-POLY1305 requires a 12-byte nonce per RFC7539")
     cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce)
     return cipher.decrypt_and_verify(ciphertext, tag)
 
@@ -294,7 +338,7 @@ def decrypt_ChaCha20_Poly1305(key: Union[str, bytes],
 
     print("JSON-B64-AES-GCM-128" == JsonCiphers.JSON_B64_AES_GCM_128)
 
-    key = get_random_bytes(32)
+    key = get_random_bytes(CHACHA20_KEY_SIZE)
     plaintext = b'Attack at dawn'
     ciphertext, tag, nonce = encrypt_ChaCha20_Poly1305(key, plaintext)
     recovered = decrypt_ChaCha20_Poly1305(key, ciphertext, tag, nonce)

hivemind_bus_client/exceptions.py

@@ -15,6 +15,10 @@ class InvalidCipher(HiveMindException):
     """unknown encryption scheme requested"""
 
 
+class InvalidKeySize(HiveMindException):
+    """ Encryption Key size does not obey specification"""
+
+
 class WrongEncryptionKey(HiveMindException):
     """ Wrong Encryption Key"""