#!/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)
    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.iv + cipher_text


def process_admin_request(data):
    block_size = len(unknown_key)
    iv = data[:block_size]
    cipher = AES.new(unknown_key, AES.MODE_CBC, iv=iv)

    plain_text = unpad(cipher.decrypt(data[block_size:]))
    for section in plain_text.split(b";"):
        key, value = section.split(b"=", 1)
        if key == b"admin":
            return True
    return False


def test_basic_inject_doesnt_work():
    print("=== Admin Inject Test ===")
    data = encryption_oracle(b";admin=true")
    if process_admin_request(data):
        print("Fail!")
    else:
        print("Pass!")


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 find_bitflip(f, cipher_text, to_flip, block_size):
    data = bytearray(cipher_text)

    # Ignore IV & first block for "unknown" prefix data
    for base_index in range(block_size * 2, len(data) - block_size):
        to_restore = {}
        for position, diff in to_flip.items():
            index = base_index + position
            to_restore[index] = data[index]
            data[index] ^= diff

        try:
            result = f(data)
        except:  # Corrupt Data is corrupt
            result = None
        if result:
            return result

        for index, byte in to_restore.items():
            data[index] = byte
    return None


test_basic_inject_doesnt_work()


print("\n=== Admin Flip ===")
to_inject = b":admin+true"
block_size = detect_block_size(encryption_oracle)
cipher_text = encryption_oracle(b"".join([(b"A" * block_size), to_inject]))
to_flip = {
    to_inject.index(b":"): ord(b":") ^ ord(b";"),
    to_inject.index(b"+"): ord(b"+") ^ ord(b"="),
}

if find_bitflip(process_admin_request, cipher_text, to_flip, block_size):
    print("Got Admin!")
else:
    print("no admin today :(")
