Update aesencrypt.py

Harriet/Harriet/aesencrypt.py

@@ -1,36 +1,26 @@
 # Red Team Operator course code template
 # payload encryption with AES
 # 
-# author: reenz0h (twitter: @sektor7net)
+# author: reenz0h (twitter: @SEKTOR7net)
 
 import sys
+from base64 import b64encode
 from Crypto.Cipher import AES
-from os import urandom
+from Crypto.Util.Padding import pad
+from Crypto.Random import get_random_bytes
 import hashlib
 
-KEY = urandom(16)
-
-def pad(s):
-#	o = lambda x: x if isinstance(x, int) else ord(x) # handle data being bytes not string
-	return s + (AES.block_size - len(s) % AES.block_size) * chr(AES.block_size - len(s) % AES.block_size)
-
-def aesenc(plaintext, key):
-
-	k = hashlib.sha256(key).digest()
-	iv = 16 * '\x00'
-	plaintext = pad(plaintext)
-	cipher = AES.new(k, AES.MODE_CBC, iv)
- #       o = lambda x: x if isinstance(x, int) else ord(x) # handle data being bytes not string
-	return cipher.encrypt(bytes(plaintext))
-
+KEY = get_random_bytes(16)
+iv = 16 * b'\x00'
+cipher = AES.new(hashlib.sha256(KEY).digest(), AES.MODE_CBC, iv)
 
 try:
-    o = lambda x: x if isinstance(x, int) else ord(x) # handle data being bytes not string
     plaintext = open(sys.argv[1], "rb").read()
 except:
     print("File argument needed! %s <raw payload file>" % sys.argv[0])
     sys.exit()
 
-ciphertext = aesenc(plaintext, KEY)
-print('AESkey[] = { 0x' + ', 0x'.join(hex(ord(x))[2:] for x in KEY) + ' };')
-print('payload[] = { 0x' + ', 0x'.join(hex(ord(x))[2:] for x in ciphertext) + ' };')
+ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))
+
+print('AESkey[] = { 0x' + ', 0x'.join(hex(x)[2:] for x in KEY) + ' };')
+print('payload[] = { 0x' + ', 0x'.join(hex(x)[2:] for x in ciphertext) + ' };')