Cryptopals Challenge 1 Set 13

cryptopals_1_13.py

from Crypto.Cipher import AES
from random import randint


def aes_128_ecb_enc(buffer, key):
    obj = AES.new(key, AES.MODE_ECB)
    return bytearray(obj.encrypt(bytes(buffer)))

def aes_128_ecb_dec(buffer, key):
    obj = AES.new(key, AES.MODE_ECB)
    return bytearray(obj.decrypt(bytes(buffer)))

def pad_pkcs7(buffer, block_size):
    if len(buffer) % block_size:
        padding = (len(buffer) / block_size + 1) * block_size - len(buffer)
    else:
        padding = 0
    # Padding size must be less than a byte 
    assert 0 <= padding <= 255
    new_buffer = bytearray()
    new_buffer[:] = buffer
    new_buffer += bytearray([chr(padding)] * padding)
    return new_buffer

def unpad_pkcs7(buffer):
    padding = buffer[-1]
    for i in range(len(buffer)-1, len(buffer)-padding, -1):
        if buffer[i] != buffer[-1]:
            return buffer
    new_buffer = bytearray()
    new_buffer[:] = buffer[:-padding]
    return new_buffer

def random_key(length):
    key = bytearray(length)
    for i in range(length):
        key[i] = chr(randint(0, 255)) 
    return key

def get_block_size(oracle):
    ciphertext_length = len(oracle(bytearray()))
    i = 1
    while True:
        data = bytearray("A" * i)
        new_ciphertext_length = len(oracle(data))
        block_size = new_ciphertext_length - ciphertext_length
        if block_size:
            return block_size
        i += 1

key = bytes(random_key(AES.block_size))

def str_to_dict(string):
    obj = {}
    for kv in string.split("&"):
        kv = kv.split("=")
        obj[kv[0]] = kv[1]
    return obj

# Oracle
def profile_for(email_buffer):
    email = bytes(email_buffer)
    email = email.replace("&", "").replace("=", "")
    profile = "email=" + email + "&uid=10&role=user"
    padded_buffer = bytes(pad_pkcs7(bytearray(profile), AES.block_size))
    return aes_128_ecb_enc(padded_buffer, key)

def dec_profile(profile):
    return bytes(unpad_pkcs7(aes_128_ecb_dec(profile, key)))

def create_admin_profile():
    block_size = get_block_size(profile_for)

    # Let's make the length of "email=...&uid=10&role=" a multiple of block_size
    # so that "user" is block aligned
    mandatory_bytes = "email=&uid=10&role="
    remaining_bytes = (len(mandatory_bytes) / block_size + 1) * block_size
    email_len = remaining_bytes - len(mandatory_bytes)
    email = "A" * email_len
    profile_prefix = profile_for(bytearray(email))[:remaining_bytes]

    # Let's make the length of "email=..." a multiple of block_size so that
    # the rest of the user inputted email is block aligned
    mandatory_bytes = "email="
    remaining_bytes = (len(mandatory_bytes) / block_size + 1) * block_size
    email_len = remaining_bytes - len(mandatory_bytes)
    email = "A" * email_len
    email += pad_pkcs7("admin", block_size)
    profile_postfix = profile_for(email)[remaining_bytes:remaining_bytes + block_size]

    profile = profile_prefix + profile_postfix
    return bytes(dec_profile(profile))

create_admin_profile()
# 'email=AAAAAAAAAAAAA&uid=10&role=admin'