@hoanghainh/komagu 1.8.0 → 2.0.1

package/operate/ai/anchoring.py

@@ -0,0 +1,130 @@
+"""Anchor model-written finding locations to the Tier 1 routes and candidates they refer to (#319).
+
+Small local models judge the right route or candidate but often write a wrong or non-existent path
+into the finding (Evaluation Run 4: `routes.js`, `app.py`, `dotnet-mvc/...`). Tier 1 already knows
+where every route handler and query / binding candidate is, so the location comes from there.
+
+Rules:
+- route findings (`route_id`): the file is the route handler's file; the model's line is kept when
+  it falls inside the handler span (handler line up to the next route in the same file, or the end
+  of the file), otherwise the handler line is used.
+- candidate findings (`candidate_id`): file and line are the candidate's.
+- evidence entries are re-pointed the same way; entries whose file does not exist and cannot be
+  re-pointed are dropped.
+- findings with an unknown `route_id` / `candidate_id` keep their location, get `confidence: low`
+  and are reported as unanchored.
+
+Standard-library only (Python 3.8+ compatible).
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+CANDIDATE_FILES = ("injection.candidates.json", "mass-assignment.candidates.json")
+
+
+@dataclass(frozen=True)
+class Anchor:
+    file: str
+    line: int
+    end_line: Optional[int]  # last line of the handler span; None for an exact candidate site
+
+
+@dataclass
+class AnchorReport:
+    moved: List[str] = field(default_factory=list)
+    unanchored: List[str] = field(default_factory=list)
+
+
+def _read(path: Path) -> Dict[str, Any]:
+    try:
+        doc = json.loads(path.read_text(encoding="utf-8"))
+        return doc if isinstance(doc, dict) else {}
+    except (OSError, ValueError):
+        return {}
+
+
+def _line_count(target_dir: Path, rel_file: str) -> int:
+    try:
+        return len((target_dir / rel_file).read_text(encoding="utf-8", errors="replace").splitlines())
+    except OSError:
+        return 0
+
+
+def load_anchors(out_dir: Path, target_dir: Path) -> Dict[str, Anchor]:
+    """route_id / candidate_id -> Anchor, from OUT/routes/*.json and the Tier 1 candidate files."""
+    handlers: List[Tuple[str, int, str]] = []
+    for path in sorted((out_dir / "routes").glob("*.json")):
+        for route in _read(path).get("routes") or []:
+            if not isinstance(route, dict):
+                continue
+            handler = route.get("handler") or {}
+            if route.get("route_id") and handler.get("file") and isinstance(handler.get("line"), int):
+                handlers.append((handler["file"], handler["line"], route["route_id"]))
+
+    anchors: Dict[str, Anchor] = {}
+    for rel_file, line, route_id in handlers:
+        later = sorted(ln for f, ln, _ in handlers if f == rel_file and ln > line)
+        end = later[0] - 1 if later else max(_line_count(target_dir, rel_file), line)
+        anchors[route_id] = Anchor(rel_file, line, end)
+
+    for name in CANDIDATE_FILES:
+        for cand in _read(out_dir / name).get("candidates") or []:
+            if isinstance(cand, dict) and cand.get("candidate_id") and cand.get("file") and isinstance(cand.get("line"), int):
+                anchors[cand["candidate_id"]] = Anchor(cand["file"], cand["line"], None)
+    return anchors
+
+
+def _anchored_line(anchor: Anchor, line: Any) -> int:
+    if anchor.end_line is None or not isinstance(line, int):
+        return anchor.line
+    return line if anchor.line <= line <= anchor.end_line else anchor.line
+
+
+def _repoint_evidence(evidence: Any, anchor: Anchor, target_dir: Path) -> List[Dict[str, Any]]:
+    kept = []
+    for ev in evidence if isinstance(evidence, list) else []:
+        if not isinstance(ev, dict):
+            continue
+        rel = str(ev.get("file") or "")
+        if rel == anchor.file or (rel and (target_dir / rel).is_file()):
+            kept.append(ev)
+        elif anchor.end_line is not None and isinstance(ev.get("line"), int) and anchor.line <= ev["line"] <= anchor.end_line:
+            kept.append({**ev, "file": anchor.file})
+    return kept or [{"file": anchor.file, "line": anchor.line, "why": "Location taken from the Tier 1 inventory"}]
+
+
+def anchor_finding(finding: Dict[str, Any], anchors: Dict[str, Anchor], target_dir: Path) -> Tuple[Dict[str, Any], Optional[bool]]:
+    """Return (finding with Tier 1 location, moved?); moved is None when the finding has no known anchor."""
+    key = finding.get("route_id") or finding.get("candidate_id")
+    anchor = anchors.get(key) if key else None
+    if anchor is None:
+        return {**finding, "confidence": "low"}, None
+    loc = finding.get("location") if isinstance(finding.get("location"), dict) else {}
+    # a wrong path with a line inside the handler span is most likely the right line in the wrong file
+    new_loc = {**loc, "file": anchor.file, "start_line": _anchored_line(anchor, loc.get("start_line"))}
+    if anchor.end_line is None:
+        new_loc.pop("end_line", None)
+    moved = (loc.get("file"), loc.get("start_line")) != (new_loc["file"], new_loc["start_line"])
+    evidence = _repoint_evidence(finding.get("evidence"), anchor, target_dir)
+    return {**finding, "location": new_loc, "evidence": evidence}, moved
+
+
+def anchor_findings(findings: List[Dict[str, Any]], anchors: Dict[str, Anchor], target_dir: Path,
+                    report: AnchorReport) -> List[Dict[str, Any]]:
+    out = []
+    for f in findings:
+        if not isinstance(f, dict):
+            out.append(f)
+            continue
+        anchored, moved = anchor_finding(f, anchors, target_dir)
+        if moved is None:
+            report.unanchored.append(str(f.get("finding_id")))
+        elif moved:
+            report.moved.append(str(f.get("finding_id")))
+        out.append(anchored)
+    return out

package/operate/ai/client.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 
 import json
 import re
+import time
 import urllib.error
 import urllib.request
 from typing import Any, Dict, List, Optional
@@ -30,6 +31,8 @@ class LLMClient:
         self.model = model
         self.api_key = api_key or "ollama"
         self.timeout = timeout
+        # Running totals across chat() calls, read by the runner for ai-run-stats.json
+        self.usage: Dict[str, float] = {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0, "seconds": 0.0}
 
     def chat(
         self,
@@ -54,10 +57,12 @@ class LLMClient:
         }
 
         req = urllib.request.Request(self.endpoint, data=data, headers=headers, method="POST")
+        started = time.monotonic()
         try:
             with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                 body = resp.read().decode("utf-8")
                 res_json = json.loads(body)
+                self._record_usage(res_json.get("usage") or {}, time.monotonic() - started)
                 choices = res_json.get("choices", [])
                 if not choices:
                     raise RuntimeError(f"LLM API returned no choices: {body}")
@@ -69,6 +74,14 @@ class LLMClient:
         except urllib.error.URLError as e:
             raise RuntimeError(f"Failed to connect to LLM endpoint at {self.endpoint}: {e.reason}") from e
 
+    def _record_usage(self, usage: Dict[str, Any], seconds: float) -> None:
+        self.usage = {
+            "calls": self.usage["calls"] + 1,
+            "prompt_tokens": self.usage["prompt_tokens"] + int(usage.get("prompt_tokens") or 0),
+            "completion_tokens": self.usage["completion_tokens"] + int(usage.get("completion_tokens") or 0),
+            "seconds": round(self.usage["seconds"] + seconds, 3),
+        }
+
     @staticmethod
     def extract_json(content: str) -> Any:
         """Extract and parse JSON from Markdown code blocks or raw text string."""

package/operate/ai/context.py

@@ -0,0 +1,105 @@
+"""Code context for the headless agents: snippets, server-side includes and bound model classes.
+
+Standard-library only (Python 3.8+ compatible).
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+def _read_snippet(target_dir: Path, rel_file: str, line: int, radius: int = 15) -> str:
+    """Read a window of code around target line from the audited repository."""
+    full_path = target_dir / rel_file
+    if not full_path.is_file():
+        return f"[File not found: {rel_file}]"
+    try:
+        lines = full_path.read_text(encoding="utf-8", errors="replace").splitlines()
+        start = max(0, line - radius - 1)
+        end = min(len(lines), line + radius)
+        numbered = [f"{i + 1:4d} | {lines[i]}" for i in range(start, end)]
+        return "\n".join(numbered)
+    except Exception as e:
+        return f"[Error reading {rel_file}: {e}]"
+
+
+# Server-side includes whose content a verifier must see to judge a control (PHP, Classic ASP, ColdFusion)
+_INCLUDE_RE = re.compile(
+    r"(?:require|include)(?:_once)?\s*\(?\s*['\"]([^'\"]+)['\"]"
+    r"|<!--\s*#include\s+(?:file|virtual)\s*=\s*\"([^\"]+)\""
+    r"|<cfinclude\s+template\s*=\s*\"([^\"]+)\"",
+    re.IGNORECASE,
+)
+
+
+def _included_files(target_dir: Path, rel_file: str) -> List[str]:
+    """Repository-relative paths of files that rel_file includes, if they exist."""
+    target_dir = target_dir.resolve()
+    path = target_dir / rel_file
+    if not path.is_file():
+        return []
+    text = path.read_text(encoding="utf-8", errors="replace")
+    found: List[str] = []
+    for m in _INCLUDE_RE.finditer(text):
+        ref = next(g for g in m.groups() if g)
+        candidate = (path.parent / ref).resolve()
+        try:
+            rel = str(candidate.relative_to(target_dir))
+        except ValueError:
+            continue
+        if candidate.is_file() and rel not in found:
+            found.append(rel)
+    return found
+
+
+_MODEL_EXTENSIONS = (".java", ".cs", ".vb", ".php", ".py", ".js", ".ts", ".rb", ".cfc")
+_NEW_OBJECT_RE = re.compile(r"\bnew\s+([A-Z][A-Za-z0-9_]*)\s*\(")
+
+
+def _model_source(target_dir: Path, name: Optional[str], near_file: str = "", near_line: int = 1) -> Optional[Dict[str, Any]]:
+    """Definition of the bound model class, so the privilege auditor can see its fields (#319).
+
+    name is the Tier 1 binding target; when it is unknown, the first `new Foo(` near the binding is used.
+    """
+    if not name and near_file:
+        m = _NEW_OBJECT_RE.search(_read_snippet(target_dir, near_file, near_line, radius=4))
+        name = m.group(1) if m else None
+    if not name or not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name):
+        return None
+    class_re = re.compile(rf"\bclass\s+{name}\b")
+    # same application and language as the binding: a Node route must not pick up a PHP model of the same name
+    suffix = Path(near_file).suffix.lower()
+    suffixes = (suffix,) if suffix in _MODEL_EXTENSIONS else _MODEL_EXTENSIONS
+    parts = Path(near_file).parts
+    root = target_dir / parts[0] if len(parts) > 1 and (target_dir / parts[0]).is_dir() else target_dir
+    for path in sorted(root.rglob("*")):
+        if path.suffix.lower() not in suffixes or not path.is_file():
+            continue
+        lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
+        for i, text in enumerate(lines):
+            if class_re.search(text):
+                rel = str(path.relative_to(target_dir))
+                return {"class": name, "file": rel, "code": _read_snippet(target_dir, rel, i + 16, radius=15)}
+    return {"class": name, "file": None, "code": "[definition not found in the repository]"}
+
+
+def _verification_context(target_dir: Path, finding: Dict[str, Any]) -> Dict[str, Any]:
+    """Code the verifier needs: the finding's site, its file header, evidence lines and included files."""
+    loc = finding.get("location") or {}
+    rel = str(loc.get("file") or "")
+    line = int(loc.get("start_line") or 1)
+    ctx: Dict[str, Any] = {"location": _read_snippet(target_dir, rel, line, radius=20)}
+    if line > 40:
+        ctx["file_header"] = _read_snippet(target_dir, rel, 13, radius=12)
+    others = []
+    for ev in finding.get("evidence") or []:
+        if isinstance(ev, dict) and ev.get("file") and ev.get("file") != rel:
+            others.append({"file": ev["file"], "code": _read_snippet(target_dir, str(ev["file"]), int(ev.get("line") or 1), radius=6)})
+    if others:
+        ctx["evidence"] = others[:3]
+    includes = [{"file": inc, "code": _read_snippet(target_dir, inc, 30, radius=30)} for inc in _included_files(target_dir, rel)[:3]]
+    if includes:
+        ctx["included_files"] = includes
+    return ctx