Merge branch 'feature/fix-aes' into 'develop/3.0.0'

Add AES encryption methods

See merge request cert/malduck!54

malduck/crypto/aes.py

@@ -57,7 +57,7 @@ def export_key(self):
 
 class AES(object):
     r"""
-    AES decryption object
+    AES encryption/decryption object
 
     :param key: Encryption key
     :type key: bytes
@@ -82,7 +82,18 @@ def __init__(self, key, iv=None, mode="cbc"):
         self.aes = Cipher(
             algorithms.AES(key), self.modes[mode](iv),
             backend=default_backend()
-        ).decryptor()
+        )
+
+    def encrypt(self, data):
+        """
+        Encrypt provided data
+
+        :param data: Buffer with data
+        :type data: bytes
+        :return: Encrypted data
+        """
+        aes_enc = self.aes.encryptor()
+        return aes_enc.update(data) + aes_enc.finalize()        
 
     def decrypt(self, data):
         """
@@ -92,7 +103,8 @@ def decrypt(self, data):
         :type data: bytes
         :return: Decrypted data
         """
-        return self.aes.update(data) + self.aes.finalize()
+        aes_dec = self.aes.decryptor()
+        return aes_dec.update(data) + aes_dec.finalize()
 
     @staticmethod
     def import_key(data):

malduck/short.py

@@ -24,6 +24,9 @@ class aes(object):
     def __init__(self, mode):
         self.mode = mode
 
+    def encrypt(self, key=None, iv=None, data=None):
+        return AES(key, iv, self.mode).encrypt(data)
+
     def decrypt(self, key=None, iv=None, data=None):
         return AES(key, iv, self.mode).decrypt(data)
 
@@ -37,8 +40,12 @@ class _cbc_(object):
 
             aes.cbc(key=b'aes128cipher_key',
                     iv=b"iv"*8,
-                    data=pkcs7(b"data_to_be_encrypted", 16))
+                    data=pkcs7(b"data_to_be_decrypted", 16))
         """
+        @staticmethod
+        def encrypt(key=None, iv=None, data=None):
+            return aes("cbc").encrypt(key, iv, data)
+
         @staticmethod
         def decrypt(key=None, iv=None, data=None):
             return aes("cbc").decrypt(key, iv, data)
@@ -56,8 +63,12 @@ class _ecb_(object):
             from malduck import aes, pkcs7
 
             aes.ecb(key=b'aes128cipher_key',
-                    data=pkcs7(b"data_to_be_encrypted", 16))
+                    data=pkcs7(b"data_to_be_decrypted", 16))
         """
+        @staticmethod
+        def encrypt(key=None, data=None):
+            return aes("ecb").encrypt(key, None, data)
+
         @staticmethod
         def decrypt(key=None, data=None):
             return aes("ecb").decrypt(key, None, data)
@@ -76,8 +87,12 @@ class _ctr_(object):
 
             aes.ctr(key=b'aes128cipher_key',
                     nonce=b"iv"*8
-                    data=pkcs7(b"data_to_be_encrypted", 16))
+                    data=pkcs7(b"data_to_be_decrypted", 16))
         """
+        @staticmethod
+        def encrypt(key=None, nonce=None, data=None):
+            return aes("ctr").encrypt(key, nonce, data)
+
         @staticmethod
         def decrypt(key=None, nonce=None, data=None):
             return aes("ctr").decrypt(key, nonce, data)

tests/test_crypto.py

@@ -10,6 +10,7 @@ def test_aes():
         b"I\x96Z\xe4\xb5\xffX\xbdT]\x93\x03\x96\xfcw\xd9"
         b"I\x96Z\xe4\xb5\xffX\xbdT]\x93\x03\x96\xfcw\xd9"
     )
+
     assert aes.ecb.decrypt(b"A"*16, data=b"C"*32) == (
         b"I\x96Z\xe4\xb5\xffX\xbdT]\x93\x03\x96\xfcw\xd9"
         b"I\x96Z\xe4\xb5\xffX\xbdT]\x93\x03\x96\xfcw\xd9"
@@ -20,11 +21,25 @@ def test_aes():
         b"\n\xd5\x19\xa7\xf6\xbc\x1b\xfe\x17\x1e\xd0@\xd5\xbf4\x9a"
     )
 
-    assert aes.ctr(
+    assert aes.ctr.decrypt(
         b"hello world12345", b"A"*16,
         b"\x803\xe3J#\xf4;\x13\x11+h\xf5\xba-\x9b\x05"
     ) == b"B"*16
 
+    assert aes.ecb.encrypt(b"A"*16, b"C"*32) == (
+        b"\xbf\x1ej>.\xc2\xdb_\x9a1&\x17\xee\xfc\x95S"
+        b"\xbf\x1ej>.\xc2\xdb_\x9a1&\x17\xee\xfc\x95S"
+    )
+
+    assert aes.cbc.encrypt(b"A"*16, b"B"*16, b"C"*32) == (
+        b"\xaa\x1a\x18\xffUa_a\"\xf2\x87He\xc8\x1b\xfc"
+        b"\xf9\xcb@\xed\xf6N\xd0-\x9d1rB\xd1\xf2Z\x00"
+    )
+
+    assert aes.ctr.encrypt(b"hello world12345", b"A"*16, b"B"*16) == (
+        b'\x803\xe3J#\xf4;\x13\x11+h\xf5\xba-\x9b\x05'
+    )
+
     assert aes.import_key(
         b"\x08\x02\x00\x00\x0ef\x00\x00\x10\x00\x00\x00" + b"A"*16
     ) == ("AES-128", b"A"*16)