from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import base64
import itertools

# Function to compute the CBC-MAC of a given message, IV, and key
def cbc_mac(message: bytes, iv: bytes, key: bytes) -> bytes:
    cipher = AES.new(key, AES.MODE_CBC, iv)
    padded_message = pad(message, 16)  # Ensure the message is padded to the block size
    return cipher.encrypt(padded_message)[-16:]

# Brute-force the key by trying all possible 16-byte keys
def brute_force_keys(target_mac, iv, message, key_size=16):
    """
    Attempts to brute-force the key used for CBC-MAC, given the target MAC.
    :param target_mac: The target MAC we're trying to guess.
    :param iv: The IV used for CBC-MAC.
    :param message: The plaintext message used to generate the MAC.
    :param key_size: The size of the key (16 bytes for AES-128).
    :return: The correct key if found, otherwise None.
    """
    # Loop through all possible 16-byte keys (this is feasible since there are only 2^128 keys)
    # For demonstration purposes, we're assuming a small brute-force keyspace (e.g., 256^16)
    for candidate_key in itertools.product(range(256), repeat=key_size):
        candidate_key = bytes(candidate_key)
        computed_mac = cbc_mac(message, iv, candidate_key)
        if computed_mac == target_mac:
            return candidate_key
    return None

# Example usage:
message = b"type=secrets&number=1337"
iv = base64.b64decode('00000000000000000000000000000000')  # Controlled IV (could be any IV)
# Assume the target MAC is obtained from the server for this message and IV
# In a real attack, you'd receive this MAC from the server
target_mac = base64.b64decode('e0f8b77b2ac5fa872f5646ac90b056ac')  # Example MAC

# Attempt to brute-force the key
guessed_key = brute_force_keys(target_mac, iv, message)

if guessed_key:
    print(f"Found key: {guessed_key.hex()}")
else:
    print("Failed to brute-force the key.")