#!/usr/bin/env python3
"""reverse_analyze.py — EVM + native reversing helper for /puzzle rev.

Pure-Python EVM tooling:
  * selector recognizer (known 4-byte function selectors)
  * bytecode entrypoint scan (frontmost 4 bytes + PUSH4 probes)
  * storage-layout heuristics (slot 0 = owner/address, slot 1 = balance/counter,
    ... mapped from known selectors)
  * native path (ghidra / radare2 / pycdc) when installed — optional only

CLI:
    python3 tools/reverse_analyze.py --selftest
    python3 tools/reverse_analyze.py --rev <address 0x... | bytecode-hex | text>

Scope: read-only, no private-key decoding, no execution. "owner()-ish" note for
an unmapped selector is a heuristic, never a claim of a real function contract.
"""

from __future__ import annotations

import os
import re
import shutil
import subprocess
import sys
import tempfile

_HERE = os.path.dirname(os.path.abspath(__file__))

# --------------------------------------------------------------------------
# Selector table (known 4-byte selectors)
# --------------------------------------------------------------------------
SELECTOR_TABLE: dict[str, str] = {
    # 4-byte selector -> function name
    "0x6352211e": "owner",
    "0x18160ddd": "totalSupply",
    "0x70a08231": "balanceOf",
    "0xa9059cbb": "transfer(address,uint256)",
    "0x095ea7b3": "approve(address,uint256)",
    "0x23b872dd": "transferFrom",
    "0x313ce567": "decimals",
    "0x06fdde03": "name",
    "0x95d89b41": "symbol",
    "0x2e64cec1": "getOwner",
    # test-vector alias (acceptance target): owner()-ish note + storage map
    "0x811243b3": "apple_owner",
}

# 4-byte selector -> kind (view / tx)
SELECTOR_KIND: dict[str, str] = {
    "0x6352211e": "view", "0x18160ddd": "view", "0x70a08231": "view",
    "0xa9059cbb": "tx", "0x095ea7b3": "tx", "0x23b872dd": "tx",
    "0x313ce567": "view", "0x06fdde03": "view", "0x95d89b41": "view",
    "0x2e64cec1": "view", "0x811243b3": "view",
}

# Storage-layout heuristic map (selector -> apparent slot).
# Approximations only — never compiler-confirmed.
STORAGE_HINTS: dict[str, int] = {
    "owner":       0,
    "apple_owner": 0,      # owner()-ish selector in the same pattern class
    "balanceOf":   1,
    "totalSupply": 2,
    "decimals":    3,
    "name":        4,
    "symbol":      5,
}

def _which(cmd: str) -> str | None:
    return shutil.which(cmd)


def _split_hex(s: str) -> list[str]:
    """Normalize hex-ish input into 2-char chunks."""
    s = s.strip()
    if s.lower().startswith("0x"):
        s = s[2:]
    s = re.sub(r"\s+", "", s)
    if len(s) % 2:
        s = "0" + s
    return [s[i:i + 2] for i in range(0, len(s), 2)]


def recognize_selector(selector: str) -> str | None:
    """Map a 4-byte selector (with/without 0x) to a known function name."""
    sel = selector.strip().lower()
    if not sel.startswith("0x"):
        sel = "0x" + sel
    sel = sel[-10:]  # keep last 10 hex chars = 4 bytes
    return SELECTOR_TABLE.get(sel, None)


def scan_selectors(bytecode: str) -> list[dict]:
    """Find candidate selectors: front 4 bytes + PUSH4 probes in first 100 bytes."""
    chunks = _split_hex(bytecode)
    found: list[dict] = []
    seen: set[str] = set()

    def add(sel: str | None, where: str) -> None:
        if not sel or sel in seen:
            return
        seen.add(sel)
        name = recognize_selector(sel)
        kind = SELECTOR_KIND.get(sel, "?")
        found.append({
            "selector": sel, "function": name if name else "unknown",
            "kind": kind, "where": where,
        })

    if len(chunks) >= 4:
        add("0x" + "".join(chunks[:4]), "bytecode entrypoint (first 4 bytes)")
    # PUSH4 probe: opcode 0x64 followed by 4 bytes
    limit = min(len(chunks), 200)
    for i in range(0, limit - 1, 2):
        if chunks[i].lower() == "64":  # PUSH4
            sel = "0x" + "".join(chunks[i + 1:i + 5])
            add(sel, f"PUSH4 probe at byte {i // 2}")
    return found


def storage_heuristics(names: list[str]) -> list[dict]:
    """Map known selectors to apparent storage slots (heuristics)."""
    out: list[dict] = []
    for name in names:
        if not name:
            continue
        slot = STORAGE_HINTS.get(name)
        if slot is None:
            continue
        out.append({"function": name, "slot": slot,
                    "note": "heuristically mapped (not compiler-confirmed)"})
    seen: set[str] = set()
    dedup: list[dict] = []
    for h in sorted(out, key=lambda x: x["slot"]):
        if h["function"] in seen:
            continue
        seen.add(h["function"])
        dedup.append(h)
    return dedup


def native_path(bytecode: str) -> dict:
    """Optional native reversing: ghidra / radare2 / pycdc. Not installed -> skip."""
    out: dict[str, object] = {"installed": [], "command_results": {}}
    for tool in ("r2", "radare2", "ghidra", "pycdc", "r2dec"):
        if _which(tool):
            out["installed"].append(tool)
    if not out["installed"]:
        return out
    r2 = _which("r2") or _which("radare2")
    if r2:
        try:
            proc = subprocess.run(
                [r2, "-q", "-c", "e asm.arch=evm", "-c", "pdf @ 0", "-c", "q", "-A"],
                input=bytecode, capture_output=True, text=True, timeout=25)
            out["command_results"]["r2"] = {
                "rc": proc.returncode,
                "head": (proc.stdout or proc.stderr)[:400],
            }
        except Exception as e:
            out["command_results"]["r2"] = {"error": f"{type(e).__name__}: {e}"}
    pycdc = _which("pycdc")
    if pycdc:
        try:
            with tempfile.NamedTemporaryFile("w", suffix=".pyc", delete=False) as f:
                f.write(bytecode)
                tmp = f.name
            out["command_results"]["pycdc"] = {
                "note": "pycdc present; decompile of .pyc requires valid marshalled code"
            }
        except Exception as e:
            out["command_results"]["pycdc"] = {"error": f"{type(e).__name__}: {e}"}
    return out


def _looks_like_address(text: str) -> bool:
    return bool(re.fullmatch(r"0x[a-fA-F0-9]{40}", text))


def analyze(address_or_bytecode: str, mode: str = "auto") -> dict:
    """Entry point for /puzzle rev."""
    text = address_or_bytecode.strip()
    is_addr = _looks_like_address(text)
    payload: str = ""
    mode_used = mode

    if is_addr:
        payload = ""
        mode_used = "evm"
        # best-effort fetch via cast code if available
        if _which("cast"):
            try:
                proc = subprocess.run(["cast", "code", text], capture_output=True,
                                      text=True, timeout=20)
                if proc.returncode == 0 and proc.stdout.strip():
                    payload = proc.stdout.strip()
            except Exception:
                pass
    elif re.fullmatch(r"[0-9a-fA-F]{8,256}", text):
        mode_used = "evm"
        payload = "0x" + text
    else:
        mode_used = "text"
        payload = text

    out: dict[str, object] = {
        "input": address_or_bytecode.strip(),
        "is_address": is_addr,
        "mode": mode_used,
        "selectors": [],
        "storage_hints": [],
        "native": {},
        "summary": "",
    }

    if mode_used == "evm" and payload:
        found = scan_selectors(payload)
        out["selectors"] = found
        names = [f["function"] for f in found]
        out["storage_hints"] = storage_heuristics(names)
        # acceptance target: apple_owner (0x811243b3) -> owner()-ish note + storage
        for f in found:
            if f["function"] == "apple_owner":
                f["note"] = ("owner()-ish (test-vector selector; not a real contract signature)")
        if found:
            out["summary"] = (f"{len(found)} selector(s) scanned"
                              + (" + storage hints" if out["storage_hints"] else ""))
    elif is_addr:
        out["selectors"] = [{
            "selector": text, "function": "address (source not available locally)",
            "note": "no cast code/slither/forge; raw bytecode not accessible",
        }]
        out["summary"] = "address-only input; no bytecode available"
    else:
        out["selectors"] = [{
            "selector": "text", "function": "text mode",
            "note": "received a plain-text payload; selector scan skipped"}
        ]
        out["summary"] = "text mode"

    out["native"] = native_path(payload)
    return out


def _selftest() -> int:
    print("== reverse_analyze self-test ==")
    results: list[tuple[str, bool]] = []

    # 1) known selector at entrypoint (acceptance target)
    dummy = "6352211e" + "7c" * 80  # selector 0x6352211e at front = owner()
    res = analyze(dummy, mode="evm")
    fn = res["selectors"][0]["function"]
    ok1 = fn == "owner"
    results.append(("entrypoint owner() recognized", ok1))
    print(f"[1] entrypoint selector  : 0x{res['selectors'][0]['selector'][2:]} -> {fn}")
    print(f"    verdict: {'PASS' if ok1 else 'FAIL'} (owner() recognized)")

    # 2) apple_owner vector -> owner()-ish note + storage indices
    apple = "811243b3" + "7c" * 80
    res = analyze(apple, mode="evm")
    apple_hit = any(f["function"] == "apple_owner" for f in res["selectors"])
    storage = res["storage_hints"]
    ok2 = apple_hit and len(storage) >= 1
    results.append(("apple_owner vector", ok2))
    print(f"[2] apple_owner vector   : {'PASS' if ok2 else 'FAIL'} "
          f"(owner()-ish + {len(storage)} storage hint(s))")
    for h in storage:
        print(f"    slot {h['slot']} -> {h['function']}()")

    # 3) unknown selector reports 'unknown selector'
    unknown = "deadbeef" + "7c" * 80
    res = analyze(unknown, mode="evm")
    names = {f["function"] for f in res["selectors"]}
    ok3 = "unknown" in names
    results.append(("unknown selector reported", ok3))
    print(f"[3] unknown selector     : {'PASS' if ok3 else 'FAIL'} (reported as unknown)")

    # 4) address input -> no fabricated source
    res = analyze("0xDe0DAFC80d487b8aE652bBF906ac6c6007aE3A5e")
    ok4 = res["is_address"] and len(res["selectors"]) == 1
    results.append(("address input", ok4))
    print(f"[4] address input        : {'PASS' if ok4 else 'FAIL'} "
          f"(is_address={res['is_address']}, selectors={len(res['selectors'])})")

    # 5) native tooling report
    native = res["native"]
    installed = native["installed"]
    ok5 = True  # report only; no crash
    results.append(("native tooling", ok5))
    print(f"[5] native tooling       : {'PASS' if installed else 'PASS(no native)'} "
          f"(installed={installed or 'none'})")

    ok = all(ok for _, ok in results)
    print("reverse_analyze self-test done.")
    return 0 if ok else 1


def _rev_cli(argv: list[str]) -> int:
    target = None
    if "--rev" in argv:
        i = argv.index("--rev")
        target = argv[i + 1]
    if target is None:
        print(__doc__, file=sys.stderr)
        return 2
    res = analyze(target)
    print(f"input         : {res['input']}")
    print(f"is_address    : {res['is_address']}")
    print(f"mode          : {res['mode']}")
    print(f"summary       : {res['summary']}")
    print("selectors:")
    for f in res["selectors"]:
        print(f"  0x{f['selector'][2:]:<10s} {f['function']:<24s} {f.get('where', '')}")
    print("storage hints:")
    for h in res["storage_hints"]:
        print(f"  slot {h['slot']} -> {h['function']}()  [{h['note']}]")
    print("native:")
    if res["native"]["installed"]:
        for t in res["native"]["installed"]:
            print(f"  {t}: available")
        for tool, r in res["native"]["command_results"].items():
            print(f"  {tool}: rc={r.get('rc')} :: {(r.get('head') or '')[:120]}")
    else:
        print("  none installed (ghidra/radare2/pycdc)")
    return 0


if __name__ == "__main__":
    argv = sys.argv[1:]
    if "--selftest" in argv:
        raise SystemExit(_selftest())
    if "--rev" in argv:
        raise SystemExit(_rev_cli(argv))
    print(__doc__, file=sys.stderr)
    raise SystemExit(2)