Cryptopals Challenge Set 2 Exercise 14 Simplified

cryptopals_1_14_simplified.py

from Crypto.Cipher import AES
from random import randint
from collections import defaultdict

def repeated_blocks(buffer, block_length=16):
    reps = defaultdict(lambda: -1)
    for i in range(0, len(buffer), block_length):
        block = bytes(buffer[i:i + block_length])
        reps[block] += 1
    return sum(reps.values())

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 random_key(length):
    key = bytearray(length)
    for i in range(length):
        key[i] = chr(randint(0, 255)) 
    return key

key = bytes(random_key(16))
random_prefix = random_key(randint(0, 15))

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

def encryption_oracle(data):
    unknown_string = bytearray((
        "Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkg\n" +
        "aGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBq\n" +
        "dXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUg\n" +
        "YnkK"
    ).decode("base64"))
    plaintext = pad_pkcs7(
        random_prefix + data + unknown_string,
        AES.block_size,
    )
    return aes_128_ecb_enc(plaintext, 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))
        if ciphertext_length != new_ciphertext_length:
            return new_ciphertext_length - ciphertext_length
        i += 1

def is_ecb_mode(buffer, block_size):
    return repeated_blocks(buffer, block_size) > 0

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

def get_prefix_padding_size(oracle, unknown_string_size, block_size):
    # Only works assuming it's ecb
    prefix_padding = bytearray()
    for i in range(unknown_string_size):
        is_ecb = is_ecb_mode(
            oracle(prefix_padding + bytearray("YELLOW SUBMARINE" * 2)),
            block_size,
        )
        if is_ecb:
            break
        prefix_padding += "A"
    return len(prefix_padding)

def get_prefix_size(oracle, block_size, prefix_padding_size):
    reps = 10
    prefix_padding = bytearray("A" * prefix_padding_size)
    buffer = oracle(prefix_padding + bytearray("YELLOW SUBMARINE" * reps))
    prev_block = count = index = None
    for i in range(0, len(buffer), block_size):
        block = buffer[i: i + block_size]
        if block == prev_block:
            count += 1
        else:
            index = i
            prev_block = block
            count = 1

        if count == reps:
            return index
    return -1

def get_unknown_string(oracle):
    block_size = get_block_size(oracle)
    unknown_string_size = get_unknown_string_size(oracle)
    prefix_padding_size = get_prefix_padding_size(oracle, unknown_string_size, block_size)
    prefix_size_rounded = get_prefix_size(oracle, block_size, prefix_padding_size)
    unknown_string_size -= prefix_size_rounded - prefix_padding_size

    unknown_string = bytearray()
    unknown_string_size_rounded = ((unknown_string_size / block_size) + 1) * block_size
    for i in range(unknown_string_size_rounded - 1, 0, -1):
        d1 = bytearray("A" * (i + prefix_padding_size))
        c1 = oracle(d1)[prefix_size_rounded:unknown_string_size_rounded + prefix_size_rounded]
        for c in range(256):
            d2 = d1[:] + unknown_string + chr(c)
            c2 = oracle(d2)[prefix_size_rounded:unknown_string_size_rounded + prefix_size_rounded]
            if c1 == c2:
                unknown_string += chr(c)
                break
    return unknown_string

get_unknown_string(encryption_oracle)
# Rollin' in my 5.0
# With my rag-top down so my hair can blow
# The girlies on standby waving just to say hi
# Did you stop? No, I just drove by