Enhanced AES test

Author: miloserdow

Diff for: src/crypto.py

@@ -52,7 +52,7 @@ def decrypt(self, data: str) -> str:
 class AESCipher(SymmetricAlgorithm):
     # Sets the key and block size
     def __init__(self, key):
-        self.key = hashlib.sha256(key.encode()).digest()
+        self._key = hashlib.sha256(key.encode()).digest() if isinstance(key, str) else None
         self.bs = 32
 
     # Pad data up to block size
@@ -63,6 +63,9 @@ def unpad(self, s):
         return s[:-ord(s[len(s) - 1:])]
 
     def encrypt(self, data: str):
+        if not isinstance(self.key, str):
+            raise ValueError('Invalid key')
+
         data = self.pad(data)
         iv = Random.new().read(AES.block_size)
         aes = AES.new(self.key, AES.MODE_CBC, iv)
@@ -71,6 +74,9 @@ def encrypt(self, data: str):
         return base64.b64encode(C).decode('utf-8')
 
     def decrypt(self, data: str) -> str:
+        if not isinstance(self.key, str):
+            raise ValueError('Invalid key')
+        
         data = base64.b64decode(data)
         iv = data[:AES.block_size]
         aes = AES.new(self.key, AES.MODE_CBC, iv)

Diff for: tests/test_aes.py

@@ -9,8 +9,9 @@ def test_cipher(self):
 
         message = ''.join(random.choice(string.ascii_uppercase + string.digits + 
 string.ascii_lowercase) for _ in range(100))
-        key = random.randint(0,32)
-
+        key = ''.join(random.choice(string.ascii_uppercase + string.digits + 
+string.ascii_lowercase) for _ in range(32))
+        
         # check cipher without key
         try:
             C = cipher.encrypt(message)