'''
Solution to https://cryptohack.org/courses/symmetric/symmetry/ 
'''

import os
import requests

from celeste.math.util import xor_bytes, xor_str

URL = 'https://aes.cryptohack.org/symmetry'

def get_encrypted_flag() -> tuple[bytes, bytes]:
    resp = requests.get(f'{URL}/encrypt_flag/')
    ciphertext = resp.json()['ciphertext']
    iv = bytes.fromhex(ciphertext[:32])
    flag = bytes.fromhex(ciphertext[32:])
    return iv, flag

def encrypt(plaintext: bytes, iv: bytes) -> bytes:
    plaintext, iv = plaintext.hex(), iv.hex()
    resp = requests.get(f'{URL}/encrypt/{plaintext}/{iv}')
    return bytes.fromhex(resp.json()['ciphertext'])

def main() -> None:
    iv, flag = get_encrypted_flag()
    # generate a random plaintext of equal length to the flag
    random_plaintext = os.urandom(len(flag))
    random_ciphertext = encrypt(random_plaintext, iv)
    # XOR random plaintext with corresponding ciphertext to
    # get the set generated by IV under the keyed AES permutation
    aes_orbit = xor_bytes(random_plaintext, random_ciphertext)
    # XOR this orbit with the flag to get the hexed flag plaintext
    flag_plaintext = xor_bytes(flag, aes_orbit)
    print('Flag:', flag_plaintext.decode())
    
    print('IV:', iv.hex())

if __name__ == '__main__':
    main()