Cleanup aes.py

Signed-off-by: Ivan Kanakarakis <[email protected]>

Author: c00kiemon5ter

src/saml2/aes.py

@@ -1,63 +1,58 @@
-#!/usr/bin/env python
 import os
-from base64 import b64encode
 from base64 import b64decode
+from base64 import b64encode
 
 from Cryptodome import Random
 from Cryptodome.Cipher import AES
 
-__author__ = 'rolandh'
 
 POSTFIX_MODE = {
-    "cbc": AES.MODE_CBC,
-    "cfb": AES.MODE_CFB,
-    "ecb": AES.MODE_CFB,
+    'cbc': AES.MODE_CBC,
+    'cfb': AES.MODE_CFB,
+    'ecb': AES.MODE_CFB,
 }
 
 BLOCK_SIZE = 16
 
 
 class AESCipher(object):
-    def __init__(self, key, iv=""):
+    def __init__(self, key, iv=None):
         """
-
         :param key: The encryption key
         :param iv: Init vector
         :return: AESCipher instance
         """
         self.key = key
         self.iv = iv
 
-    def build_cipher(self, iv="", alg="aes_128_cbc"):
+    def build_cipher(self, iv=None, alg='aes_128_cbc'):
         """
         :param iv: init vector
         :param alg: cipher algorithm
         :return: A Cipher instance
         """
-        typ, bits, cmode = alg.split("_")
+        typ, bits, cmode = alg.split('_')
 
         if not iv:
-            if self.iv:
-                iv = self.iv
-            else:
-                iv = Random.new().read(AES.block_size)
-        else:
-            assert len(iv) == AES.block_size
+            iv = self.iv if self.iv else Random.new().read(AES.block_size)
 
-        if bits not in ["128", "192", "256"]:
-            raise Exception("Unsupported key length")
-        try:
-            assert len(self.key) == int(bits) >> 3
-        except AssertionError:
-            raise Exception("Wrong Key length")
+        if len(iv) != AES.block_size:
+            raise Exception('Wrong iv size')
+
+        if bits not in ['128', '192', '256']:
+            raise Exception('Unsupported key length')
+
+        if len(self.key) != int(bits) >> 3:
+            raise Exception('Wrong Key length')
 
         try:
-            return AES.new(self.key, POSTFIX_MODE[cmode], iv), iv
+            result = AES.new(self.key, POSTFIX_MODE[cmode], iv)
         except KeyError:
-            raise Exception("Unsupported chaining mode")
-
+            raise Exception('Unsupported chaining mode')
+        else:
+            return result, iv
 
-    def encrypt(self, msg, iv=None, alg="aes_128_cbc", padding="PKCS#7",
+    def encrypt(self, msg, iv=None, alg='aes_128_cbc', padding='PKCS#7',
                 b64enc=True, block_size=BLOCK_SIZE):
         """
         :param key: The encryption key
@@ -69,51 +64,52 @@ def encrypt(self, msg, iv=None, alg="aes_128_cbc", padding="PKCS#7",
         :return: The encrypted message
         """
 
-        if padding == "PKCS#7":
+        if padding == 'PKCS#7':
             _block_size = block_size
-        elif padding == "PKCS#5":
+        elif padding == 'PKCS#5':
             _block_size = 8
         else:
             _block_size = 0
 
         if _block_size:
             plen = _block_size - (len(msg) % _block_size)
             c = chr(plen)
-            msg += c*plen
+            msg += c * plen
 
         cipher, iv = self.build_cipher(iv, alg)
         cmsg = iv + cipher.encrypt(msg)
+
         if b64enc:
-            return b64encode(cmsg)
+            enc_msg = b64encode(cmsg)
         else:
-            return cmsg
+            enc_msg = cmsg
 
+        return enc_msg
 
-    def decrypt(self, msg, iv=None, alg="aes_128_cbc", padding="PKCS#7", b64dec=True):
+    def decrypt(self, msg, iv=None, alg='aes_128_cbc', padding='PKCS#7',
+                b64dec=True):
         """
         :param key: The encryption key
         :param iv: init vector
         :param msg: Base64 encoded message to be decrypted
         :return: The decrypted message
         """
-        if b64dec:
-            data = b64decode(msg)
-        else:
-            data = msg
+        data = b64decode(msg) if b64dec else msg
 
         _iv = data[:AES.block_size]
         if iv:
             assert iv == _iv
         cipher, iv = self.build_cipher(iv, alg=alg)
         res = cipher.decrypt(data)[AES.block_size:]
-        if padding in ["PKCS#5", "PKCS#7"]:
+        if padding in ['PKCS#5', 'PKCS#7']:
             res = res[:-ord(res[-1])]
         return res
 
-if __name__ == "__main__":
-    key_ = "1234523451234545"  # 16 byte key
+
+if __name__ == '__main__':
+    key_ = '1234523451234545'  # 16 byte key
     # Iff padded, the message doesn't have to be multiple of 16 in length
-    msg_ = "ToBeOrNotTobe W.S."
+    msg_ = 'ToBeOrNotTobe W.S.'
     aes = AESCipher(key_)
     iv_ = os.urandom(16)
     encrypted_msg = aes.encrypt(key_, msg_, iv_)