#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p python3 python3Packages.pwntools

from pwn import *
from pathlib import Path

addr, port, *_ = "mercury.picoctf.net 36449".split(" ")

def chunks(l, n):
    for i in range(0, len(l), n):
        yield l[i:i + n]

def unhex(hex: str) -> list[int]:
    return [int(x, 16) for x in chunks(hex, 2)]

def xor_bytes(a: bytes, b: bytes) -> bytes:
    assert len(a) == len(b)
    return bytes(x ^ y for x, y in zip(a, b))

KEYLEN = 50000

def conn():
    if args.REMOTE:
        return remote(addr, int(port))
    else:
        return process([str(Path(__file__).parent / "otp.py")])

def main():
    r = conn()
    r.recvuntil(b'This is the encrypted flag!\n')
    enc_flag = r.recvline().strip()
    print(f"{enc_flag=}")

    flag_len = len(enc_flag) // 2
    assert flag_len % 2 == 0

    remaining_bytes_until_key_reset = KEYLEN - flag_len

    r.recvuntil(b'What data would you like to encrypt? ')
    r.sendline(b'A'*remaining_bytes_until_key_reset)

    # Ignore output, we're only trying to make the program
    # reset the key offset

    r.recvuntil(b'What data would you like to encrypt? ')
    payload = b"A" * flag_len
    r.sendline(payload)
    r.recvuntil(b'Here ya go!\n')
    response = r.recvline().strip()

    key = xor_bytes(payload, unhex(response))
    print(f"{key=}")

    flag = xor_bytes(unhex(enc_flag), key).decode().strip()
    print(f"picoCTF{{{flag}}}")

if __name__ == "__main__":
    main()