#!/usr/bin/env python3
"""FLARE-On 13 #8 crux - solver.

Usage: python3 -I solve.py
Needs: wasmtime, pycryptodome, sympy (all in ../.venv).

check(input) in main.wasm (garbled Go/js, function 3270) does:
  1. base64-decode the input (standard alphabet);
  2. AES-256-CBC decrypt + PKCS#7 unpad with
        key = cabeba7a...cdbf8723, iv = 5ff26040...ecf65d91
  3. parse the plaintext as a recursive TLV tree (func 3267):
        00                                    -> leaf (prime 2)
        01 u32 n {u32 len, child}*n u32 lp p u32 la a   (p, a little-endian)
  4. main.Kd5oA2q verifies it as a *Pratt primality certificate* of p;
  5. requires p == X, a 256-bit prime rebuilt from an obfuscated literal;
  6. on success returns the flag.
This script rebuilds X (reading the runtime-dependent bytes from a sandboxed
instance; see target_x.py), builds a Pratt certificate for X (pratt.py),
encrypts it, and asks the sandboxed check() for the result.
The sandbox (gojs_sandbox.py) is wasmtime only: no browser, network or file access.
"""
import base64
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from Crypto.Cipher import AES                      # noqa: E402
from Crypto.Util.Padding import pad                # noqa: E402
from sympy import isprime                          # noqa: E402

from gojs_sandbox import GoSandbox                 # noqa: E402
from pratt import cert                             # noqa: E402
from target_x import target_x                      # noqa: E402

WASM = os.path.join(HERE, "loot", "ext", "main.wasm")
KEY = bytes.fromhex("cabeba7ad156801d1c82381c7cb8a3c406ae2c3cd168a96cd8373684cdbf8723")
IV = bytes.fromhex("5ff26040c4c20b9f2ffdefa7ecf65d91")


def main():
    g = GoSandbox(WASM)
    g.run()
    X = int.from_bytes(target_x(g), "little")
    assert isprime(X)
    print(f"[*] target prime X = {X}")
    c = cert(X, "little")
    print(f"[*] Pratt certificate: {len(c)} bytes")
    answer = base64.b64encode(AES.new(KEY, AES.MODE_CBC, IV).encrypt(pad(c, 16))).decode()
    with open(os.path.join(HERE, "loot", "answer_input.txt"), "w") as f:
        f.write(answer + "\n")
    print(f"[*] extension input ({len(answer)} chars) -> loot/answer_input.txt")
    flag = g.call_global("check", answer)
    print(f"[+] check() -> {flag}")
    assert flag.endswith("@flare-on.com")
    with open(os.path.join(HERE, "flag.txt"), "w") as f:
        f.write(flag + "\n")


if __name__ == "__main__":
    main()