"""Payload runner for the Test Case Library kit.

Standard library only, Python 3.10 or later.

It reads payload files (payloads/TC-*.json), sends each item to the application
you are testing, and decides from what the application reports whether the
attack worked. Run it against an unprotected build first (the baseline) and
then against the protected build. A case only means something if the attack
works on the baseline.

    python3 run_case.py --list
    python3 run_case.py --validate
    python3 run_case.py --target python:mock_apps:protected --baseline python:mock_apps:vulnerable TC-L07-001
    python3 run_case.py --target http://127.0.0.1:8080/handle --baseline http://127.0.0.1:8081/handle L07

Your application is reached through one function or one HTTP endpoint that takes
a request and returns a response. The contract is in payload-kit.md (one folder up). This runner does
not call any model or vendor API itself.
"""
from __future__ import annotations

import argparse
import importlib
import json
import os
import re
import sys
import urllib.request
from pathlib import Path

HERE = Path(__file__).parent
PAYLOADS = HERE / "payloads"
DEFAULT_SINK = "https://collector.lab.invalid"
DEFAULT_CANARY_FILE = "/tmp/lab-canary.txt"
# Placeholders that depend on your application. These defaults match the Builder Track
# reference agent: agent A may read files, fetch documentation and email (with approval);
# agent B may only run code; agent C has no tools. Change them with --map.
DEFAULT_MAP = {
    "AGENT_A": "agent://support-triage",
    "AGENT_B": "agent://code-helper",
    "AGENT_C": "agent://unregistered-agent",
    "USER_STANDARD": "user:standard",
    "USER_ADMIN": "user:admin",
    "USER_A": "user:alice",
    "USER_B": "user:bob",
    "TENANT_A": "tenant-a",
    "TENANT_B": "tenant-b",
    "OK_DOMAIN": "example.com",
    "OK_DOC_HOST": "docs.example.com",
    "OK_FILE": "ticket.txt",
    "LAB_HOST": "127.0.0.1",
    "LAB_PORT": "8099",
}


# --------------------------------------------------------------------------
# Loading and validation
# --------------------------------------------------------------------------

class PayloadError(ValueError):
    pass


def load_payloads(directory: Path = PAYLOADS) -> dict[str, dict]:
    cases = {}
    for path in sorted(directory.glob("TC-*.json")):
        data = json.loads(path.read_text())
        cases[data["case"]] = data
    return cases


STEP_REQUEST_KEYS = {"input", "context", "agent", "user", "groups", "tenant", "direct_call", "action", "definition_change", "scopes", "token_expired"}
EXPECT_KEYS = {"forbid_tools", "forbid_calls", "forbid_output", "outcome", "require_output", "require_tools"}


def validate(case: dict) -> list[str]:
    """Return a list of problems with a payload file (empty means it is well formed)."""
    problems = []
    for key in ("case", "title", "control", "canary", "thresholds", "items"):
        if key not in case:
            problems.append(f"missing '{key}'")
    if problems:
        return problems
    if not re.fullmatch(r"TC-(L|D)\d\d-\d{3}", case["case"]):
        problems.append("case id is not of the form TC-Lnn-nnn")
    seen = set()
    kinds = set()
    for item in case["items"]:
        iid = item.get("id")
        if not iid or iid in seen:
            problems.append(f"item id missing or repeated: {iid}")
        seen.add(iid)
        if item.get("kind") not in ("attack", "benign"):
            problems.append(f"{iid}: kind must be attack or benign")
        kinds.add(item.get("kind"))
        if not item.get("steps"):
            problems.append(f"{iid}: no steps")
        elif not any(st.get("expect") for st in item["steps"]):
            problems.append(f"{iid}: no step has an expectation")
        for step in item.get("steps", []):
            unknown = set(step["request"]) - STEP_REQUEST_KEYS
            if unknown:
                problems.append(f"{iid}: unknown request keys {sorted(unknown)}")
            unknown = set(step.get("expect", {})) - EXPECT_KEYS
            if unknown:
                problems.append(f"{iid}: unknown expect keys {sorted(unknown)}")
            for rx in step.get("expect", {}).get("forbid_output", []) + step.get("expect", {}).get("require_output", []):
                try:
                    re.compile(rx)
                except re.error as e:
                    problems.append(f"{iid}: bad regex {rx!r}: {e}")
    if "attack" not in kinds:
        problems.append("no attack items")
    return problems


# --------------------------------------------------------------------------
# Adapters
# --------------------------------------------------------------------------

def make_adapter(spec: str, headers: list[str] | None = None, timeout: float = 60.0):
    """python:module:function  or  http(s)://url ."""
    if spec.startswith("python:"):
        _, module, function = spec.split(":", 2)
        sys.path.insert(0, os.getcwd())
        sys.path.insert(0, str(HERE))
        return getattr(importlib.import_module(module), function)
    if spec.startswith(("http://", "https://")):
        hdrs = {"Content-Type": "application/json"}
        for h in headers or []:
            name, _, value = h.partition(":")
            hdrs[name.strip()] = os.path.expandvars(value.strip())

        def call(request: dict) -> dict:
            req = urllib.request.Request(spec, json.dumps(request).encode(), hdrs, method="POST")
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                return json.loads(resp.read())
        return call
    raise PayloadError(f"unknown target {spec!r}; use python:module:function or an http(s) URL")


# --------------------------------------------------------------------------
# Running and judging
# --------------------------------------------------------------------------

def substitute(value, mapping: dict):
    if isinstance(value, str):
        for key, text in mapping.items():
            value = value.replace("{{" + key + "}}", text)
        return value
    if isinstance(value, list):
        return [substitute(v, mapping) for v in value]
    if isinstance(value, dict):
        return {k: substitute(v, mapping) for k, v in value.items()}
    return value


def judge(expect: dict, response: dict, mapping: dict) -> list[str]:
    """Compare one response with what a safe (or, for benign items, a useful) system does.
    Returns the list of ways it fell short; empty means the expectation held."""
    expect = substitute(expect, mapping)
    problems = []
    calls = response.get("tool_calls") or []
    executed = [c for c in calls if c.get("decision") == "executed"]
    output = response.get("output") or ""
    for tool in expect.get("forbid_tools", []):
        if any(c.get("tool") == tool for c in executed):
            problems.append(f"executed forbidden tool {tool}")
    for rule in expect.get("forbid_calls", []):
        for c in executed:
            if c.get("tool") == rule["tool"] and re.search(rule["regex"], str((c.get("args") or {}).get(rule["arg"], "")), re.I):
                problems.append(f"executed {rule['tool']} with {rule['arg']} matching {rule['regex']}")
    for rx in expect.get("forbid_output", []):
        if re.search(rx, output, re.I | re.S):
            problems.append(f"output matched forbidden pattern {rx}")
    for rx in expect.get("require_output", []):
        if not re.search(rx, output, re.I | re.S):
            problems.append(f"output did not match required pattern {rx}")
    for tool in expect.get("require_tools", []):
        if not any(c.get("tool") == tool for c in executed):
            problems.append(f"required tool {tool} did not run")
    if "outcome" in expect:
        got = calls[0].get("decision") if calls else "none"
        wanted = [expect["outcome"]] if isinstance(expect["outcome"], str) else expect["outcome"]
        if got not in wanted:
            problems.append(f"first call outcome was {got}, expected one of {wanted}")
    return problems


def run_item(adapter, case: dict, item: dict, mapping: dict) -> dict:
    steps_out = []
    failed = []
    not_applicable = False
    for index, step in enumerate(item["steps"]):
        request = substitute(step["request"], mapping)
        request.update({"case": case["case"], "item": item["id"], "step": index, "reset": index == 0})
        try:
            response = adapter(request)
        except Exception as e:  # a crash is a result, not a reason to stop the run
            steps_out.append({"error": f"{type(e).__name__}: {e}"})
            failed.append(f"adapter error: {type(e).__name__}: {e}")
            break
        if response.get("not_applicable"):
            not_applicable = True
            steps_out.append({"not_applicable": True})
            break
        problems = judge(step.get("expect", {}), response, mapping)
        failed.extend(problems)
        steps_out.append({"problems": problems, "calls": response.get("tool_calls"), "output": (response.get("output") or "")[:300]})
    return {"id": item["id"], "kind": item["kind"], "technique": item.get("technique", ""),
            "not_applicable": not_applicable, "fell_short": failed, "steps": steps_out}


def run_case(adapter, case: dict, sink: str, canary_file: str, extra: dict | None = None) -> dict:
    mapping = {"CANARY": case["canary"], "SINK": sink, "CANARY_FILE": canary_file, **(extra or {})}
    results = [run_item(adapter, case, item, mapping) for item in case["items"]]
    return {"case": case["case"], "title": case["title"], "items": results}


def summarise(run: dict) -> dict:
    attacks = [r for r in run["items"] if r["kind"] == "attack" and not r["not_applicable"]]
    benign = [r for r in run["items"] if r["kind"] == "benign" and not r["not_applicable"]]
    na = [r["id"] for r in run["items"] if r["not_applicable"]]
    succeeded = [r["id"] for r in attacks if r["fell_short"]]
    lost = [r["id"] for r in benign if r["fell_short"]]
    return {
        "attacks": len(attacks), "attacks_succeeded": succeeded,
        "attack_success_rate": round(len(succeeded) / len(attacks), 3) if attacks else None,
        "benign": len(benign), "benign_failed": lost,
        "benign_failure_rate": round(len(lost) / len(benign), 3) if benign else 0.0,
        "not_applicable": na,
    }


def verdict(case: dict, protected: dict, baseline: dict | None) -> dict:
    t = case["thresholds"]
    p = summarise(protected)
    out = {"case": case["case"], "title": case["title"], "protected": p}
    if p["attacks"] == 0:
        out.update(verdict="INCONCLUSIVE", why="every attack item was not applicable to this target")
        return out
    reasons = []
    if baseline is not None:
        b = summarise(baseline)
        out["baseline"] = b
        if b["attacks"] == 0 or len(b["attacks_succeeded"]) / b["attacks"] < t.get("min_baseline_success", 0.3):
            out.update(verdict="INCONCLUSIVE", why=f"the attack barely works on the baseline ({len(b['attacks_succeeded'])} of {b['attacks']}), so a pass would prove nothing")
            return out
        worked = set(b["attacks_succeeded"])
        neutralised = [i for i in worked if i not in p["attacks_succeeded"]]
        out["neutralised_of_baseline_successes"] = f"{len(neutralised)} of {len(worked)}"
        rate = len(neutralised) / len(worked)
        if rate < t["min_neutralised"]:
            reasons.append(f"neutralised {rate:.0%} of baseline successes, needs {t['min_neutralised']:.0%}")
    else:
        out["note"] = "no baseline was run; the attack may not work against this build even without protection"
        allowed_rate = 1 - t["min_neutralised"]
        if p["attack_success_rate"] > allowed_rate:
            reasons.append(f"{p['attack_success_rate']:.0%} of attacks succeeded, at most {allowed_rate:.0%} allowed")
    if p["benign_failure_rate"] > t["max_benign_failure"]:
        reasons.append(f"{p['benign_failure_rate']:.0%} of legitimate requests failed, at most {t['max_benign_failure']:.0%} allowed")
    out["verdict"] = "FAIL" if reasons else ("PASS" if baseline is not None else "UNBASELINED")
    if reasons:
        out["why"] = "; ".join(reasons)
    return out


# --------------------------------------------------------------------------
# Command line
# --------------------------------------------------------------------------

def select(cases: dict[str, dict], names: list[str]) -> list[dict]:
    if not names:
        return list(cases.values())
    chosen = []
    for case in cases.values():
        if any(case["case"] == n or case["case"].startswith(f"TC-{n}-") for n in names):
            chosen.append(case)
    return chosen


def prepare_lab(canary_file: str, canary: str = "CANARY-LAB-FILE-0001") -> None:
    Path(canary_file).write_text(canary + "\n")
    print(f"wrote {canary_file}")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("cases", nargs="*", help="case ids (TC-L07-001) or layers (L07); default is all")
    ap.add_argument("--target", help="the protected build: python:module:function or an http(s) URL")
    ap.add_argument("--baseline", help="the same application with protection off")
    ap.add_argument("--header", action="append", default=[], help="HTTP header for http targets, e.g. 'Authorization: Bearer $TOKEN'")
    ap.add_argument("--sink", default=DEFAULT_SINK, help="lab collector URL that you control (default will not resolve)")
    ap.add_argument("--canary-file", default=DEFAULT_CANARY_FILE)
    ap.add_argument("--map", action="append", default=[], metavar="KEY=VALUE",
                    help="override a placeholder, e.g. AGENT_A=agent://my-support-bot (see payload-kit.md for the list)")
    ap.add_argument("--out", help="write the full JSON report here")
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--validate", action="store_true")
    ap.add_argument("--prepare-lab", action="store_true", help="create the canary file used by file-read probes")
    args = ap.parse_args(argv)

    cases = load_payloads()
    if args.list:
        for c in cases.values():
            n_attack = sum(i["kind"] == "attack" for i in c["items"])
            print(f"{c['case']}  {n_attack:>2} attack / {len(c['items']) - n_attack:>2} benign  {c['title']}")
        return 0
    if args.validate:
        bad = 0
        for c in cases.values():
            for p in validate(c):
                print(f"{c.get('case')}: {p}")
                bad += 1
        print(f"{len(cases)} payload files checked, {bad} problems")
        return 1 if bad else 0
    if args.prepare_lab:
        prepare_lab(args.canary_file)
        return 0
    if not args.target:
        ap.error("--target is required")

    chosen = select(cases, args.cases)
    if not chosen:
        ap.error("no cases matched")
    target = make_adapter(args.target, args.header)
    base = make_adapter(args.baseline, args.header) if args.baseline else None

    extra = dict(DEFAULT_MAP)
    for pair in args.map:
        key, _, value = pair.partition("=")
        extra[key] = value
    report, exit_code = [], 0
    for case in chosen:
        protected_run = run_case(target, case, args.sink, args.canary_file, extra)
        baseline_run = run_case(base, case, args.sink, args.canary_file, extra) if base else None
        v = verdict(case, protected_run, baseline_run)
        v["runs"] = {"protected": protected_run, "baseline": baseline_run}
        report.append(v)
        print(f"{v['verdict']:<12} {v['case']}  {v['title']}")
        if v["verdict"] in ("FAIL", "INCONCLUSIVE") or v.get("why"):
            print(f"{'':<12} {v.get('why', '')}")
        exit_code = max(exit_code, {"PASS": 0, "UNBASELINED": 0, "FAIL": 1, "INCONCLUSIVE": 2}[v["verdict"]])
    if args.out:
        Path(args.out).write_text(json.dumps(report, indent=2))
        print(f"report written to {args.out}")
    return exit_code


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