#!/usr/bin/env python3

import os, sys, json, random, hashlib, hmac
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
with open("/flag.txt", "rb") as f:
    flag = f.read()
    
B = 16
opts = (
    "1) encrypt\n"
    "2) profit\n"
    "3) nyerah\n"
    ">> "
)

def exp(prk: bytes, info: bytes, L: int) -> bytes:
    return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:L]

def enc(k: bytes, data: bytes):
    iv = os.urandom(B)
    ct = AES.new(k, AES.MODE_CBC, iv).encrypt(pad(data, B))
    return iv, ct

def inp_hex(prompt: str) -> bytes:
    print(prompt, end="", flush=True)
    s = sys.stdin.readline()
    if not s:
        raise EOFError
    return bytes.fromhex(s.strip())

def rand(rng, d, lo, hi):
    while True:
        v = [rng.randint(lo, hi) for _ in range(d)]
        if any(v):
            return v

def syst(rng):
    d = rng.choice([2, 3])
    m = rng.randint(5, 8)
    x = rand(rng, d, -3, 3)
    rows = []
    for _ in range(m):
        base = rand(rng, d, -3, 3)
        r = rng.randint(1, 7)
        scaled = [r * a for a in base]
        e = rng.randint(0, 1)
        bi = sum(ai * xi for ai, xi in zip(scaled, x)) + e
        rows.append((scaled, bi))
    rng.shuffle(rows)
    A = [row for (row, _) in rows]
    B = [b for (_, b) in rows]
    pub = {"dim": d, "A": A, "b": B}
    return pub, tuple(x)

def bundle():
    rng = random.Random(os.urandom(16))
    systems = []
    hidden = []
    for _ in range(4):
        pub, x = syst(rng)
        systems.append(pub)
        hidden.append(x)
    return {"systems": systems}, tuple(hidden)

def get_key(saltx: bytes, salty: bytes, b):
    parts = []
    for x in b:
        parts.append(",".join(str(t) for t in x))
    material = "|".join(parts).encode()
    prk = hashlib.pbkdf2_hmac("sha256", material, saltx, 131072, dklen=32)
    return exp(prk, b"g3m4zzzt1q" + salty, 16)

def main():
    a, b = bundle()
    print(json.dumps(a, separators=(",", ":")), flush=True)
    saltx = os.urandom(16)
    salty   = os.urandom(16)
    key = get_key(saltx, salty, b)
    iv, ct = enc(key, flag)
    while True:
        try:
            print(opts, end="", flush=True)
            line = sys.stdin.readline()
            if not line:
                break
            try:
                choice = int(line.strip())
            except ValueError:
                print("sheesh")
                continue

            if choice == 1:
                data = inp_hex("data: ")
                iv, ct = enc(key, data)
                print(iv.hex())
                print(ct.hex())

            elif choice == 2:
                print(iv.hex())
                print(ct.hex())
                print(saltx.hex())
                print(salty.hex())

            elif choice == 3:
                print("bubay")
                return

            else:
                print("when you feel like quitting, remember why you started :v (yapping)")

        except Exception:
            print("zzz")

if __name__ == "__main__":
    main()