mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 11:49:28 +00:00
Marius CiepluchandClaude Fable 5.1
c8f7610313
66 errors had accumulated in src since the last green run (the hunt modules of groups 4-6, stats, candidates, runner). Codex typed them without behaviour change; the 17 inline import-untyped ignores it added are replaced by one override for dspy, pandas, statsmodels and scikit-learn, and the hunt package's lazy imports of benchmarks/ (never type-checked, not in files) are followed silently so src stays strict against their signatures. Co-Authored-By: Claude Fable 5.1 <[email protected]>
282 lines
14 KiB
Python
282 lines
14 KiB
Python
"""The hunt as a DSPy program: a signature over the feature row, three repository tools, and the answer check.
|
|
|
|
The program sees only the unmasked feature columns of a change, rendered in a fixed order, plus the changed
|
|
hunks. Its answer is the classifier's structure (family from the closed list, confidence, rationale, cited
|
|
lines). DSPy proposes instructions and demonstrations; the acceptance rule in `accept.py` decides, outside DSPy.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from collections.abc import Iterable, Mapping, Sequence
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import dspy
|
|
import pydantic
|
|
|
|
from .. import hunter_tools
|
|
from ..model.endpoint import _PRICES, DEEPSEEK_BASE_URL, DEFAULT_DETECTOR_MODEL
|
|
from ..model.taxonomy import FAMILY_IDS
|
|
|
|
BASELINE_INSTRUCTION = (
|
|
"You are a security reviewer looking at one pushed change. Use read_file, grep_repo and find_refs to check "
|
|
"whether the changed code introduces a vulnerability that an attacker could reach. Report only what the code "
|
|
"shows: a path, a line, the family from the list, your confidence between 0 and 1, a one-sentence rationale and "
|
|
"the lines you relied on. Report nothing when the change is ordinary."
|
|
)
|
|
|
|
# Input name -> (feature column, renderer). Order is the order the model sees them in; masked columns are absent.
|
|
INPUTS: tuple[tuple[str, str, str], ...] = (
|
|
("change_summary", "", "The files and functions the push touched."),
|
|
(
|
|
"hunks",
|
|
"feature.diff.hunks",
|
|
"Per file at head: changed line ranges as [start, end] pairs, and the lines after which the push removed code.",
|
|
),
|
|
("function_text", "feature.source.span", "The enclosing function at head, numbered lines."),
|
|
("callers", "feature.source.callers", "Where the enclosing function is called from: path and line."),
|
|
("entry_points", "feature.source.entry_points", "In-repository entry points that reach the file."),
|
|
("roles", "feature.semantic.roles", "Identifiers known as sources, sinks or sanitizers in this function."),
|
|
("def_use", "feature.semantic.def_use", "Assignments and calls inside the function, with lines."),
|
|
("obligations", "feature.semantic.obligations", "Checks the code is expected to perform before an operation."),
|
|
("families", "feature.taxonomy.families", "The closed list of vulnerability families with one-line definitions."),
|
|
)
|
|
|
|
|
|
class Finding(pydantic.BaseModel):
|
|
"""One suspicion in the classifier's answer structure."""
|
|
|
|
path: str
|
|
line: int
|
|
family: str
|
|
confidence: float = 0.5
|
|
rationale: str = ""
|
|
cited_lines: list[int] = pydantic.Field(default_factory=list)
|
|
|
|
|
|
def signature_for(mask: Iterable[str] = ()) -> type[dspy.Signature]:
|
|
"""The hunt signature with every unmasked input, in the fixed order, and the findings output."""
|
|
dropped = set(mask)
|
|
fields: dict[str, Any] = {}
|
|
for name, column, desc in INPUTS:
|
|
producer = column.split(".")[1] if column else ""
|
|
if name in dropped or column in dropped or (producer and producer in dropped):
|
|
continue
|
|
fields[name] = (str, dspy.InputField(desc=desc))
|
|
fields["findings"] = (list[Finding], dspy.OutputField(desc="JSON list; empty when the change is ordinary"))
|
|
signature: type[dspy.Signature] = dspy.Signature(fields, BASELINE_INSTRUCTION)
|
|
return signature
|
|
|
|
|
|
def numbered(lines: Sequence[str], start: int, end: int) -> str:
|
|
"""Lines start..end (1-based, inclusive) with their numbers, so cited lines can be checked."""
|
|
return "\n".join(f"{number:>5} {lines[number - 1]}" for number in range(max(start, 1), min(end, len(lines)) + 1))
|
|
|
|
|
|
def render_rows(
|
|
change: Mapping[str, Any], rows: Sequence[Mapping[str, Any]], texts: Mapping[str, Sequence[str]], *, mask: Iterable[str] = ()
|
|
) -> dict[str, str]:
|
|
"""Render one change's candidate rows into the signature's inputs. Masked inputs are not produced at all."""
|
|
signature = signature_for(mask)
|
|
wanted = set(signature.input_fields)
|
|
out: dict[str, str] = {}
|
|
if "change_summary" in wanted:
|
|
touched = sorted({(row["path"], str(row.get("feature.source.function") or "<module>")) for row in rows})
|
|
out["change_summary"] = "\n".join(f"{path}::{function}" for path, function in touched) or "(nothing readable)"
|
|
if "hunks" in wanted:
|
|
hunks: dict[str, dict[str, list]] = {}
|
|
for row in rows:
|
|
if row.get("feature.diff.hunks") or row.get("feature.diff.deletions"):
|
|
hunks[row["path"]] = {
|
|
"changed": row.get("feature.diff.hunks") or [],
|
|
"removed_after_line": row.get("feature.diff.deletions") or [],
|
|
}
|
|
out["hunks"] = json.dumps(hunks, sort_keys=True) if hunks else "(no hunks)"
|
|
if "function_text" in wanted:
|
|
blocks = []
|
|
for row in rows:
|
|
span = row.get("feature.source.span")
|
|
lines = texts.get(row["path"])
|
|
if span and lines:
|
|
blocks.append(f"# {row['path']}::{row.get('feature.source.function')}\n{numbered(lines, span[0], span[1])}")
|
|
out["function_text"] = "\n\n".join(blocks) or "(no function text)"
|
|
for name, column, _ in INPUTS[3:]:
|
|
if name not in wanted:
|
|
continue
|
|
values = {f"{row['path']}::{row.get('feature.source.function') or '<module>'}": row.get(column) for row in rows}
|
|
values = {key: value for key, value in values.items() if value not in (None, [], {})}
|
|
if name == "families": # one list for the whole change, not one per row
|
|
families = next((value for value in values.values()), None)
|
|
out[name] = "\n".join(f"{f['id']}: {f['definition']}" for f in families) if families else "(no families)"
|
|
else:
|
|
out[name] = json.dumps(values, sort_keys=True) if values else "(none)"
|
|
return out
|
|
|
|
|
|
def shown_lines(inputs: Mapping[str, str]) -> dict[str, set[int]]:
|
|
"""The numbered lines the model was shown per path, parsed back from `function_text`."""
|
|
shown: dict[str, set[int]] = {}
|
|
current = None
|
|
for line in inputs.get("function_text", "").splitlines():
|
|
if line.startswith("# ") and "::" in line:
|
|
current = line[2:].split("::", 1)[0]
|
|
elif current and line[:5].strip().isdigit():
|
|
shown.setdefault(current, set()).add(int(line[:5]))
|
|
return shown
|
|
|
|
|
|
def check(
|
|
findings: Iterable[Finding | Mapping[str, Any]], inputs: Mapping[str, str], seen: Mapping[str, set[int]] | None = None
|
|
) -> list[dict[str, Any]]:
|
|
"""Apply the answer rules: a family outside the closed list, or a cited line the model was not shown, scores
|
|
against the program by becoming `unknown` with a reason. Nothing is dropped, so the metric sees it."""
|
|
shown = shown_lines(inputs)
|
|
for path, lines in (seen or {}).items():
|
|
shown.setdefault(path, set()).update(lines)
|
|
allowed = {family for family in FAMILY_IDS if family != "unknown"}
|
|
out = []
|
|
for item in findings:
|
|
data = item.model_dump() if isinstance(item, pydantic.BaseModel) else dict(item)
|
|
reasons = []
|
|
if data.get("family") not in allowed:
|
|
reasons.append("family_not_in_list")
|
|
cited = [int(line) for line in data.get("cited_lines") or []]
|
|
if cited and not set(cited) <= shown.get(data.get("path", ""), set()):
|
|
reasons.append("cited_line_not_shown")
|
|
if reasons:
|
|
data["reported_family"] = data.get("family")
|
|
data["family"] = "unknown"
|
|
data["unknown_reason"] = ",".join(reasons)
|
|
out.append(data)
|
|
return out
|
|
|
|
|
|
def repository_tools(
|
|
root: Path, *, max_chars: int = 4000, max_matches: int = 50, seen: dict[str, set[int]] | None = None
|
|
) -> list[dspy.Tool]:
|
|
"""The three tools the shipped hunter has, bound to one checkout and clamped to it. Every line a tool returns
|
|
is recorded in `seen` (path -> line numbers), because a line the model read through a tool was shown to it."""
|
|
root = Path(root)
|
|
seen = seen if seen is not None else {}
|
|
|
|
def note(path: str, lines: Iterable[int]) -> None:
|
|
seen.setdefault(path, set()).update(int(line) for line in lines)
|
|
|
|
def read_file(path: str) -> str:
|
|
"""Read a file of the repository by its relative path, numbered lines."""
|
|
text = hunter_tools.read_file(root, path, max_chars=max_chars)
|
|
note(path, range(1, text.count("\n") + 2))
|
|
return text
|
|
|
|
def grep_repo(pattern: str) -> str:
|
|
"""Search the repository for a regular expression; returns path:line: text, one per line."""
|
|
rows = hunter_tools.grep_repo(root, pattern, max_matches=max_matches)
|
|
for row in rows:
|
|
note(str(row.get("path")), [int(row.get("line") or 0)])
|
|
return "\n".join(f"{row.get('path')}:{row.get('line')}: {row.get('text', '')}" for row in rows) or "(no matches)"
|
|
|
|
def find_refs(symbol: str) -> str:
|
|
"""Find where a symbol is referenced in the repository; returns path:line per reference."""
|
|
rows = hunter_tools.find_refs(root, symbol)
|
|
for row in rows:
|
|
note(str(row.get("path")), [int(row.get("line") or 0)])
|
|
return "\n".join(f"{row.get('path')}:{row.get('line')}" for row in rows) or "(no references)"
|
|
|
|
return [dspy.Tool(read_file), dspy.Tool(grep_repo), dspy.Tool(find_refs)]
|
|
|
|
|
|
class _Loop(dspy.ReAct):
|
|
"""ReAct where a step the adapter cannot parse ends the loop the way an invalid tool choice does, so the answer
|
|
is still extracted from what the loop gathered. DSPy 3.4 raises AdapterParseError, which is not the ValueError
|
|
its loop catches, so one malformed step otherwise discards the whole hunt."""
|
|
|
|
loop_parse_failed = False
|
|
|
|
def _call_with_potential_trajectory_truncation(
|
|
self, module: dspy.Module, trajectory: dict[str, Any], **input_args: str
|
|
) -> dspy.Prediction:
|
|
try:
|
|
return super()._call_with_potential_trajectory_truncation(module, trajectory, **input_args)
|
|
except dspy.utils.exceptions.AdapterParseError as error:
|
|
if module is self.extract:
|
|
raise
|
|
self.loop_parse_failed = True
|
|
raise ValueError("unparseable loop step") from error
|
|
|
|
|
|
class HuntProgram(dspy.Module):
|
|
"""One hunt over one change. With tools it is a bounded ReAct loop; without, one prediction (tests, dry runs)."""
|
|
|
|
def __init__(self, *, mask: Iterable[str] = (), tools: Sequence[dspy.Tool] = (), max_iters: int = 8) -> None:
|
|
super().__init__()
|
|
self.mask = tuple(sorted(set(mask)))
|
|
self.seen: dict[str, set[int]] = {}
|
|
signature = signature_for(self.mask)
|
|
self.hunt = _Loop(signature, tools=list(tools), max_iters=max_iters) if tools else dspy.Predict(signature)
|
|
|
|
def bind(self, root: Path) -> None:
|
|
"""Point the tools at one checkout and start a fresh record of what they show."""
|
|
self.seen = {}
|
|
if isinstance(self.hunt, dspy.ReAct):
|
|
for tool in repository_tools(root, seen=self.seen):
|
|
self.hunt.tools[tool.name] = tool
|
|
|
|
def forward(self, **inputs: str) -> dspy.Prediction:
|
|
if isinstance(self.hunt, _Loop):
|
|
self.hunt.loop_parse_failed = False
|
|
try:
|
|
prediction = self.hunt(**inputs)
|
|
except dspy.utils.exceptions.AdapterParseError:
|
|
# An answer the program could not shape is an empty answer that scores against it, not a broken run.
|
|
return dspy.Prediction(findings=[], raw=[], parse_failed=True, loop_parse_failed=self._loop_failed())
|
|
findings = check(getattr(prediction, "findings", []) or [], inputs, self.seen)
|
|
raw = getattr(prediction, "findings", [])
|
|
return dspy.Prediction(findings=findings, raw=raw, parse_failed=False, loop_parse_failed=self._loop_failed())
|
|
|
|
def _loop_failed(self) -> bool:
|
|
return bool(getattr(self.hunt, "loop_parse_failed", False))
|
|
|
|
|
|
def program_id(program: dspy.Module, *, mask: Iterable[str] = ()) -> str:
|
|
"""Content address of a program: its saved state plus the mask, so two identical programs share an id."""
|
|
state = program.dump_state()
|
|
payload = json.dumps({"state": state, "mask": sorted(set(mask))}, sort_keys=True, default=str)
|
|
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def build_lm(
|
|
model: str = DEFAULT_DETECTOR_MODEL, *, cache_dir: Path | None = None, temperature: float = 0.0, max_tokens: int = 4096
|
|
) -> dspy.LM:
|
|
"""The research-side model client: DSPy's LM over the DeepSeek endpoint, prices registered from our table.
|
|
|
|
The key is read from the environment only; it is never written anywhere. A cache directory per run keeps a
|
|
replay of every response next to the run's artifact."""
|
|
key = os.environ.get("DEEPSEEK_API_KEY")
|
|
if not key:
|
|
raise RuntimeError("DEEPSEEK_API_KEY is not set")
|
|
if cache_dir is not None:
|
|
# A run's own response cache, kept with its records so the run can be replayed at no cost.
|
|
Path(cache_dir).mkdir(parents=True, exist_ok=True)
|
|
dspy.configure_cache(enable_disk_cache=True, enable_memory_cache=True, disk_cache_dir=str(cache_dir))
|
|
try:
|
|
import litellm
|
|
|
|
prices = _PRICES.get(model)
|
|
if prices is not None:
|
|
litellm.register_model(
|
|
{
|
|
f"openai/{model}": {
|
|
"input_cost_per_token": prices.input_per_m / 1_000_000,
|
|
"output_cost_per_token": prices.output_per_m / 1_000_000,
|
|
"litellm_provider": "openai",
|
|
"mode": "chat",
|
|
}
|
|
}
|
|
)
|
|
except Exception as error: # noqa: BLE001 (a missing price only means cost is reported as unknown)
|
|
print(f"price registration skipped: {type(error).__name__}", flush=True)
|
|
return dspy.LM(f"openai/{model}", api_base=DEEPSEEK_BASE_URL, api_key=key, temperature=temperature, max_tokens=max_tokens, cache=True)
|