feat(scan): add initial SAST pipeline

src/openultrasast/rank.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from typing import Protocol
+from pathlib import Path
+
+from .preprocess import FileTarget
+
+
+class RankingError(ValueError):
+    """Raised when ranker output does not match the expected schema."""
+
+
+@dataclass(frozen=True)
+class RankingScore:
+    path: str
+    surface: int
+    influence: int
+    reachability: int
+    priority: float
+    rationale: str
+    model_id: str | None
+    static_boosts: list[str]
+
+
+class RankerClient(Protocol):
+    def complete_json(self, *, model: str, messages: list[dict[str, str]], timeout_seconds: int = 60) -> object:
+        raise NotImplementedError
+
+
+def rank_targets(targets: list[FileTarget], model_rankings: object | None = None, model_id: str | None = None) -> list[RankingScore]:
+    model_by_path = parse_ranker_response(model_rankings, model_id) if model_rankings is not None else {}
+    rankings: list[RankingScore] = []
+    for target in targets:
+        heuristic = heuristic_rank(target)
+        override = model_by_path.get(target.path)
+        rankings.append(_merge_rankings(heuristic, override) if override is not None else heuristic)
+    return sorted(rankings, key=lambda item: (-item.priority, item.path))
+
+
+def rank_targets_with_model(
+    targets: list[FileTarget],
+    *,
+    client: RankerClient,
+    model: str,
+    chunk_size: int = 150,
+    timeout_seconds: int = 60,
+) -> list[RankingScore]:
+    if chunk_size < 1:
+        raise RankingError("chunk_size must be at least 1")
+    model_rankings: list[dict[str, object]] = []
+    for start in range(0, len(targets), chunk_size):
+        chunk = targets[start : start + chunk_size]
+        response = client.complete_json(
+            model=model,
+            messages=[
+                {"role": "system", "content": _ranker_system_prompt()},
+                {"role": "user", "content": json.dumps([_target_prompt_payload(target) for target in chunk], sort_keys=True)},
+            ],
+            timeout_seconds=timeout_seconds,
+        )
+        for ranking in parse_ranker_response(response, model).values():
+            model_rankings.append(asdict(ranking))
+    return rank_targets(targets, {"rankings": model_rankings}, model_id=model)
+
+
+def write_rankings(rankings: list[RankingScore], path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {"rankings": [asdict(ranking) for ranking in rankings]}
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+
+
+def heuristic_rank(target: FileTarget) -> RankingScore:
+    tags = set(target.tags)
+    surface = 1
+    influence = 1
+    reachability = 1
+    boosts: list[str] = []
+
+    if "memory_unsafe" in tags:
+        surface = max(surface, 3)
+        boosts.append("memory_unsafe")
+    if "parser" in tags or "deserialization" in tags:
+        surface = max(surface, 4)
+        reachability = max(reachability, 3)
+        boosts.append("input_parser")
+    if "network_entry" in tags:
+        reachability = max(reachability, 4)
+        boosts.append("network_entry")
+    if "auth_boundary" in tags:
+        surface = max(surface, 3)
+        reachability = max(reachability, 3)
+        boosts.append("auth_boundary")
+    if "crypto" in tags:
+        influence = max(influence, 4)
+        boosts.append("crypto")
+    if "syscall_entry" in tags or "filesystem_entry" in tags:
+        surface = max(surface, 3)
+        boosts.append("privileged_boundary")
+    if target.has_fuzz_entry_point or "fuzzable" in tags:
+        surface = max(surface, 5)
+        reachability = max(reachability, 4)
+        boosts.append("fuzzable")
+    if target.loc >= 500:
+        influence = max(influence, 3)
+        boosts.append("large_file")
+
+    return RankingScore(
+        path=target.path,
+        surface=surface,
+        influence=influence,
+        reachability=reachability,
+        priority=composite_priority(surface, influence, reachability),
+        rationale=_heuristic_rationale(target, boosts),
+        model_id=None,
+        static_boosts=sorted(set(boosts)),
+    )
+
+
+def composite_priority(surface: int, influence: int, reachability: int) -> float:
+    return round(_clamp_score(surface) * 0.5 + _clamp_score(influence) * 0.2 + _clamp_score(reachability) * 0.3, 2)
+
+
+def parse_ranker_response(payload: object, model_id: str | None) -> dict[str, RankingScore]:
+    items = _ranking_items(payload)
+    rankings: dict[str, RankingScore] = {}
+    for item in items:
+        if not isinstance(item, dict):
+            raise RankingError("each ranking item must be an object")
+        path = item.get("path")
+        rationale = item.get("rationale", "model ranking")
+        if not isinstance(path, str) or not path:
+            raise RankingError("ranking item requires non-empty path")
+        if not isinstance(rationale, str) or not rationale:
+            raise RankingError(f"ranking item for {path} requires rationale")
+        surface = _required_score(item, "surface", path)
+        influence = _required_score(item, "influence", path)
+        reachability = _required_score(item, "reachability", path)
+        static_boosts = item.get("static_boosts", [])
+        if not isinstance(static_boosts, list):
+            raise RankingError(f"ranking item for {path} has invalid static_boosts")
+        rankings[path] = RankingScore(
+            path=path,
+            surface=surface,
+            influence=influence,
+            reachability=reachability,
+            priority=composite_priority(surface, influence, reachability),
+            rationale=rationale,
+            model_id=model_id,
+            static_boosts=sorted(str(boost) for boost in static_boosts),
+        )
+    return rankings
+
+
+def _ranking_items(payload: object) -> list[object]:
+    if isinstance(payload, list):
+        return payload
+    if isinstance(payload, dict):
+        rankings = payload.get("rankings")
+        if isinstance(rankings, list):
+            return rankings
+    raise RankingError("ranker response must be a list or object with rankings list")
+
+
+def _required_score(item: dict[str, object], field: str, path: str) -> int:
+    value = item.get(field)
+    if not isinstance(value, int) or isinstance(value, bool) or value < 1 or value > 5:
+        raise RankingError(f"ranking item for {path} requires integer {field} from 1 to 5")
+    return value
+
+
+def _merge_rankings(heuristic: RankingScore, model: RankingScore) -> RankingScore:
+    return RankingScore(
+        path=heuristic.path,
+        surface=max(heuristic.surface, model.surface),
+        influence=max(heuristic.influence, model.influence),
+        reachability=max(heuristic.reachability, model.reachability),
+        priority=composite_priority(
+            max(heuristic.surface, model.surface),
+            max(heuristic.influence, model.influence),
+            max(heuristic.reachability, model.reachability),
+        ),
+        rationale=model.rationale,
+        model_id=model.model_id,
+        static_boosts=sorted(set(heuristic.static_boosts) | set(model.static_boosts)),
+    )
+
+
+def _heuristic_rationale(target: FileTarget, boosts: list[str]) -> str:
+    if boosts:
+        return f"heuristic rank from tags: {', '.join(sorted(set(boosts)))}"
+    return f"baseline rank for {target.language} file"
+
+
+def _clamp_score(value: int) -> int:
+    return min(5, max(1, value))
+
+
+def _ranker_system_prompt() -> str:
+    return (
+        "Score each file for security review priority. Return JSON only with a rankings array. "
+        "Each item must include path, surface, influence, reachability, rationale, and optional static_boosts. "
+        "Scores must be integers from 1 to 5."
+    )
+
+
+def _target_prompt_payload(target: FileTarget) -> dict[str, object]:
+    return {
+        "path": target.path,
+        "language": target.language,
+        "loc": target.loc,
+        "tags": target.tags,
+        "has_fuzz_entry_point": target.has_fuzz_entry_point,
+    }