AES-ECB chosen plaintext attack

aes_ecb_cpa.py

# LICENSE: GPLv3

BLOCK_SIZE = 0x10


def send(prefix: bytes) -> list[bytes]:
    # Implement function that takes an input and returns a list of blocks

    data = ...

    return [data[i:i + BLOCK_SIZE] for i in range(0, len(data), BLOCK_SIZE)]


"""
This dude explains it well
https://crypto.stackexchange.com/a/46921
"""


# The known parts of the second block and after should be are the end of the previous block,
# but on the first block it's whatever we pad it with
# Starts with A*BLOCK_SIZE as a fake "-1st" block - used to simplify first block
known = bytearray(b"A" * BLOCK_SIZE)

# Get full encrypted secret - without any padding - used for knowing how many blocks we have to crack
ct = send(b"")

for block_idx in range(len(ct)):
    for block_byte_idx in range(BLOCK_SIZE):
        padding_len = BLOCK_SIZE - (len(known) % BLOCK_SIZE) - 1
        target = send(b"A" * padding_len)[block_idx]

        guess_prefix = known[-(BLOCK_SIZE - 1):]
        for attempt_byte in range(0x100):
            if send(guess_prefix + bytes([attempt_byte]))[0] == target:
                known.append(attempt_byte)
                # Hit target
                print("K", known[BLOCK_SIZE:], len(known) - BLOCK_SIZE)
                break
        else:
            raise RuntimeError(
                "Could not find a value that gives the target block value - did we reach padding?"
            )

# Don't print the "-1st" block
print("Output:", known[BLOCK_SIZE:])