Decryption Homework

decrypt.py

#!/usr/bin/env python3

from base64 import b64decode
from base64 import b64encode
import os

from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad

def build_keys(key):
    dictionary = ['0','1','2','3','4','5','6','7','8','9','a','b','c','d','e','f']
    keys = []
    for x in dictionary:
        for y in dictionary:
            keys.append(key+x+y)
    return keys

if __name__ == '__main__':
    # encrypt test message
    data = b"secret"
    assert(type(data) == bytes)
    key = get_random_bytes(16)
    test_cipher = AES.new(key, AES.MODE_CBC)
    ct_bytes = test_cipher.encrypt(pad(data, AES.block_size))
    test_iv = b64encode(test_cipher.iv).decode('utf-8')
    ct = b64encode(ct_bytes).decode('utf-8')
    print({'iv': test_iv, 'ciphertext': ct})
    # decrypt test message
    try:
        cipher = AES.new(key, AES.MODE_CBC, b64decode(test_iv))
        pt = unpad(cipher.decrypt(b64decode(ct)), AES.block_size)
        assert(pt == data)
    except:
        print("Incorrect decryption")

    # Get encrypted file
    file = open("./LOCKED","r")
    ciphertext = file.read()
    file.close()
    assert(type(ciphertext) == type(ct))
    ciphertext_bytes = b64decode(ciphertext)
    assert(type(ciphertext_bytes) == type(ct_bytes))
    assert(len(ciphertext_bytes)%16 == 0)
    # build keys
    given_key_part = '6c7578696f5f756e6c6f636b735f34'
    keys = build_keys(given_key_part)
    # build iv
    iv = b'\x00' * 16
    assert(type(iv) == type(b64decode(test_iv)))
    # Process
    for key in keys:
        try:
            cipher = AES.new(b64decode(key), AES.MODE_CBC, iv)
            upt = cipher.decrypt(ciphertext_bytes)
            pt = unpad(upt, AES.block_size)
            test = b64encode(pt)
            if b'MZ' in test[:10] or b'mz' in test[:10]:
                print(key)
                print(test[:10])
                file = open("./test.{}.exe".format(key), "w")
                file.write(test)
                file.close()
                assert(os.access("./test.{}.exe".format(key), os.X_OK))
        except ValueError as e:
            # print(e)
            continue
  print('End')