blacksinisterx/exploit-path-tracer-agent

1"""trace_reasoning node: the core reasoning step. Walks a candidate hop chain2(source -> intermediate functions -> sink) and judges, hop by hop, whether the3tainted value survives to the sink or gets genuinely neutralized somewhere4along the way. This -- not the Semgrep scan -- is what makes this a reasoning5agent instead of a linter: it has to read a function's actual body and decide6whether cosmetic-looking code ("trim", "strip") is real sanitization for the7specific vuln class in question.8"""9from typing import List, Literal10 11from pydantic import BaseModel, Field12 13from nodes._llm_utils import invoke_structured14from tools.fixpatterns import get as get_fix_pattern15 16SYSTEM_PROMPT = """You are a senior application security engineer performing manual taint analysis on one reported flow.17 18You will be given the vulnerability class, domain notes on what real \19sanitization looks like for that class, and a chain of hops from an \20untrusted input source to a dangerous sink. Each hop includes the actual \21source code at that point (a single line for source/sink hops, or a full \22function body for intermediate hops).23 24Rules:25- Judge based ONLY on what the code actually does, not on variable names or \26comments -- comments can be wrong or misleading; verify against the real logic.27- Cosmetic operations alone (trimming whitespace, casing, length checks) do \28NOT count as sanitization for any vuln class.29- A hop earns the "sanitized" verdict only if its logic would genuinely \30block the attack techniques relevant to THIS vulnerability class (see the \31domain notes).32- If any hop after the source genuinely neutralizes the taint before the \33sink, the overall status must be "false_positive" -- even though a naive \34source-to-sink pattern match would flag it as a real bug.35- If nothing neutralizes it, the overall status is "confirmed".36- Use "needs_review" only when the code is genuinely ambiguous (e.g. \37depends on runtime state you cannot evaluate statically).38- Return exactly one hop judgment per hop provided, in the same order given.39"""40 41 42class HopJudgment(BaseModel):43    function: str44    verdict: Literal["source", "tainted", "sanitized", "sink"]45    note: str = Field(description="One sentence: why this specific hop earned this verdict")46 47 48class TraceVerdict(BaseModel):49    status: Literal["confirmed", "false_positive", "needs_review"]50    severity: Literal["critical", "high", "medium", "low"] = Field(51        description="Impact if confirmed; use 'low' if status is false_positive"52    )53    summary: str = Field(description="2-3 sentence explanation of the overall verdict")54    hops: List[HopJudgment]55 56 57def _format_hops(hops):58    lines = []59    for i, h in enumerate(hops, 1):60        lines.append(f"Hop {i} -- function `{h['function']}` @ {h['file']}:{h['line']}")61        lines.append("```python")62        lines.append(h["snippet"])63        lines.append("```")64    return "\n".join(lines)65 66 67def _build_prompt(candidate):68    vuln_class = candidate["vuln_class"]69    domain = get_fix_pattern(vuln_class) or {}70    return (71        f"Vulnerability class: {vuln_class}\n\n"72        f"Domain notes:\n"73        f"- What this class looks like: {domain.get('description', '(none)')}\n"74        f"- What real sanitization looks like: {domain.get('real_sanitization_looks_like', '(none)')}\n\n"75        f"Hop chain ({len(candidate['hops'])} hops, source to sink):\n\n"76        f"{_format_hops(candidate['hops'])}\n\n"77        f"Judge each hop and give the overall verdict."78    )79 80 81def judge_candidate(llm, candidate) -> TraceVerdict:82    structured_llm = llm.with_structured_output(TraceVerdict)83    fallback = TraceVerdict(84        status="needs_review",85        severity="medium",86        summary="Automatic reasoning failed twice for this path; flagged for manual review instead of dropped.",87        hops=[88            HopJudgment(function=h["function"], verdict=h.get("verdict", "tainted"), note="(LLM judgment unavailable)")89            for h in candidate["hops"]90        ],91    )92    return invoke_structured(93        structured_llm,94        [95            {"role": "system", "content": SYSTEM_PROMPT},96            {"role": "user", "content": _build_prompt(candidate)},97        ],98        fallback,99    )100 101 102def merge_hop_judgments(ast_hops, judgments: List[HopJudgment]):103    """Zip the LLM's per-hop verdicts onto the AST tool's hop data (file/line/104    snippet). Defensive against a count mismatch -- structured output105    guarantees valid *shape*, not that the model returned one per hop."""106    merged = []107    for i, hop in enumerate(ast_hops):108        if i < len(judgments):109            verdict, note = judgments[i].verdict, judgments[i].note110        else:111            verdict, note = hop.get("verdict", "tainted"), "(no LLM judgment returned for this hop)"112        merged.append(113            {"file": hop["file"], "line": hop["line"], "snippet": hop["snippet"], "verdict": verdict, "note": note}114        )115    return merged116