from typing import TypedDict, List, Optional, Any
from wintermute.core import Operation, TestCaseRun
from wintermute.tickets import Ticket
class OrchestratorState(TypedDict, total=False):
    operation: Operation
    ticket: Ticket
    scope: dict[str, Any]              # {target_host, bus, address, tags, ...}
    plan_code: str                     # which TestPlan we attached
    pending_runs: list[str]            # run_ids still to execute
    completed_runs: list[str]          # run_ids that finished (any terminal)
    skipped_runs: list[str]            # run_ids marked not_applicable
    failures: list[str]                # run_ids with vulnerabilities
    report_path: Optional[str]

import re
from wintermute.tickets import Ticket
def read_ticket_node(state: OrchestratorState) -> OrchestratorState:
    """Pull the ticket, parse scope hints, seed state['scope']."""
    tid = state["ticket"]
    if isinstance(tid, str):
        ticket = Ticket.read(tid)
    else:
        ticket = tid
    desc = ticket.data.description or ""
    # Lightweight, deterministic regex parsing so the orchestrator never
    # gambles on the LLM hallucinating a target IP. The agent reasoning
    # happens in plan_node and dispatch_node, not here.
    def find(pat: str) -> str:
        m = re.search(pat, desc)
        return m.group(1) if m else ""
    scope = {
        "target_host": find(r"Target:\s*(\S+)"),
        "bus":         find(r"Bus:\s*(\S+),?"),
        "address":     find(r"Address:\s*(0x[0-9a-fA-F]+)"),
        "tags":        re.findall(r"#(\w+)", desc),  # `#i2c #eeprom #blackbox`
        "title":       ticket.data.title,
    }
    state["ticket"] = ticket
    state["scope"]  = scope
    return state

import json
from pathlib import Path
from wintermute.ai.use import simple_chat
from wintermute.core import TestPlan
PLANS_DIR = Path("TestPlans")
def plan_node(state: OrchestratorState) -> OrchestratorState:
    op = state["operation"]
    scope = state["scope"]
    # 1. Deterministic mapping for common scopes — keeps the LLM out of
    #    the path when we already know the answer.
    plan_path: Path | None = None
    if "blackbox" in scope["tags"] and "hardware" in scope["tags"]:
        plan_path = PLANS_DIR / "TP-HW-BLACKBOX-001.json"
    elif "iam" in scope["tags"] and "aws" in scope["tags"]:
        plan_path = PLANS_DIR / "TP-AWS-IAM-ROLE-001.json"
    elif "redteam" in scope["tags"] and "aws" in scope["tags"]:
        plan_path = PLANS_DIR / "TP-AWS-RED-MEGA-001.json"
    # 2. LLM fallback — describe the engagement and ask which plan applies.
    #    We pin the answer to the catalogue so the model can't invent
    #    a plan code that does not exist on disk.
    if plan_path is None:
        catalogue = "\n".join(
            f"- {p.stem}: {json.loads(p.read_text())['name']}"
            for p in sorted(PLANS_DIR.glob("*.json"))
        )
        prompt = (
            f"Engagement scope: {json.dumps(scope)}\n"
            f"Operation devices: {[d.hostname for d in op.devices]}\n"
            f"Cloud accounts: {[a.name for a in op.cloud_accounts]}\n\n"
            f"Pick the single best test plan for this scope from:\n{catalogue}\n\n"
            "Reply ONLY with the plan code (e.g. TP-HW-BLACKBOX-001)."
        )
        # Cheap path — no tool use, just one classification call.
        plan_code = simple_chat(state["router"], prompt, task_tag="cheap").strip()
        plan_path = PLANS_DIR / f"{plan_code}.json"
        if not plan_path.is_file():
            raise RuntimeError(f"plan_node: model picked a non-existent plan: {plan_code}")
    # 3. Attach + generate
    plan = TestPlan.from_dict(json.loads(plan_path.read_text()))
    op.addTestPlan(plan)
    new_runs = op.generateTestRuns()
    op.save()
    state["plan_code"] = plan.code
    state["pending_runs"] = [r.run_id for r in new_runs if r.status.value == "not_run"]
    state["completed_runs"] = []
    state["skipped_runs"] = []
    state["failures"] = []
    return state

from wintermute.core import RunStatus
DISPATCH_SYSTEM = """You are the dispatcher for a sanctioned hardware
penetration test. You do NOT execute tests yourself. Your sole output is
a JSON object mapping each run_id to one of:
  "execute"  - run it now
  "skip"     - mark not_applicable; the bound target makes the test moot
  "defer"    - mark blocked; needs human input or a permit
Constraints:
- Skip a test only when the bound target clearly makes it irrelevant.
- Defer a test only when destructive (e.g., DA lockout, glitching).
- Default is "execute".
"""
def dispatch_node(state: OrchestratorState) -> OrchestratorState:
    op = state["operation"]
    pending = state["pending_runs"]
    # Build a compact run summary for the LLM — code, target, status, any
    # tags from the parent test case that hint at destructiveness.
    summary = []
    for run_id in pending:
        run = next(r for r in op.test_runs if r.run_id == run_id)
        tc  = next((t for t in op.iterTestCases() if t.code == run.test_case_code), None)
        summary.append({
            "run_id": run_id,
            "test_case": tc.code if tc else "?",
            "name": tc.name if tc else "?",
            "tags": tc.target_scope.tags if tc else [],
            "execution_mode": tc.execution_mode.value if tc else "?",
            "bound": [{"alias": b.alias, "kind": b.kind, "object_id": b.object_id}
                      for b in run.bound],
        })
    prompt = (
        f"Engagement scope: {state['scope']}\n"
        f"Pending runs: {json.dumps(summary, indent=2)}\n\n"
        "Reply with a JSON object: {\"<run_id>\": \"execute|skip|defer\", ...}"
    )
    raw = simple_chat(state["router"], prompt,
                      task_tag="cheap")  # dispatch is cheap reasoning
    decisions: dict[str, str] = json.loads(raw)
    new_pending: list[str] = []
    for run_id in pending:
        decision = decisions.get(run_id, "execute")
        run = next(r for r in op.test_runs if r.run_id == run_id)
        if decision == "skip":
            run.status = RunStatus.not_applicable
            run.finish()
            state["skipped_runs"].append(run_id)
        elif decision == "defer":
            run.status = RunStatus.blocked
            state["completed_runs"].append(run_id)   # done for now
        else:
            new_pending.append(run_id)
    state["pending_runs"] = new_pending
    op.save()
    return state

def execute_runs_node(state: OrchestratorState) -> OrchestratorState:
    """For each pending run, dispatch a sub-agent. Detailed in Part 7."""
    from wintermute.subagents import run_subagent_for_test_case_run  # Part 7
    op = state["operation"]
    while state["pending_runs"]:
        run_id = state["pending_runs"].pop(0)
        run = next(r for r in op.test_runs if r.run_id == run_id)
        try:
            run_subagent_for_test_case_run(op, run, router=state["router"])
        except Exception as exc:
            run.status = RunStatus.blocked
            run.notes = (run.notes + "\n" if run.notes else "") + f"sub-agent error: {exc}"
            run.finish()
        if run.status.value == "failed":
            state["failures"].append(run_id)
        state["completed_runs"].append(run_id)
        op.save()                       # checkpoint after each run
    return state

from wintermute.reports import Report, ReportSpec
from wintermute.backends.docx_reports import DocxTplPerVulnBackend
Report.register_backend("docx", DocxTplPerVulnBackend(
    template_dir="templates",
    main_template="report_main.docx",
    vuln_template="report_vuln.docx",
), make_default=True)
def report_node(state: OrchestratorState) -> OrchestratorState:
    op = state["operation"]
    spec = ReportSpec(
        title=f"{op.operation_name} — Hardware Security Assessment",
        author="Wintermute Orchestrator",
        summary=(
            f"{state['plan_code']} executed against {len(op.devices)} device(s). "
            f"{len(state['completed_runs'])} runs completed, "
            f"{len(state['failures'])} findings, "
            f"{len(state['skipped_runs'])} N/A, "
            f"{sum(1 for r in op.test_runs if r.status.value == 'blocked')} blocked."
        ),
    )
    out = f"./reports/{op.operation_name}.docx"
    Path("reports").mkdir(exist_ok=True)
    Report.save(spec, [op], out)
    state["report_path"] = out
    return state

from langgraph.graph import END, StateGraph
graph = StateGraph(OrchestratorState)
graph.add_node("read_ticket", read_ticket_node)
graph.add_node("plan", plan_node)
graph.add_node("dispatch", dispatch_node)
graph.add_node("execute", execute_runs_node)
graph.add_node("report", report_node)
graph.set_entry_point("read_ticket")
graph.add_edge("read_ticket", "plan")
graph.add_edge("plan", "dispatch")
graph.add_edge("dispatch", "execute")
graph.add_edge("execute", "report")
graph.add_edge("report", END)
compiled = graph.compile()
result = compiled.invoke({
    "operation": op,                                       # built in Part 2
    "ticket": "T000001",                                   # the I²C ticket
    "router": init_router(),
})
print(result["report_path"])

onoSendai > operation load acme-iotcam-2026-Q2
onoSendai [acme-iotcam-2026-Q2] > setup_ticket_backend bugzilla \
    --url https://bugs.acme.local --api-key $BZ_KEY \
    --product IoT-Camera --component Security
[*] Bugzilla backend registered.
onoSendai [acme-iotcam-2026-Q2] > orchestrator run T000001
[*] read_ticket: title="I2C EEPROM extraction on MCIO" tags=#i2c #eeprom #blackbox
[*] plan_node: attached TP-HW-BLACKBOX-001 (24 cases) — generated 17 runs
[*] dispatch_node: 9 execute, 6 defer, 2 skip
[*] execute_runs_node: dispatching sub-agents...
    ...
[*] report_node: ./reports/acme-iotcam-2026-Q2.docx