Separate AES cipher

Signed-off-by: Ivan Kanakarakis <ivan.kanak@gmail.com>

Author: c00kiemon5ter

src/pyop/crypto.py

@@ -0,0 +1,74 @@
+import base64
+import hashlib
+
+from Cryptodome import Random
+from Cryptodome.Cipher import AES
+
+
+class _AESCipher(object):
+    """
+    This class will perform AES encryption/decryption with a keylength of 256.
+
+    @see: http://stackoverflow.com/questions/12524994/encrypt-decrypt-using-pycrypto-aes-256
+    """
+
+    def __init__(self, key):
+        """
+        Constructor
+
+        :type key: str
+
+        :param key: The key used for encryption and decryption. The longer key the better.
+        """
+        self.bs = 32
+        self.key = hashlib.sha256(key.encode()).digest()
+
+    def encrypt(self, raw):
+        """
+        Encryptes the parameter raw.
+
+        :type raw: bytes
+        :rtype: str
+
+        :param: bytes to be encrypted.
+
+        :return: A base 64 encoded string.
+        """
+        raw = self._pad(raw)
+        iv = Random.new().read(AES.block_size)
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return base64.urlsafe_b64encode(iv + cipher.encrypt(raw))
+
+    def decrypt(self, enc):
+        """
+        Decryptes the parameter enc.
+
+        :type enc: bytes
+        :rtype: bytes
+
+        :param: The value to be decrypted.
+        :return: The decrypted value.
+        """
+        enc = base64.urlsafe_b64decode(enc)
+        iv = enc[:AES.block_size]
+        cipher = AES.new(self.key, AES.MODE_CBC, iv)
+        return self._unpad(cipher.decrypt(enc[AES.block_size:]))
+
+    def _pad(self, b):
+        """
+        Will padd the param to be of the correct length for the encryption alg.
+
+        :type b: bytes
+        :rtype: bytes
+        """
+        return b + (self.bs - len(b) % self.bs) * chr(self.bs - len(b) % self.bs).encode("UTF-8")
+
+    @staticmethod
+    def _unpad(b):
+        """
+        Removes the padding performed by the method _pad.
+
+        :type b: bytes
+        :rtype: bytes
+        """
+        return b[:-ord(b[len(b) - 1:])]

src/pyop/storage.py

@@ -1,15 +1,15 @@
 # -*- coding: utf-8 -*-
 
 from abc import ABC, abstractmethod
-from Cryptodome import Random
-from Cryptodome.Cipher import AES
 import copy
 import json
-import base64
-import hashlib
 import logging
 from datetime import datetime
-from urllib.parse import urlparse, parse_qs
+from urllib.parse import urlparse
+from urllib.parse import parse_qs
+
+from .crypto import _AESCipher
+
 
 logger = logging.getLogger(__name__)
 
@@ -241,7 +241,7 @@ def __init__(self, collection, encryption_key, alg=None):
         if not alg or alg.lower() == "aes256":
             self.cipher = _AESCipher(encryption_key)
         else:
-            return ValueError(f"Invalid encryption algorithm: {alg}")
+            raise ValueError(f"Invalid encryption algorithm: {alg}")
 
     def __setitem__(self, key, value):
         pass
@@ -282,75 +282,6 @@ def _unpack(self, value):
         return unpacked_val
 
 
-class _AESCipher(object):
-    """
-    This class will perform AES encryption/decryption with a keylength of 256.
-
-    @see: http://stackoverflow.com/questions/12524994/encrypt-decrypt-using-pycrypto-aes-256
-    """
-
-    def __init__(self, key):
-        """
-        Constructor
-
-        :type key: str
-
-        :param key: The key used for encryption and decryption. The longer key the better.
-        """
-        self.bs = 32
-        self.key = hashlib.sha256(key.encode()).digest()
-
-    def encrypt(self, raw):
-        """
-        Encryptes the parameter raw.
-
-        :type raw: bytes
-        :rtype: str
-
-        :param: bytes to be encrypted.
-
-        :return: A base 64 encoded string.
-        """
-        raw = self._pad(raw)
-        iv = Random.new().read(AES.block_size)
-        cipher = AES.new(self.key, AES.MODE_CBC, iv)
-        return base64.urlsafe_b64encode(iv + cipher.encrypt(raw))
-
-    def decrypt(self, enc):
-        """
-        Decryptes the parameter enc.
-
-        :type enc: bytes
-        :rtype: bytes
-
-        :param: The value to be decrypted.
-        :return: The decrypted value.
-        """
-        enc = base64.urlsafe_b64decode(enc)
-        iv = enc[:AES.block_size]
-        cipher = AES.new(self.key, AES.MODE_CBC, iv)
-        return self._unpad(cipher.decrypt(enc[AES.block_size:]))
-
-    def _pad(self, b):
-        """
-        Will padd the param to be of the correct length for the encryption alg.
-
-        :type b: bytes
-        :rtype: bytes
-        """
-        return b + (self.bs - len(b) % self.bs) * chr(self.bs - len(b) % self.bs).encode("UTF-8")
-
-    @staticmethod
-    def _unpad(b):
-        """
-        Removes the padding performed by the method _pad.
-
-        :type b: bytes
-        :rtype: bytes
-        """
-        return b[:-ord(b[len(b) - 1:])]
-
-
 class MongoDB(object):
     """Simple wrapper to get pymongo real objects from the settings uri"""