#!/usr/bin/env python3
import importlib.util
import json
import re
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]


def load_verifier():
    path = ROOT / "public" / "verify_receipt.py"
    spec = importlib.util.spec_from_file_location("scalaralias_verify", path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


verifier = load_verifier()


def malleate_receipt(receipt: dict) -> dict:
    changed = json.loads(json.dumps(receipt))
    changed["signature"]["s"] = str(int(changed["signature"]["s"]) + verifier.Q)
    return changed


def binding_key(receipt: dict) -> bytes:
    return verifier.hashlib.sha256(
        b"KiteLedger sealed export binding\0" + verifier.canonical_json(receipt)
    ).digest()


def keystream(key: bytes, nonce: bytes, length: int) -> bytes:
    out = b""
    counter = 0
    while len(out) < length:
        out += verifier.hashlib.sha256(
            key + b"scalaralias-stream" + nonce + counter.to_bytes(4, "big")
        ).digest()
        counter += 1
    return out[:length]


def xor_bytes(left: bytes, right: bytes) -> bytes:
    return bytes(a ^ b for a, b in zip(left, right))


def open_export(receipt: dict, sealed: dict) -> str:
    if not verifier.verify_receipt(receipt):
        raise ValueError("receipt does not verify")

    key = binding_key(receipt)
    nonce = bytes.fromhex(sealed["nonce_hex"])
    ciphertext = bytes.fromhex(sealed["ciphertext_hex"])
    tag = bytes.fromhex(sealed["tag_hex"])
    expected = verifier.hashlib.sha256(key + b"export-tag" + nonce + ciphertext).digest()[:16]
    if expected != tag:
        raise ValueError("sealed export tag mismatch")
    plaintext = xor_bytes(ciphertext, keystream(key, nonce, len(ciphertext))).decode()
    match = re.search(r"caplag\{[^}]+\}", plaintext)
    if not match:
        raise ValueError("flag not found")
    return match.group(0)


def main() -> None:
    receipt = json.loads((ROOT / "public" / "receipt.json").read_text())
    sealed = json.loads((ROOT / "public" / "sealed_export.json").read_text())
    print(open_export(malleate_receipt(receipt), sealed))


if __name__ == "__main__":
    main()