verifier.py

#!/usr/bin/env python3
# verifier.py
# Verify a Groth16 proof stored in XRPL memos by transaction hash.
# Supports modules: VerifySign, Allowed, Freshness, Full
# If --vk is omitted, the verification key is inferred from --module and --setup-root
# or auto-detected from the on-chain memos.

import sys
import os
import json
import argparse
import subprocess
import tempfile
from typing import Dict, List, Tuple

from xrpl.clients import JsonRpcClient
from xrpl.models.requests import Tx

TESTNET_RPC = "https://s.altnet.rippletest.net:51234"


# ---------- helpers ----------
def hex_to_str(h: str) -> str:
    try:
        return bytes.fromhex((h or "")).decode("utf-8")
    except Exception:
        return ""

def extract_memos_any(tx_info: dict) -> List[dict]:
    """
    Return the raw 'Memos' array from common response shapes.
    """
    if isinstance(tx_info.get("Memos"), list):
        return tx_info["Memos"]
    for k in ("tx", "transaction", "tx_json", "validated_transaction"):
        v = tx_info.get(k)
        if isinstance(v, dict) and isinstance(v.get("Memos"), list):
            return v["Memos"]
    return []

def fetch_tx(tx_hash: str, rpc_url: str) -> Dict:
    client = JsonRpcClient(rpc_url)
    return client.request(Tx(transaction=tx_hash, binary=False)).result

def get_meta_memo(memos: List[dict], key: str) -> str:
    """
    Return decoded MemoData for the first memo whose MemoType == key.
    """
    for entry in memos:
        m = entry.get("Memo", {})
        mtype = hex_to_str(m.get("MemoType", ""))
        if mtype == key:
            return hex_to_str(m.get("MemoData", ""))
    return ""

def collect_chunked_hex(memos: List[dict], base_label: str) -> str:
    chunks = []
    for entry in memos:
        m = entry.get("Memo", {})
        t = hex_to_str(m.get("MemoType", ""))
        if t.startswith(base_label + "[") and t.endswith("]"):
            try:
                bracket = t[t.index("[") + 1 : -1]
                idx_str, _total = bracket.split("/", 1)
                idx = int(idx_str)
            except Exception:
                continue
            chunks.append((idx, m.get("MemoData", "") or ""))
    if not chunks:
        return ""
    chunks.sort(key=lambda x: x[0])
    return "".join(hex_data for _, hex_data in chunks)

def try_get_proof_and_public_json(memos: List[dict]) -> Tuple[str, str]:
    """
    Supports two formats:
      - Plain JSON in single memos: MemoType=ProofData/PublicData, MemoData=utf8 hex of JSON
      - Gzipped hex split across chunks: MemoType=ProofGZ[i/total]/PublicGZ[i/total], MemoData=hex(zlib.compress(json))
    Returns (proof_json_str, public_json_str).
    """
    proof_plain = get_meta_memo(memos, "ProofData")
    public_plain = get_meta_memo(memos, "PublicData")
    if proof_plain and public_plain:
        json.loads(proof_plain)
        json.loads(public_plain)
        return proof_plain, public_plain

    proof_hex_gz = collect_chunked_hex(memos, "ProofGZ")
    public_hex_gz = collect_chunked_hex(memos, "PublicGZ")
    if not proof_hex_gz or not public_hex_gz:
        raise ValueError("Could not find proof/public data memos (neither plain nor gzipped chunked).")
    try:
        proof_bytes = bytes.fromhex(proof_hex_gz)
        public_bytes = bytes.fromhex(public_hex_gz)
        import zlib as _z
        proof_json_str = _z.decompress(proof_bytes).decode("utf-8")
        public_json_str = _z.decompress(public_bytes).decode("utf-8")
        json.loads(proof_json_str)
        json.loads(public_json_str)
        return proof_json_str, public_json_str
    except Exception as e:
        raise ValueError(f"Failed to decode gzipped memos: {e}")

def write_text(path: str, s: str) -> None:
    with open(path, "w", encoding="utf-8") as f:
        f.write(s)

def infer_vk_path(module: str, setup_root: str) -> str:
    fname = f"{module}_verification_key.json"
    return os.path.join(setup_root, module, fname)

def infer_module(memos: List[dict], public_json_str: str, module_arg: str) -> str:
    """
    Module must be provided explicitly via CLI. No auto-detection.
    """
    return module_arg

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Verify Groth16 proof from XRPL memos by transaction hash.")
    p.add_argument("tx_hash", help="XRPL transaction hash")
    p.add_argument("--rpc", default=TESTNET_RPC, help="XRPL JSON-RPC URL (default: testnet)")
    p.add_argument("--vk", default=None, help="Path to verification_key.json (optional)")
    p.add_argument("--module", choices=["VerifySign", "Allowed", "Freshness", "Full"], required=True,
                   help="Circuit/module name.")
    p.add_argument("--setup-root", default="setup", help="Root dir of setup artifacts (default: setup)")
    return p.parse_args()


# ---------- main ----------
def main():
    args = parse_args()

    tx_hash = args.tx_hash.strip()

    # 1) Fetch tx & memos first (to allow auto module detection)
    tx_info = fetch_tx(tx_hash, args.rpc)
    if not tx_info:
        print("Transaction not found (empty result).")
        sys.exit(3)

    validated = tx_info.get("validated")
    if validated is None:
        validated = (tx_info.get("tx") or tx_info.get("transaction") or tx_info.get("tx_json") or {}).get("validated")
    print("Validated:", validated)

    memos = extract_memos_any(tx_info)
    if not memos:
        print("No memos found in transaction.")
        sys.exit(4)
    print(f"Found {len(memos)} memos.")

    # 2) Read Proof/public JSON from memos

    proof_json_str, public_json_str = try_get_proof_and_public_json(memos)

    # 3) Determine module / verification key
    module_name = infer_module(memos, public_json_str, args.module)
    vk_path = args.vk or infer_vk_path(module_name, args.setup_root)
    print(f"Module: {module_name}")

    if not os.path.exists(vk_path):
        print(f"verification_key.json not found at: {vk_path}")
        sys.exit(2)

    # 4) Write files & verify with snarkjs
    with tempfile.TemporaryDirectory() as d:
        proof_path = os.path.join(d, "proof.json")
        public_path = os.path.join(d, "public.json")
        write_text(proof_path, proof_json_str)
        write_text(public_path, public_json_str)

        cmd = ["snarkjs", "groth16", "verify", vk_path, public_path, proof_path]
        print("Running:", " ".join(cmd))
        try:
            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, text=True)
            print(out.strip())
            if "OK!" in out:
                print("Proof verified successfully against verification_key.json.")
                sys.exit(0)
            else:
                print("snarkjs did not print OK!")
                sys.exit(5)
        except subprocess.CalledProcessError as e:
            print("Verification failed.")
            print(e.output)
            sys.exit(e.returncode)

if __name__ == "__main__":
    main()