#!/usr/bin/env python3
"""vault_scan.py — ARION audit-leg pre/post-deploy scanner for MuseBankVault
(Robinhood Chain). Stdlib only. Runs the falsifier sheet in FALSIFIERS.md
against a deployed address and/or a published source file, and emits one
verdict row per check (append-only JSONL + human table).

Usage:
  python3 vault_scan.py --address 0x... [--source vault.sol] [--out verdicts.jsonl]
  python3 vault_scan.py --source vault.sol            # pre-deploy source review
  python3 vault_scan.py --self-test                   # scan a known RHC contract

Verdict vocabulary (receipts-machine grammar): CONFIRM | FALSIFY | UNBOUND.
Every row carries the pull line a stranger needs to re-run the check.
"""
import json, os, re, sys, time, urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "..", "tools"))
try:
    from eth_call import keccak256
except Exception:
    keccak256 = None

RPCS = ["https://rpc.ordofi.network", "https://rpc.mainnet.chain.robinhood.com",
        "https://robinhood-rpc.publicnode.com", "https://robinhood.drpc.org"]
H = {"Content-Type": "application/json", "User-Agent": "ARION-musebank-audit/1.0"}

# EIP-1967 storage slots
SLOT_IMPL = "0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc"
SLOT_ADMIN = "0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103"
SLOT_BEACON = "0xa3f0ad74e5423aebfd80d3ef4346578335a9a72aeaee59ff6cb3582b35133d50"

# Opcode names of interest (walk skips PUSH data so these are real opcodes)
DANGER_OPS = {0xf1: "CALL", 0xf4: "DELEGATECALL", 0xfa: "STATICCALL",
              0xff: "SELFDESTRUCT", 0xf5: "CREATE2", 0xf0: "CREATE",
              0xf2: "CALLCODE"}

# Admin-gated selectors that would falsify "no owner, no admin keys, no pause"
ADMIN_SIGS = {
    "owner()": "owner key exists",
    "transferOwnership(address)": "ownership transferable",
    "renounceOwnership()": "ownable pattern",
    "pause()": "pause button",
    "unpause()": "pause button",
    "setPaused(bool)": "pause button",
    "upgradeTo(address)": "upgradeable impl",
    "upgradeToAndCall(address,bytes)": "upgradeable impl",
    "setImplementation(address)": "upgradeable impl",
    "setFee(uint256)": "mutable fee",
    "setFeeBps(uint256)": "mutable fee",
    "setDevWallet(address)": "mutable fee wallet",
    "setKeeper(address)": "mutable keeper",
    "setBlacklist(address,bool)": "blacklist",
    "blacklist(address)": "blacklist",
    "freeze(address)": "freeze switch",
    "emergencyWithdraw()": "admin sweep",
    "sweep(address)": "admin sweep",
    "withdrawAll()": "admin sweep",
}

# Access-control / mutability patterns for source review (line-numbered)
SRC_PATTERNS = [
    (r"\bonlyOwner\b|\bonlyAdmin\b|\bonlyRole\b|\brequire\(msg\.sender\s*==", "access gate present"),
    (r"\bpaus\w*\b|Pausable", "pause machinery"),
    (r"\bupgrade|proxy|ERC1967|UUPS|TransparentUpgradeable", "upgrade machinery"),
    (r"\bselfdestruct|suicide\(", "selfdestruct"),
    (r"\bdelegatecall", "delegatecall"),
    (r"\bblacklist|\bfreeze|\bbanned", "blacklist/freeze"),
    (r"\bsetFee|\bfeeBps\s*=|\bdevShare\s*=", "mutable fee/dev share"),
    (r"\btransferTax|_tax|feeOnTransfer", "fee-on-transfer hook"),
]

_i = 0
def rpc(method, params, budget=25):
    global _i
    body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode()
    t0, delay, last = time.time(), 0.4, None
    while True:
        url = RPCS[_i % len(RPCS)]; _i += 1
        try:
            d = json.load(urllib.request.urlopen(
                urllib.request.Request(url, data=body, headers=H), timeout=30))
            if "result" in d:
                time.sleep(0.05)
                return d["result"]
            last = d.get("error")
        except Exception as e:
            last = e
        if time.time() - t0 > budget:
            raise RuntimeError(f"rpc budget: {last}")
        time.sleep(delay); delay = min(delay * 1.7, 2.5)


def op_walk(code: bytes):
    """Yield real opcodes (skip PUSH payload bytes)."""
    i = 0
    while i < len(code):
        op = code[i]
        yield i, op
        if 0x60 <= op <= 0x7f:  # PUSH1..PUSH32
            i += 1 + (op - 0x5f)
        else:
            i += 1


def push4_sels(code: bytes):
    """Collect PUSH4 immediates (candidate function selectors in dispatch)."""
    sels = set()
    i = 0
    while i < len(code):
        op = code[i]
        if op == 0x63 and i + 4 < len(code):
            sels.add(code[i+1:i+5].hex())
            i += 5
        elif 0x60 <= op <= 0x7f:
            i += 1 + (op - 0x5f)
        else:
            i += 1
    return sels


def row(check, clause, result, evidence, pull):
    return {"check": check, "clause": clause, "result": result,
            "evidence": evidence, "pull": pull, "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}


def scan_deployed(addr):
    rows = []
    pull_rpc = f"eth_* vs {'/'.join(RPCS[:1])} (+{len(RPCS)-1} fallbacks)"
    code = rpc("eth_getCode", [addr, "latest"])
    if not code or code in ("0x", "0x0"):
        rows.append(row("deployed", "§2 vault exists", "UNBOUND",
                        "no code at address", f"eth_getCode {addr}"))
        return rows, b""
    cb = bytes.fromhex(code[2:])
    rows.append(row("deployed", "§2 vault exists", "CONFIRM",
                    f"{len(cb)} bytes of code", f"eth_getCode {addr}"))

    # EIP-1967 slots — any nonzero impl/admin/beacon slot falsifies "no upgrades"
    for name, slot in (("impl", SLOT_IMPL), ("admin", SLOT_ADMIN), ("beacon", SLOT_BEACON)):
        v = rpc("eth_getStorageAt", [addr, slot, "latest"])
        nonzero = v and int(v, 16) != 0
        rows.append(row(f"eip1967-{name}", "§2/§5 no upgrades", "FALSIFY" if nonzero else "CONFIRM",
                        f"slot {slot[:12]}… = {v}", f"eth_getStorageAt {addr} {slot}"))

    # Dangerous real opcodes
    hits = {}
    for off, op in op_walk(cb):
        if op in DANGER_OPS:
            hits.setdefault(DANGER_OPS[op], []).append(off)
    for opname, offs in hits.items():
        verdict = "FALSIFY" if opname in ("DELEGATECALL", "SELFDESTRUCT", "CALLCODE") else "FLAG"
        clause = "§2/§5 no upgrades·immutable" if opname in ("DELEGATECALL", "SELFDESTRUCT") else "§4 depositor funds safe"
        rows.append(row(f"opcode-{opname}", clause, verdict,
                        f"{len(offs)} sites e.g. byte-offset {offs[0]}",
                        f"eth_getCode {addr} + opcode walk"))

    # Admin selector presence in dispatch
    if keccak256:
        sels = push4_sels(cb)
        found = [(s, why) for s, why in ADMIN_SIGS.items()
                 if keccak256(s.encode())[:4].hex() in sels]
        if found:
            rows.append(row("admin-selectors", "§2/§5 no owner·admin keys", "FALSIFY",
                            "; ".join(f"{s} ({w})" for s, w in found),
                            f"eth_getCode {addr} + PUSH4 walk"))
        else:
            rows.append(row("admin-selectors", "§2/§5 no owner·admin keys", "CONFIRM",
                            f"0 of {len(ADMIN_SIGS)} admin sigs in dispatch",
                            f"eth_getCode {addr} + PUSH4 walk"))
    return rows, cb


def scan_source(path):
    rows = []
    lines = open(path, encoding="utf-8", errors="replace").read().splitlines()
    pull = f"source file {path} (sha256 via `shasum -a 256 {path}`)"
    for pat, why in SRC_PATTERNS:
        hits = [n + 1 for n, l in enumerate(lines) if re.search(pat, l, re.I)]
        verdict = "FALSIFY" if hits else "CONFIRM"
        ev = f"lines {hits[:8]}{'…' if len(hits) > 8 else ''} match /{pat}/" if hits else f"no match /{pat}/"
        rows.append(row(f"src:{why}", "§2/§5 no owner·admin·upgrade", verdict, ev, pull))
    return rows


def main():
    args = sys.argv[1:]
    if "--self-test" in args:
        # TMT token on RHC — expect: deployed CONFIRM, no proxy slots, no admin sels.
        args = ["--address", "0x974e13c26a2143d6161a08fb00ac842063a99ba3"]
    addr = next((a.split("=", 1)[1] for a in args if a.startswith("--address=")), None)
    if not addr and "--address" in args:
        addr = args[args.index("--address") + 1]
    src = next((a.split("=", 1)[1] for a in args if a.startswith("--source=")), None)
    if not src and "--source" in args:
        src = args[args.index("--source") + 1]
    outp = next((a.split("=", 1)[1] for a in args if a.startswith("--out=")), None)
    if not outp and "--out" in args:
        outp = args[args.index("--out") + 1]

    rows = []
    if addr:
        r, _ = scan_deployed(addr.lower())
        rows += r
    if src:
        rows += scan_source(src)
    if not rows:
        print(__doc__); sys.exit(1)

    for r in rows:
        print(f"{r['result']:8} {r['check']:22} {r['clause']:34} {r['evidence'][:80]}")
    if outp:
        with open(outp, "a") as f:
            for r in rows:
                f.write(json.dumps(r) + "\n")
        print(f"# {len(rows)} rows appended -> {outp}")


if __name__ == "__main__":
    main()
