#!/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()

menu = (
    "1) encrypt\n"
    "2) profit\n"
    "3) nyerah\n"
    ">> "
)


k = 1024  
n = 169

rng = random.SystemRandom()

def rand(k_bits: int) -> int:
    return rng.randrange(1, 1 << k_bits)

def gen(n: int, k_bits: int):
    a = [rand(k_bits) for _ in range(n)]
    return a

def b2b(b: bytes) -> list[int]:
    out = []
    for byte in b:
        for i in range(8):
            out.append((byte >> i) & 1)
    return out

def pack(bits) -> bytes:
    out = bytearray()
    for i in range(0, len(bits), 8):
        chunk = bits[i:i+8]
        val = 0
        for j, bit in enumerate(chunk):
            val |= (bit & 1) << j
        out.append(val)
    return bytes(out)

def gen_key(x_bits: list[int], saltx: bytes, salty: bytes, length: int = 16) -> bytes:
    material = pack(x_bits)
    prk = hashlib.pbkdf2_hmac("sha256", material, saltx, 131072, dklen=32)
    return hmac.new(prk, b"afk-players-wanted-now" + salty + b"\x01", hashlib.sha256).digest()[:length]

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

def add(a, x_bits) -> int:
    return sum(ai for ai, xi in zip(a, x_bits) if xi)

def main():
    a = gen(n, k)
    x_bits = b2b(flag)
    if len(x_bits) < n:
        x_bits += [rng.randrange(0, 2) for _ in range(n - len(x_bits))]
    else:
        x_bits = x_bits[:n]

    s = add(a, x_bits)

    saltx = os.urandom(16)
    salty = os.urandom(16)
    key = gen_key(x_bits, saltx, salty)

    iv, ct = enc(key, flag)

    header = {
        "n": n,
        "k_bits": k,
    }
    print(json.dumps(header, separators=(",", ":")), flush=True)

    while True:
        try:
            print(menu, end="", flush=True)
            line = sys.stdin.readline()
            if not line:
                break
            try:
                choice = int(line.strip())
            except ValueError:
                print("sheesh")
                continue

            if choice == 1:
                print("data: ", end="", flush=True)
                dline = sys.stdin.readline()
                if not dline:
                    break
                try:
                    data = bytes.fromhex(dline.strip())
                except Exception:
                    print("hmmm")
                    continue
                iv, ct = enc(key, data)
                print(iv.hex())
                print(ct.hex())

            elif choice == 2:
                print(json.dumps({"a": a}, separators=(",", ":")))
                print(str(s))
                print(saltx.hex())
                print(salty.hex())
                print(iv.hex())
                print(ct.hex())

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

            else:
                print("tidak ada yang mustahil, hehehe")

        except Exception:
            print("zzz")

if __name__ == "__main__":
    main()