diff --git a/.kiro/specs/three-stage-scan/tasks.md b/.kiro/specs/three-stage-scan/tasks.md index 3778b0e..e8559e7 100644 --- a/.kiro/specs/three-stage-scan/tasks.md +++ b/.kiro/specs/three-stage-scan/tasks.md @@ -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. diff --git a/src/openultrasast/tool_hunter.py b/src/openultrasast/tool_hunter.py new file mode 100644 index 0000000..338ace7 --- /dev/null +++ b/src/openultrasast/tool_hunter.py @@ -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() diff --git a/tests/test_tool_hunter.py b/tests/test_tool_hunter.py new file mode 100644 index 0000000..6840bd4 --- /dev/null +++ b/tests/test_tool_hunter.py @@ -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