#!/usr/bin/env python3
"""contract_audit.py — OWASP Smart Contract Top 10 detector + risk scorer.

Pure-Python static analyzer (no solc/forge/slither dependency required).
Designed to be shelled out from the Node bot (/puzzle sc <file.sol | "address 0x...">).

Usage:
    python3 tools/contract_audit.py --selftest
    python3 tools/contract_audit.py --audit /tmp/tester.sol
    python3 tools/contract_audit.py --audit "address 0x..."

Detector set (OWASP SC01-SC10):
  SC01 reentrancy (single / cross-function)
  SC02 unchecked-return / unchecked-call
  SC03 integer overflow / underflow
  SC04 arithmetic / division-by-zero
  SC05 tx.origin misuse (CRITICAL)
  SC06 unsafe delegatecall
  SC07 selfdestruct / suicide (CRITICAL)
  SC08 timestamp manipulation
  SC09 bad randomness (blockhash / difficulty)
  SC10 access control (onlyOwner / tx.origin / missing guard)

Output is JSON-ish text to stdout; the bot parses findings/score/verdict.
No source fetching of addresses is possible without external tools (cast code,
slither, forge). Fallback: addresses are reported with a note, no fabricated
source-level findings.
"""

from __future__ import annotations

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

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

# --------------------------------------------------------------------------
# Detector registry: (id, label, severity, regexes, note, filter_fn)
# severity: HIGH | CRITICAL
# filter_fn(line) -> bool  : True => skip this line for this detector
# --------------------------------------------------------------------------
DETECTORS: list[tuple[str, str, str, list[str], str, object]] = [
    (
        "SC01", "reentrancy",
        "HIGH",
        [
            r"msg\.sender\.call\s*\{",          # value transfer *to* msg.sender
            r"\.send\s*\(.*msg\.sender",        # .send(...) -> msg.sender
            r"\.transfer\s*\(.*msg\.sender",    # .transfer(...) -> msg.sender
        ],
        "value transfer into a public function without (visible) reentrancy guard",
        None,
    ),
    (
        "SC02", "unchecked-return-value",
        "HIGH",
        [
            r"\.call\s*(\.value)?\s*\(",         # .call(...) / .call.value(...) without bool unpack
            r"\.delegatecall\s*\(",              # no bool unpack
            r"\.send\s*\(",                      # send() returns bool, often ignored
            r"\.transfer\s*\(",                  # transfer() returns bool, often ignored
        ],
        "external call result is not checked (unchecked-return-value)",
        lambda line: bool(re.search(r"=\s*\(\s*bool\s+\w+\s*,\s*\w+\s*=", line)),
    ),
    (
        "SC03", "integer-overflow-underflow",
        "HIGH",
        [
            r"(?<![a-zA-Z0-9_])[-+]\s*[a-zA-Z_]\w*",   # a + b / a - b (no SafeMath / not in unchecked)
            r"(?<![a-zA-Z0-9_])[-+*/]\s*[a-zA-Z_]\w*",  # binary arithmetic
        ],
        "arithmetic without SafeMath / unchecked{} (pre-0.8 arithmetic overflow)",
        lambda line: "//" in line and not re.search(r"(?<![a-zA-Z0-9_])([-+*/])(?![a-zA-Z0-9_])", line),
    ),
    (
        "SC04", "arithmetic-division-by-zero",
        "HIGH",
        [
            r"/\s*0\b",                              # division by literal zero
            r"/\s*\(\s*0\s*\)",                      # division by literal zero (paren)
        ],
        "division by zero / division-by-zero hazard",
        None,
    ),
    (
        "SC05", "tx-origin-misuse",
        "CRITICAL",
        [
            r"\btx\.origin\b",
        ],
        "tx.origin used for auth (accepts attacks from anyone sharing the session)",
        None,
    ),
    (
        "SC06", "unsafe-delegatecall",
        "CRITICAL",
        [
            r"\bdelegatecall\s*\(",                                 # to unknown/untrusted target
        ],
        "delegatecall to an untrusted/unknown target (state overwrite)",
        None,
    ),
    (
        "SC07", "selfdestruct-suicide",
        "CRITICAL",
        [
            r"\b(selfdestruct|suicide|selfDestruct)\s*\(",          # suicide is the old keyword name
        ],
        "selfdestruct / suicidenable exit (lost funds + griefing)",
        None,
    ),
    (
        "SC08", "timestamp-manipulation",
        "HIGH",
        [
            r"\bblock\.timestamp\b",
        ],
        "block.timestamp used for state-critical logic (miner ~3s mints, averaging)",
        None,
    ),
    (
        "SC09", "bad-randomness",
        "HIGH",
        [
            r"\bblockhash\b",
            r"\bdifficulty\b",
        ],
        "blockhash/difficulty used for randomness (miner/externally controllable)",
        None,
    ),
    (
        "SC10", "access-control-missing",
        "HIGH",
        [
            r"onlyOwner",                       # presence = good; missing guard otherwise
            r"tx\.origin\b",                     # covered by SC05
            r"require\s*\(.*msg\.sender\s*==\s*",
        ],
        "access control: onlyOwner/msg.sender guard present (or missing)",
        None,
    ),
]


def _find_matches(source: str) -> list[dict]:
    """Raw line-level matches: one entry per (detector, line)."""
    hits: list[dict] = []
    for rid, label, sev, regexes, note, flt in DETECTORS:
        for lineno, line in enumerate(source.splitlines(), start=1):
            line_clean = line.strip()
            if not line_clean or line_clean.startswith("//"):
                continue
            for rx in regexes:
                if re.search(rx, line):
                    if flt is not None and flt(line):
                        continue
                    hits.append({
                        "id": rid, "name": label, "severity": sev,
                        "line": lineno, "snippet": line.strip()[:120],
                        "note": note,
                    })
                    break
    return hits


def _collapse_hits(raw_hits: list[dict]) -> list[dict]:
    """Group raw hits by detector id; keep highest severity."""
    merged: dict[str, dict] = {}
    for h in raw_hits:
        m = merged.get(h["id"])
        if m is None:
            merged[h["id"]] = dict(h)
        else:
            if h["severity"] == "CRITICAL" and m["severity"] != "CRITICAL":
                merged[h["id"]] = dict(h)
    return list(merged.values())


def _external_available(cmd: str) -> bool:
    return shutil.which(cmd) is not None


def fetch_source_for_address(addr: str) -> tuple[str | None, list[str]]:
    """Best-effort source fetch for `address` via available external tools."""
    hints: list[str] = []
    for tool in ("cast", "slither", "forge"):
        if _external_available(tool):
            hints.append(tool)
    if not hints:
        return None, ["no external tool (cast/slither/forge) installed; source unavailable"]
    try:
        proc = subprocess.run(["cast", "code", addr], capture_output=True, text=True, timeout=20)
        if proc.returncode == 0:
            out = proc.stdout.strip()
            if out and not out.startswith("{"):
                return out, hints
    except Exception as e:
        hints.append(f"cast-code failed: {e}")
    return None, hints


def analyze(source: str = "", addr: str | None = None, external: bool = False) -> dict:
    """Run the OWASP SC01-SC10 detector set over `source`."""
    if addr:
        external_s, external_hints = fetch_source_for_address(addr)
        if external_s is None:
            source = ""
            source_kind = f"no source loaded (tried: {'/'.join(external_hints)})"
        else:
            source = external_s
            source_kind = f"bytecode fetched via {' '.join(external_hints)}"
    else:
        source_kind = "local source file"

    hits = _find_matches(source)
    findings = _collapse_hits(hits)

    # ---- access-control *decision* for SC10 ----
    has_only_owner = bool(re.search(r"onlyOwner|Ownable", source, re.I))
    has_guard = bool(re.search(r"nonReentrant|ReentrancyGuard", source, re.I))
    has_constructor_owner = bool(
        re.search(r"constructor\s*\([^)]*\)\s*\{[^}]*owner", source, re.S)
    )
    if has_only_owner:
        sc10_note = "onlyOwner/Ownable guard present"
    elif has_guard:
        sc10_note = "reentrancy guard present (no owner pinning)"
    elif has_constructor_owner:
        sc10_note = "constructor pins owner"
    else:
        sc10_note = "MISSING access-control guard (no owner pinning found)"

    for f in findings:
        if f["id"] == "SC10":
            f["note"] = sc10_note

    # ---- scoring (0-100) ----
    critical = [f for f in findings if f["severity"] == "CRITICAL"]
    n = len(findings)
    score = min(100, 10 * n)
    if critical or n >= 3:
        score = min(100, score + 30)
    verdict = "HIGH" if score >= 50 else "LOW"

    return {
        "verdict": verdict,
        "score": score,
        "finding_count": n,
        "critical_count": len(critical),
        "subjects": [f["id"] for f in findings],
        "findings": findings,
        "source_kind": source_kind,
        "note": sc10_note,
    }


# --------------------------------------------------------------------------
# Built-in test vector: tester.sol
# --------------------------------------------------------------------------
_TESTER_SOL = '''// SPDX-License-Identifier: MIT
pragma solidity ^0.6.0;

contract Tester {
    address public owner;
    mapping(address => uint256) public balances;

    constructor() public {
        owner = tx.origin;   // SC05
    }

    function withdraw(uint256 amount) public {
        require(balances[msg.sender] >= amount, "no funds");
        (bool success, ) = msg.sender.call{value: amount}(new bytes(0)); // SC02 unchecked return
        balances[msg.sender] -= amount;
        require(success, "transfer failed");
    }

    function riskySend(address to, uint256 amount) public {
        to.call.value(amount)(new bytes(0));   // SC02: unchecked call, no bool check
    }

    function add(uint256 a, uint256 b) public pure returns (uint256) {
        uint256 c = a + b;   // SC03
        return c;
    }

    function transfer(address to, uint256 amount) public {
        uint256 balance = balances[msg.sender];
        require(balance >= amount, "insufficient");
        uint256 newBalance = balance - amount; // SC03: underflow if amount > balance
        balances[msg.sender] = newBalance;
        balances[to] += amount;
    }
}
'''

_SELFTEST_EXPECT = {
    "SC01": True, "SC02": True, "SC03": True, "SC05": True,
}


def _selftest() -> int:
    print("== contract_audit self-test ==")
    fd, path = tempfile.mkstemp(suffix=".sol", prefix="tester_")
    try:
        with os.fdopen(fd, "w") as f:
            f.write(_TESTER_SOL)
        result = analyze(_TESTER_SOL)
        print(f"file          : {path}")
        print(f"verdict       : {result['verdict']}  (score {result['score']})")
        print(f"finding_count : {result['finding_count']}  (want >=3)")
        print("findings (id | severity | line):")
        for f in result["findings"]:
            print(f"  {f['id']:10s} | {f['severity']:8s} | L{f['line']:<4} | {f['snippet']!r}")
        # positive controls
        ids = {f["id"] for f in result["findings"]}
        checks = [
            ("SC01", "SC01 reentrancy found"),
            ("SC02", "SC02 unchecked-return found"),
            ("SC03", "SC03 overflow/underflow found"),
            ("SC05", "SC05 tx.origin found"),
            (">=3", f"{result['finding_count']} >= 3"),
            ("score>=50", f"{result['score']} >= 50"),
        ]
        ok = True
        for key, msg in checks:
            good = {"SC01": "SC01" in ids, "SC02": "SC02" in ids,
                    "SC03": "SC03" in ids, "SC05": "SC05" in ids,
                    ">=3": result["finding_count"] >= 3,
                    "score>=50": result["score"] >= 50}[key]
            print(f"[{'PASS' if good else 'FAIL'}] {msg}")
            ok = ok and good
        print(f"source_kind   : {result['source_kind']}")
        print("contract_audit self-test done.")
        return 0 if ok else 1
    finally:
        os.unlink(path)


def _audit_cli(argv: list[str]) -> int:
    target = None
    if "--audit" in argv:
        i = argv.index("--audit")
        target = argv[i + 1]
    if target is None:
        print(__doc__, file=sys.stderr)
        return 2
    addr = None
    if target.lower().startswith(("address", "addr")):
        parts = target.split()
        addr = parts[-1]
        if not addr.startswith("0x"):
            print(f"error: expected 'address 0x<40 hex>', got {target!r}", file=sys.stderr)
            return 2
        source, hints = fetch_source_for_address(addr)
        if source is None:
            print(f"address {addr}: source unavailable (try: {', '.join(hints)})")
            return 0
        print(f"address {addr}: source fetched via {' '.join(hints)}")
        text = source
    else:
        if not os.path.exists(target):
            print(f"error: file not found: {target}", file=sys.stderr)
            return 2
        with open(target, "r", encoding="utf-8") as f:
            text = f.read()
    result = analyze(text, addr=addr)
    print(f"target        : {target}")
    print(f"verdict       : {result['verdict']}")
    print(f"risk_score    : {result['score']}/100")
    print(f"finding_count : {result['finding_count']} (critical {result['critical_count']})")
    print("findings:")
    for f in result["findings"]:
        print(f"  [{f['id']}] {f['severity']:8s} L{f['line']:<5} | {f['snippet']!r}")
        print(f"           note: {f['note']}")
    return 0


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