import requests

from celeste.constants import PRINTABLE
from celeste.attacks import paddingoracle

from Crypto.Util.Padding import pad

import string

NIBBLESET = [n.to_bytes() for n in range(256)]

HOST = 'https://aes.cryptohack.org'
ENDPOINT = '/ecb_oracle/encrypt/{0}/'
URI = HOST + ENDPOINT

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

FLAG_LEN = 26  # flag length, calculated by hand
SPACER = b'\x0f' # arbitrary spacing character

BLOCK_BYTES = 16
BLOCK_NIBBLES = 32 # measured in nibbles


def mkreq(hextext: str) -> dict[str, str]:
    resp = requests.get(URI.format(hextext))
    if resp.status_code != 200:
        raise Exception(f'[!] resp failed! {resp} : {resp.text}')
    return resp.json()

def encrypt(b: bytes) -> bytes:
    resp = mkreq(b.hex())
    return bytes.fromhex(resp['ciphertext'])

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

               

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