diff --git a/malduck/crypto/aes.py b/malduck/crypto/aes.py index 8657d8e..89ba3d5 100644 --- a/malduck/crypto/aes.py +++ b/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): diff --git a/malduck/short.py b/malduck/short.py index 58f5d40..65ec69a 100644 --- a/malduck/short.py +++ b/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) diff --git a/tests/test_crypto.py b/tests/test_crypto.py index 65a1054..747d9c3 100644 --- a/tests/test_crypto.py +++ b/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)