exploit_runner.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import re
import shutil
import signal
import struct
import subprocess
import sys
import traceback
from datetime import datetime, timezone
from pathlib import Path
from types import ModuleType
from typing import Protocol, cast

_ALLOWED_PROOF_TYPES = frozenset(
    {"shell", "arbitrary_read", "arbitrary_write", "vulnerability_trigger"}
)
_ABS_PATH_TOKEN_RE = re.compile(r"(?P<path>/[^\s\"'`]+)")
_REASON_PASS = "attempt_pass"
_REASON_FAIL = "attempt_fail"
_REASON_PARTIAL = "attempt_partial"
_BUNDLE_PASS = "repro_pass"
_BUNDLE_FAIL = "repro_fail"
_BUNDLE_PARTIAL = "repro_partial"

_REDACTION_RULES: tuple[tuple[re.Pattern[str], str], ...] = (
    (re.compile(r"(?i)(authorization\s*:\s*bearer\s+)([^\s]+)"), r"\1[REDACTED]"),
    (re.compile(r"(?i)\b(bearer\s+)([^\s]+)"), r"\1[REDACTED]"),
    (
        re.compile(
            r"(?i)\b(password|passwd|token|api[_-]?key|secret)\b(\s*[:=]\s*)([^\s,;]+)"
        ),
        r"\1\2[REDACTED]",
    ),
)


class PoCResultLike(Protocol):
    success: bool
    proof_type: str
    proof_evidence: str
    timestamp: str


class PoCInterfaceRuntime(Protocol):
    chain_id: str
    target_service: str

    def setup(
        self,
        target_ip: str,
        target_port: int,
        *,
        context: dict[str, object],
    ) -> None: ...

    def execute(self) -> PoCResultLike: ...

    def cleanup(self) -> None: ...


def _utc_now() -> str:
    return (
        datetime.now(tz=timezone.utc)
        .replace(microsecond=0)
        .isoformat()
        .replace("+00:00", "Z")
    )


def _safe_chain_id(chain_id: str) -> str:
    safe = re.sub(r"[^A-Za-z0-9._-]", "_", chain_id.strip())
    return safe if safe else "unknown"


def _sanitize_paths(text: str) -> str:
    def repl(match: re.Match[str]) -> str:
        token = match.group("path")
        p = Path(token)
        return p.name if p.name else "<abs-path>"

    return _ABS_PATH_TOKEN_RE.sub(repl, text)


def redact_sensitive_tokens(text: str) -> str:
    redacted = text
    for pattern, replacement in _REDACTION_RULES:
        redacted = pattern.sub(replacement, redacted)
    return redacted


def _sanitize_log_text(text: str) -> str:
    return redact_sensitive_tokens(_sanitize_paths(text))


def _coerce_proof_evidence(value: object) -> str:
    """Coerce a PoC's ``proof_evidence`` attribute to the string shape the
    runner promises downstream parsers (``poc_validation`` parses
    ``readback_hash=<value>`` tokens out of this string).

    The ``PoCResultLike`` Protocol (line 40) declares
    ``proof_evidence: str``, but LLM-generated plugins routinely drift --
    the 2026-04-13 R7000 plugins used ``Dict[str, Any]`` instead, which
    crashed ``_sanitize_paths`` with ``TypeError: expected string or
    bytes-like object, got 'dict'`` and prevented ``evidence_bundle.json``
    from ever being written. We tolerate dict / list / other shapes by
    JSON-serialising them; plain ``str`` pass through unchanged. Anything
    that fails to serialise is last-resorted to ``repr``.
    """
    if isinstance(value, str):
        return value
    try:
        return json.dumps(value, sort_keys=True, ensure_ascii=True, default=str)
    except Exception:
        return repr(value)


def _write_text(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    _ = path.write_text(content, encoding="utf-8")


def _write_json(path: Path, payload: dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    _ = path.write_text(
        json.dumps(payload, indent=2, sort_keys=True, ensure_ascii=True) + "\n",
        encoding="utf-8",
    )


def _read_json_obj(path: Path) -> dict[str, object]:
    if not path.is_file():
        return {}
    try:
        raw = cast(object, json.loads(path.read_text(encoding="utf-8")))
    except Exception:
        return {}
    if not isinstance(raw, dict):
        return {}
    return cast(dict[str, object], raw)


def _relative_to_run(run_dir: Path, path: Path) -> str:
    return path.resolve().relative_to(run_dir.resolve()).as_posix()


def _compute_sha256(path: Path) -> str:
    hasher = hashlib.sha256()
    with path.open("rb") as fh:
        while True:
            chunk = fh.read(1024 * 1024)
            if not chunk:
                break
            hasher.update(chunk)
    return hasher.hexdigest()


def _load_module_from_path(path: Path) -> ModuleType:
    """Dynamically load a PoC plugin module from a file path.

    The module MUST be registered in ``sys.modules`` before ``exec_module``
    runs. Python 3.12+ ``dataclasses._is_type`` resolves forward references
    via ``sys.modules.get(cls.__module__).__dict__``; if the module is not
    registered, ``.get()`` returns ``None`` and ``.__dict__`` crashes with
    ``AttributeError: 'NoneType' object has no attribute '__dict__'``. This
    was the silent failure mode in the 2026-04-13 R7000 run where every
    LLM-generated plugin (all using ``@dataclass`` for their PoCResult)
    failed to load, leaving ``exploits/chain_*/evidence_bundle.json``
    empty and blocking the verdict_verified cascade.

    See https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
    """
    module_name = (
        f"_private_poc_{hashlib.sha256(str(path).encode('utf-8')).hexdigest()[:16]}"
    )
    spec = importlib.util.spec_from_file_location(module_name, str(path))
    if spec is None or spec.loader is None:
        raise RuntimeError("failed to prepare module loader")
    module = importlib.util.module_from_spec(spec)
    # Register before exec_module so dataclass forward-ref resolution works.
    sys.modules[module_name] = module
    try:
        spec.loader.exec_module(module)
    except Exception:
        # Roll back the registration so a retry with the same path does not
        # observe a half-initialised module.
        sys.modules.pop(module_name, None)
        raise
    return module


def _instantiate_poc(module: ModuleType) -> PoCInterfaceRuntime:
    instance: object
    build = getattr(module, "build_poc", None)
    if callable(build):
        instance = build()
    else:
        cls = getattr(module, "PoC", None)
        if callable(cls):
            instance = cls()
        else:
            raise TypeError("private plugin must expose build_poc() or PoC class")

    for method in ("setup", "execute", "cleanup"):
        if not callable(getattr(instance, method, None)):
            raise TypeError(f"plugin missing required method: {method}")
    chain_id = getattr(instance, "chain_id", None)
    if not isinstance(chain_id, str) or not chain_id:
        raise TypeError("plugin must define non-empty chain_id")
    target_service = getattr(instance, "target_service", None)
    if not isinstance(target_service, str) or not target_service:
        raise TypeError("plugin must define non-empty target_service")
    return cast(PoCInterfaceRuntime, instance)


def _load_private_plugin(
    exploit_dir: Path, chain_id: str
) -> tuple[PoCInterfaceRuntime, Path]:
    matches: list[tuple[PoCInterfaceRuntime, Path]] = []
    for path in sorted(exploit_dir.rglob("*.py"), key=lambda p: p.as_posix()):
        if not path.is_file():
            continue
        module = _load_module_from_path(path)
        poc = _instantiate_poc(module)
        if getattr(poc, "chain_id") == chain_id:
            matches.append((poc, path))

    if not matches:
        raise FileNotFoundError(f"no private plugin found for chain-id: {chain_id}")
    if len(matches) > 1:
        raise RuntimeError("multiple private plugins matched chain-id")
    return matches[0]


def _load_target_details(run_dir: Path) -> tuple[str, int, dict[str, object]]:
    interfaces = _read_json_obj(
        run_dir / "stages" / "dynamic_validation" / "network" / "interfaces.json"
    )
    ports = _read_json_obj(
        run_dir / "stages" / "dynamic_validation" / "network" / "ports.json"
    )

    target_ip = "127.0.0.1"
    interfaces_any = interfaces.get("interfaces")
    if isinstance(interfaces_any, list):
        for iface_any in cast(list[object], interfaces_any):
            if not isinstance(iface_any, dict):
                continue
            iface_obj = cast(dict[str, object], iface_any)
            ipv4_any = iface_obj.get("ipv4")
            if isinstance(ipv4_any, list):
                for candidate in cast(list[object], ipv4_any):
                    if isinstance(candidate, str) and candidate:
                        target_ip = candidate
                        break
            if target_ip != "127.0.0.1":
                break

    target_port = 80
    open_ports_any = ports.get("open_ports")
    if isinstance(open_ports_any, list):
        for candidate in cast(list[object], open_ports_any):
            if isinstance(candidate, int) and candidate > 0:
                target_port = candidate
                break

    context: dict[str, object] = {
        "dynamic_validation": (
            "present"
            if (run_dir / "stages" / "dynamic_validation").is_dir()
            else "missing"
        ),
        "target_ip": target_ip,
        "target_port": target_port,
        "run_dir": str(run_dir),
    }
    return target_ip, target_port, context


def _execute_attempt(
    poc: PoCInterfaceRuntime,
    *,
    attempt_idx: int,
    chain_dir: Path,
) -> dict[str, object]:
    timestamp = _utc_now()
    status = "partial"
    reason_code = _REASON_PARTIAL
    proof_type = "unknown"
    proof_evidence = ""

    try:
        result = poc.execute()
        timestamp_raw = result.timestamp
        if timestamp_raw:
            timestamp = timestamp_raw
        proof_type_raw = result.proof_type
        proof_type = proof_type_raw or "unknown"
        # LLM-generated plugins sometimes violate the PoCResultLike Protocol
        # and hand us a dict/list for proof_evidence. Coerce to str so the
        # downstream sanitiser + readback_hash parser never crash.
        proof_evidence_raw = cast(object, result.proof_evidence)
        proof_evidence = _coerce_proof_evidence(proof_evidence_raw)
        success = bool(result.success)

        if success and proof_type in _ALLOWED_PROOF_TYPES:
            status = "pass"
            reason_code = _REASON_PASS
        elif success:
            status = "partial"
            reason_code = _REASON_PARTIAL
            proof_evidence = (
                proof_evidence
                + " | invalid_proof_type: expected shell|arbitrary_read|arbitrary_write"
            )
        else:
            status = "fail"
            reason_code = _REASON_FAIL
    except Exception as exc:
        status = "partial"
        reason_code = _REASON_PARTIAL
        proof_type = "error"
        proof_evidence = f"execute_exception: {type(exc).__name__}: {exc}"
        tb = _sanitize_log_text(traceback.format_exc())
        _write_text(chain_dir / f"execution_log_{attempt_idx}.txt", tb)
        return {
            "attempt": attempt_idx,
            "reason_code": reason_code,
            "status": status,
            "timestamp": timestamp,
            "proof_type": proof_type,
            "proof_evidence": _sanitize_log_text(proof_evidence),
            "transition_evidence": _transition_evidence_for_attempt(
                poc,
                attempt_idx=attempt_idx,
                attempt_status=status,
                proof_evidence=_sanitize_log_text(proof_evidence),
            ),
        }

    sanitized_evidence = _sanitize_log_text(proof_evidence)
    log_lines = [
        f"attempt={attempt_idx}",
        f"status={status}",
        f"reason_code={reason_code}",
        f"timestamp={timestamp}",
        f"proof_type={proof_type}",
        f"proof_evidence={sanitized_evidence}",
    ]
    _write_text(
        chain_dir / f"execution_log_{attempt_idx}.txt", "\n".join(log_lines) + "\n"
    )
    return {
        "attempt": attempt_idx,
        "reason_code": reason_code,
        "status": status,
        "timestamp": timestamp,
        "proof_type": proof_type,
        "proof_evidence": sanitized_evidence,
        "transition_evidence": _transition_evidence_for_attempt(
            poc,
            attempt_idx=attempt_idx,
            attempt_status=status,
            proof_evidence=sanitized_evidence,
        ),
    }


def _transition_evidence_for_attempt(
    poc: PoCInterfaceRuntime,
    *,
    attempt_idx: int,
    attempt_status: str,
    proof_evidence: str,
) -> list[dict[str, object]]:
    """Project optional Plan IR transitions into evidence-bundle rows.

    AutoPoC templates can carry ``_PLAN_IR`` on the PoC instance. The runner
    does not execute each transition separately, but preserving planned
    transition IDs with observed attempt status lets downstream verifiers and
    analysts reason about multi-step chains without parsing plugin source.
    """
    plan_ir = getattr(poc, "_PLAN_IR", None)
    if not isinstance(plan_ir, dict):
        return []
    transitions_any = plan_ir.get("transitions")
    if not isinstance(transitions_any, list):
        return []

    evidence: list[dict[str, object]] = []
    for idx, item_any in enumerate(cast(list[object], transitions_any), start=1):
        if not isinstance(item_any, dict):
            continue
        item = cast(dict[str, object], item_any)
        transition_id = item.get("transition_id")
        if not isinstance(transition_id, str) or not transition_id:
            transition_id = f"transition_{idx:02d}"
        action = item.get("action")
        channel_type = item.get("channel_type")
        target = item.get("target")
        verifier = item.get("verifier")
        row: dict[str, object] = {
            "attempt": attempt_idx,
            "transition_id": transition_id,
            "action": action if isinstance(action, str) else "unknown",
            "channel_type": (
                channel_type if isinstance(channel_type, str) else "unknown"
            ),
            "target": target if isinstance(target, str) else "unknown",
            "verifier": verifier if isinstance(verifier, str) else "unknown",
            "status": "observed_attempt_pass" if attempt_status == "pass" else "planned_or_unproven",
            "proof_evidence_ref": "attempt.proof_evidence",
        }
        if "readback_hash=" in proof_evidence:
            row["readback_hash_present"] = True
        evidence.append(row)
    return evidence


def _write_minimal_pcap(path: Path, *, linktype: int = 1) -> None:
    """Write a valid pcap global header with zero packets."""
    header = struct.pack(
        "<IHHIIII",
        0xA1B2C3D4,  # magic
        2,  # major
        4,  # minor
        0,  # thiszone
        0,  # sigfigs
        65535,  # snaplen
        int(linktype),  # network linktype (1=Ethernet)
    )
    _ = path.write_bytes(header)


def _capture_pcap(
    pcap_path: Path,
    target_ip: str,
    target_port: int,
    *,
    timeout_s: float = 10.0,
) -> dict[str, str]:
    """Capture network traffic during PoC execution.

    Uses tcpdump when available; falls back to a minimal valid PCAP header
    when tcpdump is missing or fails.  This keeps downstream parsers happy
    regardless of runtime environment.
    """
    tcpdump_bin = shutil.which("tcpdump")
    if not tcpdump_bin:
        _write_minimal_pcap(pcap_path)
        return {"status": "placeholder", "reason": "tcpdump_unavailable"}

    bpf_filter = f"host {target_ip} and port {target_port}"
    cmd = [
        tcpdump_bin,
        "-i",
        "any",
        "-n",
        "-U",
        "-w",
        str(pcap_path),
        "-c",
        "100",
        bpf_filter,
    ]

    try:
        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
        )
    except Exception as exc:
        _write_minimal_pcap(pcap_path)
        return {
            "status": "placeholder",
            "reason": f"tcpdump_start_failed:{type(exc).__name__}",
        }

    try:
        # Drain tcpdump's stdio so it doesn't block on a full PIPE buffer;
        # we don't care about the stderr bytes themselves.
        _ = proc.communicate(timeout=timeout_s)
    except subprocess.TimeoutExpired:
        # Gracefully terminate tcpdump
        try:
            proc.send_signal(signal.SIGTERM)
        except Exception:
            pass
        try:
            proc.wait(timeout=3.0)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass

    # Validate captured file
    if pcap_path.is_file() and pcap_path.stat().st_size >= 24:
        # Verify PCAP magic number
        try:
            header = pcap_path.read_bytes()[:4]
            if header in (b"\xd4\xc3\xb2\xa1", b"\xa1\xb2\xc3\xd4"):
                return {"status": "captured", "reason": "tcpdump_ok"}
        except Exception:
            pass

    # Fallback: tcpdump ran but produced invalid/empty output
    _write_minimal_pcap(pcap_path)
    return {"status": "placeholder", "reason": "tcpdump_capture_empty"}


def run_exploit(
    *,
    run_dir: Path,
    exploit_dir: Path,
    chain_id: str,
    repro: int,
) -> int:
    if not run_dir.is_dir():
        print(f"[FAIL] invalid_contract: run_dir is not a directory: {run_dir}")
        return 1
    if not exploit_dir.is_dir():
        print("[FAIL] invalid_contract: --exploit-dir must be an existing directory")
        return 1
    if repro < 1:
        print("[FAIL] invalid_contract: --repro must be >= 1")
        return 1

    safe_id = _safe_chain_id(chain_id)
    chain_dir = run_dir / "exploits" / f"chain_{safe_id}"
    chain_dir.mkdir(parents=True, exist_ok=True)

    try:
        poc, plugin_file = _load_private_plugin(exploit_dir.resolve(), chain_id)
    except Exception as exc:
        print(f"[FAIL] private_plugin_load_failed: {_sanitize_log_text(str(exc))}")
        return 1

    digest = _compute_sha256(plugin_file)
    _write_text(chain_dir / "poc_sha256.txt", digest + "\n")

    target_ip, target_port, setup_context = _load_target_details(run_dir)

    attempts: list[dict[str, object]] = []
    setup_ok = True
    setup_error = ""
    cleanup_error = ""

    try:
        poc.setup(target_ip, target_port, context=setup_context)
    except Exception as exc:
        setup_ok = False
        setup_error = _sanitize_log_text(f"{type(exc).__name__}: {exc}")

    try:
        if setup_ok:
            for idx in range(1, repro + 1):
                attempts.append(
                    _execute_attempt(poc, attempt_idx=idx, chain_dir=chain_dir)
                )
        else:
            for idx in range(1, repro + 1):
                timestamp = _utc_now()
                log_text = "\n".join(
                    [
                        f"attempt={idx}",
                        "status=partial",
                        f"reason_code={_REASON_PARTIAL}",
                        f"timestamp={timestamp}",
                        "proof_type=setup_error",
                        f"proof_evidence={setup_error}",
                    ]
                )
                _write_text(chain_dir / f"execution_log_{idx}.txt", log_text + "\n")
                attempts.append(
                    {
                        "attempt": idx,
                        "reason_code": _REASON_PARTIAL,
                        "status": "partial",
                        "timestamp": timestamp,
                        "proof_type": "setup_error",
                        "proof_evidence": setup_error,
                    }
                )
    finally:
        try:
            poc.cleanup()
        except Exception as exc:
            cleanup_error = _sanitize_log_text(f"{type(exc).__name__}: {exc}")

    pcap_path = chain_dir / "network_capture.pcap"
    pcap_meta = _capture_pcap(pcap_path, target_ip, target_port, timeout_s=10.0)

    passed = sum(1 for item in attempts if item.get("status") == "pass")
    if passed == repro:
        repro_status = "pass"
        repro_reason = _BUNDLE_PASS
    elif any(item.get("status") == "partial" for item in attempts):
        repro_status = "partial"
        repro_reason = _BUNDLE_PARTIAL
    else:
        repro_status = "fail"
        repro_reason = _BUNDLE_FAIL

    execution_logs = [
        _relative_to_run(run_dir, chain_dir / f"execution_log_{idx}.txt")
        for idx in range(1, repro + 1)
    ]
    transition_evidence = [
        transition
        for attempt in attempts
        for transition in cast(
            list[object],
            attempt.get("transition_evidence", [])
            if isinstance(attempt.get("transition_evidence"), list)
            else [],
        )
        if isinstance(transition, dict)
    ]
    bundle: dict[str, object] = {
        "schema_version": "exploit-evidence-v1",
        "chain_id": chain_id,
        "exploitability_tier": "dynamic_repro",
        "generated_at": _utc_now(),
        "reproducibility": {
            "attempted": repro,
            "passed": passed,
            "reason_code": repro_reason,
            "requested": repro,
            "status": repro_status,
        },
        "attempts": attempts,
        "transition_evidence": transition_evidence,
        "artifacts": {
            "execution_logs": execution_logs,
            "network_capture": _relative_to_run(run_dir, pcap_path),
            "poc_sha256": _relative_to_run(run_dir, chain_dir / "poc_sha256.txt"),
        },
        "pcap": pcap_meta,
        "policy": {
            "private_plugin_path_recorded": False,
            "source_copied_to_run_dir": False,
        },
        "runtime": {
            "cleanup_error": cleanup_error,
            "setup_error": setup_error,
            "target_ip": target_ip,
            "target_port": target_port,
        },
    }
    _write_json(chain_dir / "evidence_bundle.json", bundle)

    if repro_status == "pass":
        print(f"[OK] exploit evidence captured: {_relative_to_run(run_dir, chain_dir)}")
        return 0
    print(
        f"[FAIL] exploit_repro_{repro_status}: {_relative_to_run(run_dir, chain_dir)}"
    )
    return 1


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Run private exploit plugin and capture evidence-only artifacts into run_dir."
        )
    )
    _ = parser.add_argument("--run-dir", required=True, help="Run directory path")
    _ = parser.add_argument(
        "--exploit-dir",
        required=True,
        help="Private exploit directory (external injection only)",
    )
    _ = parser.add_argument("--chain-id", required=True, help="Target exploit chain id")
    _ = parser.add_argument(
        "--repro", type=int, default=3, help="Reproduction attempts"
    )
    args = parser.parse_args(argv)

    run_dir_raw = getattr(args, "run_dir", None)
    exploit_dir_raw = getattr(args, "exploit_dir", None)
    chain_id_raw = getattr(args, "chain_id", None)
    repro_raw = getattr(args, "repro", None)

    if (
        not isinstance(run_dir_raw, str)
        or not run_dir_raw
        or not isinstance(exploit_dir_raw, str)
        or not exploit_dir_raw
        or not isinstance(chain_id_raw, str)
        or not chain_id_raw
        or not isinstance(repro_raw, int)
    ):
        print("[FAIL] invalid_contract: invalid CLI arguments")
        return 1

    return run_exploit(
        run_dir=Path(run_dir_raw).resolve(),
        exploit_dir=Path(exploit_dir_raw).resolve(),
        chain_id=chain_id_raw,
        repro=repro_raw,
    )


if __name__ == "__main__":
    raise SystemExit(main())