#!/usr/bin/env python3
import sys
from random import randint
from Crypto.Cipher import AES


def random_bytes(size):
    i = randint(0, 2 ** (8 * size))
    return i.to_bytes(size, sys.byteorder)


unknown_key = random_bytes(16)


def unpad(d):
    if not valid_padding(d):
        raise ValueError("Invalid Padding!")
    return d[: -d[-1]]


def pad(d, block_size=16):
    r = block_size - (len(d) % block_size)
    if r > 0:
        p = bytes([r]) * r
    else:
        p = bytes([block_size]) * block_size
    return d + p


def valid_padding(d, block_size=16):
    last = d[-1]
    if not d.endswith(bytes([last]) * last):
        return False

    if last > block_size:
        return False
    return True


def encryption_oracle(data):
    cipher = AES.new(unknown_key, AES.MODE_CBC, iv=unknown_key)
    cipher_text = cipher.encrypt(
        pad(
            b"".join(
                [
                    b"comment1=cooking%20MCs;userdata=",
                    data.replace(b";", b"%3B").replace(b"=", b"%3D"),
                    b";comment2=%20like%20a%20pound%20of%20bacon",
                ]
            )
        )
    )
    return cipher_text


def process_admin_request(data):
    try:
        cipher = AES.new(unknown_key, AES.MODE_CBC, iv=unknown_key)

        plain_text = unpad(cipher.decrypt(data))
        text = plain_text.decode()
        for section in text.split(";"):
            key, value = section.split("=", 1)
            if key == "admin":
                return True
    except UnicodeDecodeError:
        raise ValueError(f"Invalid Decryption", plain_text)
    return False


def xor(a, b):
    r = int.from_bytes(a, sys.byteorder) ^ int.from_bytes(b, sys.byteorder)
    return r.to_bytes(len(a), sys.byteorder)


def detect_block_size(f):
    padding = b""
    orig_length = len(f(padding))
    current_length = orig_length

    while current_length == orig_length:
        padding += b"A"
        current_length = len(f(padding))

    return current_length - orig_length


def get_iv_equals_key_package(cipher_text, block_size=16):
    cipher_text_first_block = cipher_text[block_size : block_size * 2]
    cipher_text_last_2_blocks = cipher_text[-(block_size * 2) :]
    return b"".join(
        [
            cipher_text_first_block,
            b"\x00" * block_size,  # Negate xor for next block & trigger ascii error
            cipher_text_first_block,
            cipher_text_last_2_blocks,  # To pass padding validation
        ]
    )


def main():
    block_size = detect_block_size(encryption_oracle)
    cipher_text = encryption_oracle(b"Stuff")
    modified_cipher_text = get_iv_equals_key_package(cipher_text)

    # Extract the invalid plain text from an error attempt
    corrupt_plain_text = None
    try:
        process_admin_request(modified_cipher_text)
    except ValueError as e:
        expected_msg_contents = "Invalid Decryption"
        if expected_msg_contents in str(e):
            corrupt_plain_text = e.args[1]
        else:
            raise

    if corrupt_plain_text is not None:
        found_key = xor(
            corrupt_plain_text[:block_size],
            corrupt_plain_text[block_size * 2 : block_size * 3],
        )
        if found_key == unknown_key:
            print(f"Found the key: {found_key}")
            cipher = AES.new(found_key, AES.MODE_CBC, iv=found_key)
            msg = unpad(cipher.decrypt(cipher_text)).decode()
            print(f"Decrypted Message: {msg}")
        else:
            print("The key is not in this castle. :(")
    else:
        print("Didn't trigger an ASCII decode error :(")


if __name__ == "__main__":
    main()
