Open_Ultra_SAST/src/openultrasast/tool_hunter.py

"""Bounded tool-using hunter over map hotspots. Findings are suspicion only."""

from __future__ import annotations

import json
import os
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,
    entry_points,
    find_refs,
    flows,
    grep_repo,
    obligations,
    read_definition,
    read_file,
)
from .provider.openrouter import OpenRouterChatClient, OpenRouterError

_DEFAULT_MAX_CHARS = 4000
_DEFAULT_MAX_MATCHES = 50
DEFAULT_MAX_STEPS = 4
CLIENT_ENV = "OPENULTRASAST_HUNTER_CLIENT"
_FINDING_ID_PREFIX = "tool-hunter:"
DEFAULT_TOOL_NAMES = ("read_file", "grep_repo", "find_refs")
_HUNTER_TOOL_NAMES = frozenset(DEFAULT_TOOL_NAMES) | {"read_definition", "entry_points", "flows", "obligations"}
_SEVERITIES = frozenset({"critical", "high", "medium", "low", "info"})
_SCRIPTED_FLAGS = frozenset({"scripted", "script"})
_DUMP_FLAGS = frozenset({"dump", "dump-only"})
_UNSAFE_FLAGS = frozenset({"unsafe", "unsafe-snippet"})
_OPENROUTER_FLAGS = frozenset({"openrouter", "live"})
_UNSAFE_SNIPPET = 'client = docker.DockerClient(base_url="unix://var/run/docker.sock")\n'
_SCRIPTED_FINDINGS = [
    {
        "path": "app.py",
        "line": 3,
        "title": "Dynamic execution of attacker-controlled data",
        "rationale": "eval/execute follows request input across the hotspot.",
        "evidence_level": "static_corroboration",
    }
]

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"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_definition",
            "description": "Resolve one named function to its file, line span and body. Use it instead of guessing what a callee does.",
            "parameters": {
                "type": "object",
                "properties": {"symbol": {"type": "string"}, "path": {"type": "string"}, "max_chars": {"type": "integer"}},
                "required": ["symbol"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "entry_points",
            "description": "List the handlers this repository exposes, with their access classification. Identifiers and line spans only.",
            "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "flows",
            "description": "List adjudicated source-to-sink records inside one function. Identifiers and line numbers only.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}, "function": {"type": "string"}},
                "required": ["path", "function"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "obligations",
            "description": "List what one function owes and has not discharged: the operation, the missing discharger and why it is owed.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}, "function": {"type": "string"}},
                "required": ["path", "function"],
            },
        },
    },
]

_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, ...] = ()
    reasoning: str | None = None  # provider chain-of-thought; replayed on tool turns, never shown or scored
    mean_logprob: float | None = None  # averaged token log-probability when the caller asked for it


class ChatClient(Protocol):
    def complete(
        self,
        *,
        model: str,
        messages: list[dict[str, object]],
        tools: list[dict[str, object]],
        timeout_seconds: int = 60,
        json_object: bool = False,  # the final answer turn asks for JSON; adapters without it ignore the flag
    ) -> ChatResponse:
        raise NotImplementedError


class ScriptedChatClient:
    """Deterministic ChatClient for tests and OPENULTRASAST_HUNTER_CLIENT injection."""

    def __init__(self, turns: list[ChatResponse], *, cycle: bool = False) -> None:
        self._turns = turns
        self._cycle = cycle  # a stub that runs out of answers turns every multi-run test into a silent zero
        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,
        **options: object,
    ) -> ChatResponse:
        # `options` carries the adapter's extras (json_object, logprobs) so a test can assert what was asked for.
        self.calls.append({"model": model, "messages": list(messages), "tools": tools, "timeout_seconds": timeout_seconds, **options})
        index = len(self.calls) - 1
        if index >= len(self._turns):
            if not self._cycle or not self._turns:
                return ChatResponse(content="")
            index %= len(self._turns)
        return self._turns[index]


class OpenRouterHunterClient:
    """Adapt OpenRouterChatClient to the tool-hunter ChatClient protocol."""

    def __init__(self, client: OpenRouterChatClient) -> None:
        self._client = client

    @classmethod
    def from_env(cls) -> OpenRouterHunterClient:
        return cls(OpenRouterChatClient.from_env())

    def complete(
        self,
        *,
        model: str,
        messages: list[dict[str, object]],
        tools: list[dict[str, object]],
        timeout_seconds: int = 60,
        json_object: bool = False,
    ) -> ChatResponse:
        del json_object  # OpenRouter routes to many providers; the prompt asks for JSON and the parser tolerates prose
        message = self._client.complete_chat(model=model, messages=messages, tools=tools, timeout_seconds=timeout_seconds)
        return chat_response_from_message(message)


def scripted_hunter_client(mode: str = "scripted") -> ScriptedChatClient:
    grep_turn = ChatResponse(
        tool_calls=(ToolCall(id="call_grep", name="grep_repo", arguments={"pattern": "eval|execute|query", "max_matches": 10}),)
    )
    if mode in _DUMP_FLAGS:
        return ScriptedChatClient([ChatResponse(content=json.dumps(_SCRIPTED_FINDINGS))])
    if mode in _UNSAFE_FLAGS:
        payload = [{**_SCRIPTED_FINDINGS[0], "snippet": _UNSAFE_SNIPPET}]
        return ScriptedChatClient([grep_turn, ChatResponse(content=json.dumps(payload))])
    return ScriptedChatClient([grep_turn, ChatResponse(content=json.dumps(_SCRIPTED_FINDINGS))])


def resolve_hunter_client() -> ChatClient | None:
    """Return a hunter client: scripted when asked, DeepSeek when its key is set, else OpenRouter.

    DeepSeek is the chat provider; OpenRouter carries embeddings. This resolved OpenRouter only, so with the
    project's keys the hunter could not reach the chat model at all -- every call went to an endpoint whose key
    is scoped to embeddings and came back 401.
    """
    flag = os.environ.get(CLIENT_ENV, "").strip().lower()
    if flag in _SCRIPTED_FLAGS:
        return scripted_hunter_client("scripted")
    if flag in _DUMP_FLAGS:
        return scripted_hunter_client("dump-only")
    if flag in _UNSAFE_FLAGS:
        return scripted_hunter_client("unsafe-snippet")
    if flag not in _OPENROUTER_FLAGS and os.environ.get("DEEPSEEK_API_KEY"):
        from .model.endpoint import DEEPSEEK_BASE_ENV, DEEPSEEK_BASE_URL, DEFAULT_DETECTOR_MODEL, ChatEndpoint, DeepSeekChatClient, price_of

        base_url = os.environ.get(DEEPSEEK_BASE_ENV, DEEPSEEK_BASE_URL).rstrip("/")
        # Priced, so a budget can see what it spends: unpriced, the meter read $0.00 for any number of calls.
        endpoint = ChatEndpoint(provider="deepseek", base_url=base_url, thinking=False, prices=price_of(DEFAULT_DETECTOR_MODEL))
        return DeepSeekChatClient(OpenRouterChatClient(api_key=os.environ["DEEPSEEK_API_KEY"], base_url=base_url), endpoint=endpoint)
    if flag in _OPENROUTER_FLAGS or (not flag and os.environ.get("OPENROUTER_API_KEY")):
        try:
            return OpenRouterHunterClient.from_env()
        except OpenRouterError:
            return None
    return None


def chat_response_from_message(message: Mapping[str, object]) -> ChatResponse:
    content = message.get("content")
    text = content if isinstance(content, str) else None
    raw_calls = message.get("tool_calls") or ()
    calls: list[ToolCall] = []
    if isinstance(raw_calls, Sequence) and not isinstance(raw_calls, str | bytes):
        for index, item in enumerate(raw_calls):
            parsed = _tool_call_from_openrouter(item, index)
            if parsed is not None:
                calls.append(parsed)
    reasoning = message.get("reasoning_content") or message.get("reasoning")
    return ChatResponse(
        content=text,
        tool_calls=tuple(calls),
        reasoning=reasoning if isinstance(reasoning, str) and reasoning else None,
        mean_logprob=_mean_logprob(message.get("logprobs")),
    )


def _mean_logprob(payload: object) -> float | None:
    """The mean token log-probability of a reply, used by the classifier to abstain when it is unsure."""
    if not isinstance(payload, Mapping):
        return None
    tokens = payload.get("content")
    if not isinstance(tokens, Sequence) or isinstance(tokens, str | bytes):
        return None
    values = [float(item["logprob"]) for item in tokens if isinstance(item, Mapping) and isinstance(item.get("logprob"), int | float)]
    return sum(values) / len(values) if values else None


def _tool_call_from_openrouter(item: object, index: int) -> ToolCall | None:
    if not isinstance(item, Mapping):
        return None
    function = item.get("function")
    if not isinstance(function, Mapping):
        return None
    name = function.get("name")
    if not isinstance(name, str) or not name:
        return None
    call_id = item.get("id")
    identifier = call_id if isinstance(call_id, str) and call_id else f"call_{index}"
    arguments = function.get("arguments", {})
    if isinstance(arguments, str):
        try:
            loaded = json.loads(arguments)
        except json.JSONDecodeError:
            loaded = {}
        parsed_args = loaded if isinstance(loaded, dict) else {}
    elif isinstance(arguments, Mapping):
        parsed_args = dict(arguments)
    else:
        parsed_args = {}
    return ToolCall(id=identifier, name=name, arguments=parsed_args)


def run_tool_hunter(
    root: Path,
    hotspots: list[Hotspot],
    *,
    client: ChatClient,
    model: str,
    max_steps: int,
    system_prompt: str = _SYSTEM_PROMPT,
    user_prompt: str | None = None,
    context_files: Sequence[str] = (),
    tools: Sequence[str] | None = None,
    max_chars: int = _DEFAULT_MAX_CHARS,
    tags: Sequence[str] = ("tool-hunter",),
) -> list[StaticFinding]:
    """One bounded hunt. Every keyword defaults to today's behaviour; a per-family detector overrides the
    prompt, the context it starts from, the tools it may call and the tags its findings carry."""
    if max_steps <= 0:
        return []
    schemas = _tool_schemas(tools)
    messages: list[dict[str, object]] = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": _user_message(root, hotspots, user_prompt, context_files, max_chars)},
    ]
    used_tools = False
    for _step in range(max_steps):
        response = client.complete(model=model, messages=messages, tools=schemas)
        tool_calls = tuple(response.tool_calls or ())
        if tool_calls:
            messages.append(_assistant_tool_message(response.content, tool_calls, response.reasoning))
            for call in tool_calls:
                if call.name in _HUNTER_TOOL_NAMES:
                    used_tools = True
                messages.append(_tool_result_message(root, call, max_chars))
            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 []
        findings = _findings_from_content(root, hotspots, content, tags)
        if findings or _looks_like_json(content):
            return findings
        return _final_answer(client, model, messages, schemas, content, root, hotspots, tags)
    return _final_answer(client, model, messages, schemas, "", root, hotspots, tags) if used_tools else []


def _final_answer(
    client: ChatClient,
    model: str,
    messages: list[dict[str, object]],
    schemas: list[dict[str, object]],
    prose: str,
    root: Path,
    hotspots: Sequence[Hotspot],
    tags: Sequence[str],
) -> list[StaticFinding]:
    """One more turn, asking the provider for JSON, when the hunt ended in prose or ran out of steps.

    Measured live against DeepSeek: the model spends its step budget on tools and then answers in a paragraph that
    names the sink, the argument and the line — which parses as nothing at all. A detector that found the bug and
    reported zero is indistinguishable from one that found nothing, and every family would baseline at zero for
    that reason alone. This turn formats an answer the model has already reached; it never asks for a new one, and
    a hunt that called no tool never gets here.
    """
    body = prose.strip()
    tail: list[dict[str, object]] = [{"role": "assistant", "content": body}] if body else []
    messages = [
        *messages,
        *tail,
        {
            "role": "user",
            "content": (
                "Answer now, as JSON only: an array of objects with path, line, title, rationale and family. "
                "Report only what you established with the tools; return [] if that is nothing."
            ),
        },
    ]
    # Plain text, no tools. In JSON mode with tools attached, DeepSeek answered the final turn with the format
    # itself (`{"type": "json_object"}`), wrapped the array (`{"type": "json_object", "value": [...]}`) or emitted
    # its tool-call markup as text -- and a hunt that had found the bug reported nothing. Measured on the v2 sink
    # verifier: 0 of 7 declared vulnerable functions confirmed in JSON mode, 5 of 7 asked plainly (6 of 7 at 8
    # steps). The parser already takes a bare or fenced array.
    del schemas
    response = client.complete(model=model, messages=messages, tools=[], json_object=False)
    return _findings_from_content(root, hotspots, response.content or "", tags)


def _looks_like_json(content: str) -> bool:
    stripped = content.strip()
    fenced = re.fullmatch(r"```(?:json)?\s*(.*?)\s*```", stripped, flags=re.DOTALL)
    if fenced is not None:
        stripped = fenced.group(1).strip()
    return stripped.startswith(("[", "{"))


def _tool_schemas(names: Sequence[str] | None) -> list[dict[str, object]]:
    wanted = tuple(names) if names is not None else DEFAULT_TOOL_NAMES
    by_name = {str(schema["function"]["name"]): schema for schema in HUNTER_TOOLS}  # type: ignore[index]
    return [by_name[name] for name in wanted if name in by_name]


def _user_message(root: Path, hotspots: Sequence[Hotspot], user_prompt: str | None, context_files: Sequence[str], max_chars: int) -> str:
    body = user_prompt if user_prompt is not None else _hotspot_prompt(hotspots)
    parts = [body]
    for relative in context_files:
        try:
            text = read_file(root, relative, max_chars=max_chars)
        except (PathEscapesRepo, OSError):
            continue
        parts.append(f"--- {relative}\n{text}")
    return "\n\n".join(parts)


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], reasoning: str | None = None) -> dict[str, object]:
    message: dict[str, object] = {
        "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
        ],
    }
    if reasoning:
        # The provider drops its own chain of thought unless it comes back on the next turn (DeepSeek tool calls).
        message["reasoning_content"] = reasoning
    return message


def _tool_result_message(root: Path, call: ToolCall, max_chars: int = _DEFAULT_MAX_CHARS) -> dict[str, object]:
    return {
        "role": "tool",
        "tool_call_id": call.id,
        "name": call.name,
        "content": _run_tool(root, call, max_chars),
    }


def _run_tool(root: Path, call: ToolCall, max_chars: int = _DEFAULT_MAX_CHARS) -> str:
    try:
        payload = _dispatch_tool(root, call.name, _arguments(call.arguments), max_chars)
    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], default_max_chars: int = _DEFAULT_MAX_CHARS) -> 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 == "read_definition":
        symbol = _require_str(arguments, "symbol")
        scope = arguments.get("path")
        found = read_definition(
            root,
            symbol,
            max_chars=_optional_int(arguments, "max_chars", default_max_chars),
            path=str(scope) if isinstance(scope, str) and scope else None,
        )
        return found if found is not None else {"symbol": symbol, "found": False}
    if name == "entry_points":
        scope = arguments.get("path")
        return entry_points(root, path=str(scope) if isinstance(scope, str) and scope else None)
    if name == "flows":
        return flows(root, path=_require_str(arguments, "path"), function=_require_str(arguments, "function"))
    if name == "obligations":
        return obligations(root, path=_require_str(arguments, "path"), function=_require_str(arguments, "function"))
    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, tags: Sequence[str] = ("tool-hunter",)
) -> 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"]
    elif isinstance(payload, dict) and isinstance(payload.get("value"), list):
        # DeepSeek in JSON mode wrapped the array it was asked for as `{"type": "json_object", "value": [...]}`.
        items = payload["value"]
    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, tags)
        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], tags: Sequence[str] = ("tool-hunter",)
) -> 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")
    snippet = item.get("snippet")
    proposed_snippet = snippet if isinstance(snippet, str) and snippet else None
    family = item.get("family")
    # The family is the only key the class-aware scorer credits, and it is the model's own answer: a fabricated
    # name is counted as noise there rather than silently dropped here.
    family_tags = [f"family:{family}"] if isinstance(family, str) and family.strip() else []
    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=[*tags, *family_tags],
        ranking_priority=float(scores.get(relative, 0.0)),
        proposed_snippet=proposed_snippet,
    )


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()