mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
feat(three-stage-scan): add tool-using hunter that records suspicion only
This commit is contained in:
@@ -62,7 +62,7 @@ This plan ships the three-stage scan without rebuilding CWE policy or the regex
|
||||
- _Boundary: ComplexityMapBuilder_
|
||||
- _Depends: 2.2_
|
||||
|
||||
- [ ] 3. Add repo-bound hunter tools and a scripted tool loop
|
||||
- [x] 3. Add repo-bound hunter tools and a scripted tool loop
|
||||
- [x] 3.1 (P) Clamp tool paths and expose read, grep, and reference search
|
||||
- Read, grep, and reference listing resolve only under the scan root; `..` and absolute escapes raise a path-escape error.
|
||||
- Grep runs with the standard library over enumerated source files.
|
||||
|
||||
@@ -0,0 +1,291 @@
|
||||
"""Bounded tool-using hunter over map hotspots. Findings are suspicion only."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from .complexity.map import Hotspot
|
||||
from .findings import StaticFinding
|
||||
from .hunter_tools import PathEscapesRepo, clamp_repo_path, find_refs, grep_repo, read_file
|
||||
|
||||
_DEFAULT_MAX_CHARS = 4000
|
||||
_DEFAULT_MAX_MATCHES = 50
|
||||
_FINDING_ID_PREFIX = "tool-hunter:"
|
||||
_HUNTER_TOOL_NAMES = frozenset({"read_file", "grep_repo", "find_refs"})
|
||||
_SEVERITIES = frozenset({"critical", "high", "medium", "low", "info"})
|
||||
|
||||
HUNTER_TOOLS: list[dict[str, object]] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file inside the scanned repository root.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Path relative to the repository root."},
|
||||
"max_chars": {"type": "integer"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "grep_repo",
|
||||
"description": "Search repository source text with a Python regular expression.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"pattern": {"type": "string"},
|
||||
"max_matches": {"type": "integer"},
|
||||
},
|
||||
"required": ["pattern"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "find_refs",
|
||||
"description": "List references to a symbol inside the scanned repository.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"symbol": {"type": "string"}},
|
||||
"required": ["symbol"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
_SYSTEM_PROMPT = (
|
||||
"You are a tool-using security hunter. Tools: read_file, grep_repo, find_refs. "
|
||||
"Stay inside the scanned repository. Use a tool before concluding. "
|
||||
"After using tools, reply with a JSON array of objects with path, line, title, and rationale. "
|
||||
"These findings are suspicions, not static corroboration."
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolCall:
|
||||
id: str
|
||||
name: str
|
||||
arguments: dict[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ChatResponse:
|
||||
content: str | None = None
|
||||
tool_calls: tuple[ToolCall, ...] = ()
|
||||
|
||||
|
||||
class ChatClient(Protocol):
|
||||
def complete(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]],
|
||||
timeout_seconds: int = 60,
|
||||
) -> ChatResponse:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def run_tool_hunter(
|
||||
root: Path,
|
||||
hotspots: list[Hotspot],
|
||||
*,
|
||||
client: ChatClient,
|
||||
model: str,
|
||||
max_steps: int,
|
||||
) -> list[StaticFinding]:
|
||||
if max_steps <= 0:
|
||||
return []
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "system", "content": _SYSTEM_PROMPT},
|
||||
{"role": "user", "content": _hotspot_prompt(hotspots)},
|
||||
]
|
||||
used_tools = False
|
||||
for _step in range(max_steps):
|
||||
response = client.complete(model=model, messages=messages, tools=HUNTER_TOOLS)
|
||||
tool_calls = tuple(response.tool_calls or ())
|
||||
if tool_calls:
|
||||
messages.append(_assistant_tool_message(response.content, tool_calls))
|
||||
for call in tool_calls:
|
||||
if call.name in _HUNTER_TOOL_NAMES:
|
||||
used_tools = True
|
||||
messages.append(_tool_result_message(root, call))
|
||||
continue
|
||||
content = response.content or ""
|
||||
if not content.strip() or not used_tools:
|
||||
# A findings dump is not evidence unless the loop invoked a repo tool.
|
||||
return []
|
||||
return _findings_from_content(root, hotspots, content)
|
||||
return []
|
||||
|
||||
|
||||
def _hotspot_prompt(hotspots: Sequence[Hotspot]) -> str:
|
||||
payload = [
|
||||
{
|
||||
"path": hotspot.path,
|
||||
"function_name": hotspot.function_name,
|
||||
"score": hotspot.score,
|
||||
"band": hotspot.band,
|
||||
"rationale": hotspot.rationale,
|
||||
}
|
||||
for hotspot in hotspots
|
||||
]
|
||||
return (
|
||||
"Inspect these complexity-map hotspots with the repository tools. "
|
||||
"Return JSON findings only after using a tool.\n"
|
||||
f"{json.dumps(payload, sort_keys=True)}"
|
||||
)
|
||||
|
||||
|
||||
def _assistant_tool_message(content: str | None, tool_calls: Sequence[ToolCall]) -> dict[str, object]:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": call.name,
|
||||
"arguments": json.dumps(call.arguments, sort_keys=True) if not isinstance(call.arguments, str) else call.arguments,
|
||||
},
|
||||
}
|
||||
for call in tool_calls
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _tool_result_message(root: Path, call: ToolCall) -> dict[str, object]:
|
||||
return {
|
||||
"role": "tool",
|
||||
"tool_call_id": call.id,
|
||||
"name": call.name,
|
||||
"content": _run_tool(root, call),
|
||||
}
|
||||
|
||||
|
||||
def _run_tool(root: Path, call: ToolCall) -> str:
|
||||
try:
|
||||
payload = _dispatch_tool(root, call.name, _arguments(call.arguments))
|
||||
except Exception as exc: # noqa: BLE001 — tool failures stay in-band so the scan continues
|
||||
payload = {"error": f"{type(exc).__name__}: {exc}"}
|
||||
return json.dumps(payload, sort_keys=True)
|
||||
|
||||
|
||||
def _dispatch_tool(root: Path, name: str, arguments: Mapping[str, object]) -> object:
|
||||
if name == "read_file":
|
||||
path = _require_str(arguments, "path")
|
||||
max_chars = _optional_int(arguments, "max_chars", _DEFAULT_MAX_CHARS)
|
||||
return {"path": path, "content": read_file(root, path, max_chars=max_chars)}
|
||||
if name == "grep_repo":
|
||||
pattern = _require_str(arguments, "pattern")
|
||||
max_matches = _optional_int(arguments, "max_matches", _DEFAULT_MAX_MATCHES)
|
||||
return grep_repo(root, pattern, max_matches=max_matches)
|
||||
if name == "find_refs":
|
||||
symbol = _require_str(arguments, "symbol")
|
||||
return find_refs(root, symbol, None)
|
||||
raise ValueError(f"unknown tool: {name}")
|
||||
|
||||
|
||||
def _arguments(raw: object) -> dict[str, object]:
|
||||
if isinstance(raw, str):
|
||||
parsed = json.loads(raw)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
raise ValueError("tool arguments must be a JSON object")
|
||||
if isinstance(raw, Mapping):
|
||||
return dict(raw)
|
||||
raise ValueError("tool arguments must be an object")
|
||||
|
||||
|
||||
def _require_str(arguments: Mapping[str, object], key: str) -> str:
|
||||
value = arguments.get(key)
|
||||
if not isinstance(value, str) or not value:
|
||||
raise ValueError(f"{key} is required")
|
||||
return value
|
||||
|
||||
|
||||
def _optional_int(arguments: Mapping[str, object], key: str, default: int) -> int:
|
||||
value = arguments.get(key, default)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return default
|
||||
return value
|
||||
|
||||
|
||||
def _findings_from_content(root: Path, hotspots: Sequence[Hotspot], content: str) -> list[StaticFinding]:
|
||||
payload = _load_json(content)
|
||||
if isinstance(payload, list):
|
||||
items = payload
|
||||
elif isinstance(payload, dict) and isinstance(payload.get("findings"), list):
|
||||
items = payload["findings"]
|
||||
else:
|
||||
return []
|
||||
scores = {hotspot.path: hotspot.score for hotspot in hotspots}
|
||||
findings: list[StaticFinding] = []
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
finding = _finding_from_item(root, item, scores)
|
||||
if finding is not None:
|
||||
findings.append(finding)
|
||||
return sorted(findings, key=lambda item: item.finding_id)
|
||||
|
||||
|
||||
def _load_json(content: str) -> object | None:
|
||||
stripped = content.strip()
|
||||
fenced = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", stripped, flags=re.DOTALL)
|
||||
if fenced is not None:
|
||||
stripped = fenced.group(1).strip()
|
||||
try:
|
||||
parsed: object = json.loads(stripped)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
return parsed
|
||||
|
||||
|
||||
def _finding_from_item(root: Path, item: Mapping[str, object], scores: Mapping[str, float]) -> StaticFinding | None:
|
||||
raw_path = item.get("path")
|
||||
if not isinstance(raw_path, str) or not raw_path:
|
||||
return None
|
||||
try:
|
||||
relative = _clamp_relative_path(root, raw_path)
|
||||
except PathEscapesRepo:
|
||||
return None
|
||||
line = item.get("line")
|
||||
line_no = line if isinstance(line, int) and not isinstance(line, bool) else None
|
||||
title = item.get("title")
|
||||
rationale = item.get("rationale")
|
||||
function_name = item.get("function_name")
|
||||
severity = item.get("severity")
|
||||
return StaticFinding(
|
||||
finding_id=f"{_FINDING_ID_PREFIX}{relative}:{line_no if line_no is not None else 0}",
|
||||
path=relative,
|
||||
title=title if isinstance(title, str) and title else "tool hunter suspicion",
|
||||
severity=severity if isinstance(severity, str) and severity in _SEVERITIES else "medium",
|
||||
confidence="low",
|
||||
evidence_level="suspicion",
|
||||
rationale=rationale if isinstance(rationale, str) and rationale else "Reported by the tool-using hunter.",
|
||||
line=line_no,
|
||||
function_name=function_name if isinstance(function_name, str) and function_name else None,
|
||||
reachability_status="unknown",
|
||||
reachability_evidence=[],
|
||||
reachability_conditions=[],
|
||||
tags=["tool-hunter"],
|
||||
ranking_priority=float(scores.get(relative, 0.0)),
|
||||
)
|
||||
|
||||
|
||||
def _clamp_relative_path(root: Path, user_path: str) -> str:
|
||||
resolved_root = root.resolve()
|
||||
return clamp_repo_path(resolved_root, user_path).relative_to(resolved_root).as_posix()
|
||||
@@ -0,0 +1,166 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from openultrasast.complexity.map import Hotspot
|
||||
from openultrasast.tool_hunter import ChatResponse, ToolCall, run_tool_hunter
|
||||
|
||||
SPLIT_SINK = """\
|
||||
term = request.args["q"]
|
||||
query = "select * from items where title like '%" + term + "%'"
|
||||
return db.execute(query)
|
||||
"""
|
||||
|
||||
_DUMP_FINDINGS = [
|
||||
{
|
||||
"path": "app.py",
|
||||
"line": 3,
|
||||
"title": "SQL injection",
|
||||
"rationale": "query is built from request input then executed",
|
||||
"evidence_level": "static_corroboration",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class ScriptedChatClient:
|
||||
def __init__(self, turns: list[ChatResponse]) -> None:
|
||||
self._turns = turns
|
||||
self.calls: list[dict[str, object]] = []
|
||||
|
||||
def complete(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]],
|
||||
timeout_seconds: int = 60,
|
||||
) -> ChatResponse:
|
||||
self.calls.append({"model": model, "messages": list(messages), "tools": tools, "timeout_seconds": timeout_seconds})
|
||||
index = len(self.calls) - 1
|
||||
if index >= len(self._turns):
|
||||
return ChatResponse(content="")
|
||||
return self._turns[index]
|
||||
|
||||
|
||||
def _repo_with_split_sink(tmp_path: Path) -> Path:
|
||||
root = tmp_path / "repo"
|
||||
root.mkdir()
|
||||
(root / "app.py").write_text(SPLIT_SINK)
|
||||
return root
|
||||
|
||||
|
||||
def _hotspot(path: str = "app.py") -> Hotspot:
|
||||
return Hotspot(
|
||||
path=path,
|
||||
function_name="search",
|
||||
score=7.5,
|
||||
band="high",
|
||||
signals={"loc": 3},
|
||||
rationale="reachable split-sink",
|
||||
test_hint=None,
|
||||
inventory_finding_ids=(),
|
||||
)
|
||||
|
||||
|
||||
def _tool_names(tools: object) -> set[str]:
|
||||
names: set[str] = set()
|
||||
if not isinstance(tools, list):
|
||||
return names
|
||||
for item in tools:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
function = item.get("function")
|
||||
if isinstance(function, dict) and isinstance(function.get("name"), str):
|
||||
names.add(str(function["name"]))
|
||||
elif isinstance(item.get("name"), str):
|
||||
names.add(str(item["name"]))
|
||||
return names
|
||||
|
||||
|
||||
def test_json_dump_without_tool_call_emits_no_findings(tmp_path: Path) -> None:
|
||||
root = _repo_with_split_sink(tmp_path)
|
||||
client = ScriptedChatClient([ChatResponse(content=json.dumps(_DUMP_FINDINGS))])
|
||||
|
||||
findings = run_tool_hunter(root, [_hotspot()], client=client, model="test-hunter", max_steps=4)
|
||||
|
||||
assert client.calls, "the hunter loop must run when a model client is provided"
|
||||
assert _tool_names(client.calls[0]["tools"]) >= {"read_file", "grep_repo", "find_refs"}
|
||||
assert findings == []
|
||||
|
||||
|
||||
def test_grep_then_json_yields_suspicion_with_tool_hunter_prefix(tmp_path: Path) -> None:
|
||||
root = _repo_with_split_sink(tmp_path)
|
||||
client = ScriptedChatClient(
|
||||
[
|
||||
ChatResponse(tool_calls=(ToolCall(id="call_grep", name="grep_repo", arguments={"pattern": r"query\s*=", "max_matches": 10}),)),
|
||||
ChatResponse(content=json.dumps(_DUMP_FINDINGS)),
|
||||
]
|
||||
)
|
||||
|
||||
findings = run_tool_hunter(root, [_hotspot()], client=client, model="test-hunter", max_steps=4)
|
||||
|
||||
assert len(client.calls) >= 2
|
||||
assert _tool_names(client.calls[0]["tools"]) >= {"read_file", "grep_repo", "find_refs"}
|
||||
tool_trace = json.dumps(client.calls[1]["messages"])
|
||||
assert "query =" in tool_trace
|
||||
assert len(findings) == 1
|
||||
assert findings[0].evidence_level == "suspicion"
|
||||
assert findings[0].finding_id.startswith("tool-hunter:")
|
||||
assert findings[0].path == "app.py"
|
||||
assert findings[0].line == 3
|
||||
assert findings[0].title == "SQL injection"
|
||||
|
||||
|
||||
def test_loop_dispatches_read_file_grep_repo_and_find_refs(tmp_path: Path) -> None:
|
||||
root = _repo_with_split_sink(tmp_path)
|
||||
client = ScriptedChatClient(
|
||||
[
|
||||
ChatResponse(
|
||||
tool_calls=(
|
||||
ToolCall(id="call_read", name="read_file", arguments={"path": "app.py", "max_chars": 4000}),
|
||||
ToolCall(id="call_grep", name="grep_repo", arguments={"pattern": r"db\.execute", "max_matches": 5}),
|
||||
ToolCall(id="call_refs", name="find_refs", arguments={"symbol": "query"}),
|
||||
)
|
||||
),
|
||||
ChatResponse(content=json.dumps(_DUMP_FINDINGS)),
|
||||
]
|
||||
)
|
||||
|
||||
findings = run_tool_hunter(root, [_hotspot()], client=client, model="test-hunter", max_steps=4)
|
||||
|
||||
assert findings
|
||||
assert findings[0].evidence_level == "suspicion"
|
||||
second_messages = json.dumps(client.calls[1]["messages"])
|
||||
assert "term = request.args" in second_messages
|
||||
assert "db.execute" in second_messages
|
||||
assert "query" in second_messages
|
||||
|
||||
|
||||
def test_path_escape_stays_in_band_and_still_counts_as_tool_use(tmp_path: Path) -> None:
|
||||
root = _repo_with_split_sink(tmp_path)
|
||||
client = ScriptedChatClient(
|
||||
[
|
||||
ChatResponse(tool_calls=(ToolCall(id="call_escape", name="read_file", arguments={"path": "../etc/passwd"}),)),
|
||||
ChatResponse(content=json.dumps(_DUMP_FINDINGS)),
|
||||
]
|
||||
)
|
||||
|
||||
findings = run_tool_hunter(root, [_hotspot()], client=client, model="test-hunter", max_steps=4)
|
||||
|
||||
assert findings
|
||||
assert findings[0].evidence_level == "suspicion"
|
||||
assert "error" in json.dumps(client.calls[1]["messages"]).lower()
|
||||
|
||||
|
||||
def test_max_steps_stops_before_ungrounded_json(tmp_path: Path) -> None:
|
||||
root = _repo_with_split_sink(tmp_path)
|
||||
client = ScriptedChatClient(
|
||||
[
|
||||
ChatResponse(tool_calls=(ToolCall(id="call_grep", name="grep_repo", arguments={"pattern": "query"}),)),
|
||||
ChatResponse(content=json.dumps(_DUMP_FINDINGS)),
|
||||
]
|
||||
)
|
||||
|
||||
findings = run_tool_hunter(root, [_hotspot()], client=client, model="test-hunter", max_steps=1)
|
||||
|
||||
assert findings == []
|
||||
assert len(client.calls) == 1
|
||||
Reference in New Issue
Block a user