from Cryptodome.Cipher import AES
from Cryptodome.Util.Padding import pad, unpad
from Cryptodome.Random import get_random_bytes
#from secret import flag

flag = f"CRYPTO25({'A'*36})"
#Total flag len is 46
#The block size is 16
#16 32 48
assert (len(flag) == len("CRYPTO25{}") + 36)
key = get_random_bytes(24)
flag = flag.encode()

# the encrypted payload is the given Data + Flag
def encrypt() -> bytes:
    data = bytes.fromhex(input("> "))
    payload = data + flag

    cipher = AES.new(key=key, mode=AES.MODE_ECB)
    print(cipher.encrypt(pad(payload, AES.block_size)).hex())


def main():
    menu = \
        "What do you want to do?\n" + \
        "quit - quit the program\n" + \
        "enc - encrypt something\n" + \
        "help - show this menu again\n" + \
        "> "

    while True:
        cmd = input(menu).strip()

        if cmd == "quit":
            break
        elif cmd == "help":
            continue
        elif cmd == "enc":
            encrypt()


if __name__ == '__main__':
    main()