Merge pull request #515 from c00kiemon5ter/fix-example-dependencies

Use cryptography in place of Cryptodome for aes

Author: c00kiemon5ter

src/saml2/aes.py

@@ -1,64 +1,70 @@
-#!/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
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher
+from cryptography.hazmat.primitives.ciphers import algorithms
+from cryptography.hazmat.primitives.ciphers import modes
 
-__author__ = 'rolandh'
 
 POSTFIX_MODE = {
-    "cbc": AES.MODE_CBC,
-    "cfb": AES.MODE_CFB,
-    "ecb": AES.MODE_CFB,
+    'cbc': modes.CBC,
+    'cfb': modes.CFB,
+    'ecb': modes.ECB,
 }
 
-BLOCK_SIZE = 16
+AES_BLOCK_SIZE = int(algorithms.AES.block_size / 8)
 
 
 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.lower().split('_')
+        bits = int(bits)
 
         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 = os.urandom(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: {}'.format(len(iv)))
+
+        if bits not in algorithms.AES.key_sizes:
+            raise Exception('Unsupported key length: {}'.format(bits))
+
+        if len(self.key) != bits / 8:
+            raise Exception('Wrong Key length: {}'.format(len(self.key)))
 
         try:
-            return AES.new(self.key, POSTFIX_MODE[cmode], iv), iv
+            mode = POSTFIX_MODE[cmode]
         except KeyError:
-            raise Exception("Unsupported chaining mode")
+            raise Exception('Unsupported chaining mode: {}'.format(cmode))
+
+        cipher = Cipher(
+                algorithms.AES(self.key),
+                mode(iv),
+                backend=default_backend())
 
+        return cipher, iv
 
-    def encrypt(self, msg, iv=None, alg="aes_128_cbc", padding="PKCS#7",
-                b64enc=True, block_size=BLOCK_SIZE):
+    def encrypt(self, msg, iv=None, alg='aes_128_cbc', padding='PKCS#7',
+                b64enc=True, block_size=AES_BLOCK_SIZE):
         """
         :param key: The encryption key
         :param iv: init vector
@@ -69,57 +75,71 @@ 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
+            c = chr(plen).encode()
+            msg += c * plen
 
         cipher, iv = self.build_cipher(iv, alg)
-        cmsg = iv + cipher.encrypt(msg)
+        encryptor = cipher.encryptor()
+        cmsg = iv + encryptor.update(msg) + encryptor.finalize()
+
         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]
+        _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"]:
-            res = res[:-ord(res[-1])]
+        decryptor = cipher.decryptor()
+        res = decryptor.update(data)[AES_BLOCK_SIZE:] + decryptor.finalize()
+        if padding in ['PKCS#5', 'PKCS#7']:
+            idx = bytearray(res)[-1]
+            res = res[:-idx]
         return res
 
-if __name__ == "__main__":
-    key_ = "1234523451234545"  # 16 byte key
+
+def run_test():
+    key = b'1234523451234545'  # 16 byte key
+    iv = os.urandom(AES_BLOCK_SIZE)
     # Iff padded, the message doesn't have to be multiple of 16 in length
-    msg_ = "ToBeOrNotTobe W.S."
-    aes = AESCipher(key_)
-    iv_ = os.urandom(16)
-    encrypted_msg = aes.encrypt(key_, msg_, iv_)
-    txt = aes.decrypt(key_, encrypted_msg, iv_)
-    assert txt == msg_
-
-    encrypted_msg = aes.encrypt(key_, msg_, 0)
-    txt = aes.decrypt(key_, encrypted_msg, 0)
-    assert txt == msg_
+    original_msg = b'ToBeOrNotTobe W.S.'
+    aes = AESCipher(key)
+
+    encrypted_msg = aes.encrypt(original_msg, iv)
+    decrypted_msg = aes.decrypt(encrypted_msg, iv)
+    assert decrypted_msg == original_msg
+
+    encrypted_msg = aes.encrypt(original_msg)
+    decrypted_msg = aes.decrypt(encrypted_msg)
+    assert decrypted_msg == original_msg
+
+    aes = AESCipher(key, iv)
+    encrypted_msg = aes.encrypt(original_msg)
+    decrypted_msg = aes.decrypt(encrypted_msg)
+    assert decrypted_msg == original_msg
+
+
+if __name__ == '__main__':
+    run_test()

src/saml2/authn.py

@@ -120,7 +120,7 @@ def __init__(self, srv, mako_template, template_lookup, pwd, return_to):
         self.return_to = return_to
         self.active = {}
         self.query_param = "upm_answer"
-        self.aes = AESCipher(self.srv.symkey, srv.iv)
+        self.aes = AESCipher(self.srv.symkey.encode(), srv.iv)
 
     def __call__(self, cookie=None, policy_url=None, logo_url=None,
                  query="", **kwargs):
@@ -171,7 +171,8 @@ def verify(self, request, **kwargs):
         try:
             self._verify(_dict["password"][0], _dict["login"][0])
             timestamp = str(int(time.mktime(time.gmtime())))
-            info = self.aes.encrypt("::".join([_dict["login"][0], timestamp]))
+            msg = "::".join([_dict["login"][0], timestamp])
+            info = self.aes.encrypt(msg.encode())
             self.active[info] = timestamp
             cookie = make_cookie(self.cookie_name, info, self.srv.seed)
             return_to = create_return_url(self.return_to, _dict["query"][0],
@@ -191,7 +192,8 @@ def authenticated_as(self, cookie=None, **kwargs):
                 info, timestamp = parse_cookie(self.cookie_name,
                                                self.srv.seed, cookie)
                 if self.active[info] == timestamp:
-                    uid, _ts = self.aes.decrypt(info).split("::")
+                    msg = self.aes.decrypt(info).decode()
+                    uid, _ts = msg.split("::")
                     if timestamp == _ts:
                         return {"uid": uid}
             except Exception: