import string

from celeste.attacks import paddingoracle

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

CHARSET = [c.encode() for c in string.printable]

KEY = b'you wont get me!'
FLAG = b'imbaud{omg_you_catched_me}'
CIPHER = AES.new(KEY, AES.MODE_ECB)

def encrypt(b: bytes, debug=False) -> bytes:
    padded = pad(b + FLAG, 16)
    if debug:
        print(padded)
    # print(padded)
    return CIPHER.encrypt(padded)

def main() -> None:
    paddingoracle.crack(encrypt, pad, CHARSET, 16, batch_size=50, debug=True)

if __name__ == '__main__':
    try:
        main()
    except (KeyboardInterrupt, EOFError):
        print('\n[!] Interrupt')