Add tests for AES cipher

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

Author: c00kiemon5ter

src/saml2/aes.py

@@ -103,27 +103,3 @@ def decrypt(self, msg, alg='aes_128_cbc', padding='PKCS#7', b64dec=True):
             idx = bytearray(res)[-1]
             res = res[:-idx]
         return res
-
-
-def run_test():
-    key = b'1234523451234545'  # 16 byte key
-    # Iff padded, the message doesn't have to be multiple of 16 in length
-    original_msg = b'ToBeOrNotTobe W.S.'
-    aes = AESCipher(key)
-
-    encrypted_msg = aes.encrypt(original_msg)
-    decrypted_msg = aes.decrypt(encrypted_msg)
-    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)
-    encrypted_msg = aes.encrypt(original_msg)
-    decrypted_msg = aes.decrypt(encrypted_msg)
-    assert decrypted_msg == original_msg
-
-
-if __name__ == '__main__':
-    run_test()

tests/test_92_aes.py

@@ -0,0 +1,74 @@
+import os
+
+import saml2.aes
+
+
+class TestAES():
+    def test_aes_defaults(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(16)
+        aes = saml2.aes.AESCipher(key)
+
+        encrypted_msg = aes.encrypt(original_msg)
+        decrypted_msg = aes.decrypt(encrypted_msg)
+        assert decrypted_msg == original_msg
+
+    def test_aes_128_cbc(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(16)
+        aes = saml2.aes.AESCipher(key)
+        alg = 'aes_128_cbc'
+
+        encrypted_msg = aes.encrypt(original_msg, alg=alg)
+        decrypted_msg = aes.decrypt(encrypted_msg, alg=alg)
+        assert decrypted_msg == original_msg
+
+    def test_aes_128_cfb(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(16)
+        aes = saml2.aes.AESCipher(key)
+        alg = 'aes_128_cfb'
+
+        encrypted_msg = aes.encrypt(original_msg, alg=alg)
+        decrypted_msg = aes.decrypt(encrypted_msg, alg=alg)
+        assert decrypted_msg == original_msg
+
+    def test_aes_192_cbc(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(24)
+        aes = saml2.aes.AESCipher(key)
+        alg = 'aes_192_cbc'
+
+        encrypted_msg = aes.encrypt(original_msg, alg=alg)
+        decrypted_msg = aes.decrypt(encrypted_msg, alg=alg)
+        assert decrypted_msg == original_msg
+
+    def test_aes_192_cfb(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(24)
+        aes = saml2.aes.AESCipher(key)
+        alg = 'aes_192_cfb'
+
+        encrypted_msg = aes.encrypt(original_msg, alg=alg)
+        decrypted_msg = aes.decrypt(encrypted_msg, alg=alg)
+        assert decrypted_msg == original_msg
+
+    def test_aes_256_cbc(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(32)
+        aes = saml2.aes.AESCipher(key)
+        alg = 'aes_256_cbc'
+
+        encrypted_msg = aes.encrypt(original_msg, alg=alg)
+        decrypted_msg = aes.decrypt(encrypted_msg, alg=alg)
+        assert decrypted_msg == original_msg
+
+    def test_aes_256_cfb(self):
+        original_msg = b'ToBeOrNotTobe W.S.'
+        key = os.urandom(32)
+        aes = saml2.aes.AESCipher(key)
+        alg = 'aes_256_cfb'
+
+        encrypted_msg = aes.encrypt(original_msg, alg=alg)
+        decrypted_msg = aes.decrypt(encrypted_msg, alg=alg)
+        assert decrypted_msg == original_msg