blacksinisterx/exploit-path-tracer-agent

1"""LangGraph harness for the exploit-path tracer.2 3ingest -> static_scan -> build_call_graph -> trace_reasoning -> poc_drafting4-> fix_drafting -> persist_findings5 6static_scan / build_call_graph are thin wrappers around tools/semgrep_tool.py7and tools/callgraph_tool.py (already verified against the fixture -- see8PLAN.md). trace_reasoning / poc_drafting / fix_drafting are the LLM nodes9(nodes/), where the actual reasoning happens.10"""11import os12from typing import List, Optional, TypedDict13 14from langgraph.graph import END, StateGraph15 16import db17from nodes.fix_drafting import draft_fix18from nodes.poc_drafting import draft_poc19from nodes.trace_reasoning import judge_candidate, merge_hop_judgments20from tools import callgraph_tool, semgrep_tool21 22 23class GraphState(TypedDict):24    target_dir: str25    scan_id: Optional[str]26    sources: List[dict]27    sinks: List[dict]28    candidates: List[dict]29    findings: List[dict]30 31 32def get_llm():33    from langchain_groq import ChatGroq34 35    return ChatGroq(model=os.environ.get("GROQ_MODEL", "llama-3.3-70b-versatile"), temperature=0)36 37 38def ingest(state: GraphState) -> dict:39    db.update_scan_status(state.get("scan_id"), "running")40    db.write_progress(state.get("scan_id"), "ingest", f"Scanning {state['target_dir']}")41    return {}42 43 44def static_scan(state: GraphState) -> dict:45    db.write_progress(state.get("scan_id"), "static_scan", "Running Semgrep taint scan...")46    result = semgrep_tool.scan(state["target_dir"])47    db.write_progress(48        state.get("scan_id"),49        "static_scan",50        f"Found {len(result['sources'])} sources, {len(result['sinks'])} candidate sinks",51    )52    return {"sources": result["sources"], "sinks": result["sinks"]}53 54 55def build_call_graph(state: GraphState) -> dict:56    db.write_progress(state.get("scan_id"), "build_call_graph", "Tracing call-graph hops...")57    index = callgraph_tool.build_index(state["target_dir"])58    candidates = []59    for sink in state["sinks"]:60        for source in state["sources"]:61            hops = callgraph_tool.find_hops(index, source, sink)62            if hops:63                candidates.append({"vuln_class": sink["vuln_class"], "hops": hops})64    db.write_progress(state.get("scan_id"), "build_call_graph", f"Found {len(candidates)} candidate path(s)")65    return {"candidates": candidates}66 67 68def trace_reasoning(state: GraphState) -> dict:69    llm = get_llm()70    findings = []71    total = len(state["candidates"])72    for i, cand in enumerate(state["candidates"], 1):73        db.write_progress(74            state.get("scan_id"),75            "trace_reasoning",76            f"Reasoning about path {i}/{total} ({cand['vuln_class']})...",77        )78        verdict = judge_candidate(llm, cand)79        findings.append(80            {81                "vuln_class": cand["vuln_class"],82                "hops": merge_hop_judgments(cand["hops"], verdict.hops),83                "status": verdict.status,84                "severity": verdict.severity,85                "summary": verdict.summary,86            }87        )88    return {"findings": findings}89 90 91def poc_drafting(state: GraphState) -> dict:92    llm = get_llm()93    findings = state["findings"]94    for f in findings:95        if f["status"] == "confirmed":96            db.write_progress(state.get("scan_id"), "poc_drafting", f"Drafting PoC for {f['vuln_class']}...")97            f["poc"] = draft_poc(llm, f)98    return {"findings": findings}99 100 101def fix_drafting(state: GraphState) -> dict:102    llm = get_llm()103    findings = state["findings"]104    for f in findings:105        if f["status"] == "confirmed":106            db.write_progress(state.get("scan_id"), "fix_drafting", f"Drafting fix for {f['vuln_class']}...")107            f["fix_diff"] = draft_fix(llm, f)108    return {"findings": findings}109 110 111def persist_findings(state: GraphState) -> dict:112    for f in state["findings"]:113        db.write_finding(state.get("scan_id"), f)114    db.update_scan_status(state.get("scan_id"), "completed")115    db.write_progress(state.get("scan_id"), "persist_findings", "Scan complete.")116    return {}117 118 119def build_graph():120    g = StateGraph(GraphState)121    g.add_node("ingest", ingest)122    g.add_node("static_scan", static_scan)123    g.add_node("build_call_graph", build_call_graph)124    g.add_node("trace_reasoning", trace_reasoning)125    g.add_node("poc_drafting", poc_drafting)126    g.add_node("fix_drafting", fix_drafting)127    g.add_node("persist_findings", persist_findings)128 129    g.set_entry_point("ingest")130    g.add_edge("ingest", "static_scan")131    g.add_edge("static_scan", "build_call_graph")132    g.add_edge("build_call_graph", "trace_reasoning")133    g.add_edge("trace_reasoning", "poc_drafting")134    g.add_edge("poc_drafting", "fix_drafting")135    g.add_edge("fix_drafting", "persist_findings")136    g.add_edge("persist_findings", END)137 138    return g.compile()139