@hoanghainh/komagu 2.0.0 → 2.0.1

package/operate/tests/test_ai_anchoring.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Finding locations anchored to Tier 1 routes / candidates, and one retry on empty replies (#319)."""
+
+from __future__ import annotations
+
+import json
+import shutil
+import tempfile
+import unittest
+from pathlib import Path
+
+from operate.ai.anchoring import AnchorReport, anchor_findings, load_anchors
+from operate.ai.runner import HeadlessAIRunner
+
+
+def _write(path: Path, data: object) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(data) if not isinstance(data, str) else data, encoding="utf-8")
+
+
+class TestAnchoring(unittest.TestCase):
+    def setUp(self) -> None:
+        self.tmp = Path(tempfile.mkdtemp(prefix="komagu_anchor_"))
+        self.target, self.out = self.tmp / "repo", self.tmp / "out"
+        _write(self.target / "node-express/routes/admin.js", "\n".join(f"// line {i}" for i in range(1, 21)) + "\n")
+        _write(self.target / "python-flask/profile.py", "\n".join(f"# line {i}" for i in range(1, 31)) + "\n")
+        _write(self.out / "routes/routes-001.json", {"routes": [
+            {"route_id": "R-1", "handler": {"file": "node-express/routes/admin.js", "line": 8}},
+            {"route_id": "R-2", "handler": {"file": "node-express/routes/admin.js", "line": 12}},
+            {"route_id": "R-3", "handler": {"file": "node-express/routes/admin.js", "line": 16}}]})
+        _write(self.out / "mass-assignment.candidates.json", {"candidates": [
+            {"candidate_id": "M-14", "file": "python-flask/profile.py", "line": 10}]})
+        self.anchors = load_anchors(self.out, self.target)
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmp, ignore_errors=True)
+
+    def run_one(self, finding: dict) -> tuple:
+        report = AnchorReport()
+        return anchor_findings([finding], self.anchors, self.target, report)[0], report
+
+    def test_handler_spans_end_before_the_next_route(self) -> None:
+        self.assertEqual((self.anchors["R-1"].line, self.anchors["R-1"].end_line), (8, 11))
+        self.assertEqual((self.anchors["R-3"].line, self.anchors["R-3"].end_line), (16, 20))
+
+    def test_wrong_path_with_line_in_span_keeps_the_line(self) -> None:
+        f, report = self.run_one({"finding_id": "BFLA-R-1-1", "route_id": "R-1",
+                                  "location": {"file": "routes.js", "start_line": 9},
+                                  "evidence": [{"file": "routes.js", "line": 9, "why": "no role check"}]})
+        self.assertEqual(f["location"], {"file": "node-express/routes/admin.js", "start_line": 9})
+        self.assertEqual(f["evidence"], [{"file": "node-express/routes/admin.js", "line": 9, "why": "no role check"}])
+        self.assertEqual(report.moved, ["BFLA-R-1-1"])
+
+    def test_line_outside_span_falls_back_to_handler(self) -> None:
+        f, _ = self.run_one({"finding_id": "BFLA-R-1-1", "route_id": "R-1",
+                             "location": {"file": "node-express/routes/admin.js", "start_line": 17}})
+        self.assertEqual(f["location"]["start_line"], 8)
+
+    def test_candidate_finding_takes_candidate_site(self) -> None:
+        f, _ = self.run_one({"finding_id": "MASS-M-14-1", "candidate_id": "M-14",
+                             "location": {"file": "app.py", "start_line": 8, "end_line": 9},
+                             "evidence": [{"file": "app.py", "line": 8, "why": "binds request body"}]})
+        self.assertEqual(f["location"], {"file": "python-flask/profile.py", "start_line": 10})
+        self.assertEqual(f["evidence"][0]["file"], "python-flask/profile.py")
+
+    def test_correct_location_is_not_reported_as_moved(self) -> None:
+        f, report = self.run_one({"finding_id": "BFLA-R-2-1", "route_id": "R-2",
+                                  "location": {"file": "node-express/routes/admin.js", "start_line": 13}})
+        self.assertEqual(f["location"]["start_line"], 13)
+        self.assertEqual((report.moved, report.unanchored), ([], []))
+
+    def test_unknown_reference_is_unanchored_and_low_confidence(self) -> None:
+        f, report = self.run_one({"finding_id": "BFLA-R-9-1", "route_id": "R-9", "confidence": "high",
+                                  "location": {"file": "x.js", "start_line": 1}})
+        self.assertEqual(f["location"], {"file": "x.js", "start_line": 1})
+        self.assertEqual(f["confidence"], "low")
+        self.assertEqual(report.unanchored, ["BFLA-R-9-1"])
+
+    def test_input_is_not_mutated(self) -> None:
+        original = {"finding_id": "BFLA-R-1-1", "route_id": "R-1", "location": {"file": "routes.js", "start_line": 9}}
+        snapshot = json.dumps(original, sort_keys=True)
+        self.run_one(original)
+        self.assertEqual(json.dumps(original, sort_keys=True), snapshot)
+
+
+class TestModelSource(unittest.TestCase):
+    """The privilege auditor sees the bound model's fields, from the same app and language (#319)."""
+
+    def setUp(self) -> None:
+        from operate.ai.runner import _model_source
+        self.model_source = _model_source
+        self.tmp = Path(tempfile.mkdtemp(prefix="komagu_model_"))
+        _write(self.tmp / "java-legacy/src/model/Account.java", "public class Account {\n    private String role;\n}\n")
+        _write(self.tmp / "java-legacy/src/web/ProfileServlet.java",
+               "class ProfileServlet {\n  void post() {\n    Account account = new Account();\n    BeanUtils.populate(account, params);\n  }\n}\n")
+        _write(self.tmp / "php-laravel/app/Models/User.php", "<?php\nclass User extends Model {\n  protected $guarded = [];\n}\n")
+        _write(self.tmp / "node-express/routes/profile.js", "router.post('/p', (req, res) => User.update(req.body));\n")
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmp, ignore_errors=True)
+
+    def test_target_class_is_found_with_its_fields(self) -> None:
+        m = self.model_source(self.tmp, "Account", "java-legacy/src/web/ProfileServlet.java", 4)
+        self.assertEqual(m["file"], "java-legacy/src/model/Account.java")
+        self.assertIn("private String role", m["code"])
+
+    def test_class_is_inferred_from_new_near_the_binding(self) -> None:
+        m = self.model_source(self.tmp, None, "java-legacy/src/web/ProfileServlet.java", 4)
+        self.assertEqual((m["class"], m["file"]), ("Account", "java-legacy/src/model/Account.java"))
+
+    def test_other_app_or_language_is_not_used(self) -> None:
+        m = self.model_source(self.tmp, "User", "node-express/routes/profile.js", 1)
+        self.assertEqual((m["class"], m["file"]), ("User", None))
+        self.assertIn("not found", m["code"])
+
+    def test_unknown_target_without_new_gives_none(self) -> None:
+        self.assertIsNone(self.model_source(self.tmp, None, "node-express/routes/profile.js", 1))
+        self.assertIsNone(self.model_source(self.tmp, "../etc", "node-express/routes/profile.js", 1))
+
+
+class _ScriptedClient:
+    """Returns the scripted replies in order; an Exception item is raised."""
+
+    model = "scripted"
+    endpoint = "test://"
+
+    def __init__(self, replies: list) -> None:
+        self.replies = list(replies)
+        self.calls: list = []
+
+    def chat(self, messages, temperature=0.1):  # noqa: ANN001
+        self.calls.append(messages)
+        item = self.replies.pop(0)
+        if isinstance(item, Exception):
+            raise item
+        return item
+
+    @staticmethod
+    def extract_json(content: str):  # noqa: ANN205
+        from operate.ai.client import LLMClient
+        return LLMClient.extract_json(content)
+
+
+class TestRetry(unittest.TestCase):
+    def setUp(self) -> None:
+        self.tmp = Path(tempfile.mkdtemp(prefix="komagu_retry_"))
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmp, ignore_errors=True)
+
+    def runner(self, replies: list) -> HeadlessAIRunner:
+        return HeadlessAIRunner(self.tmp, self.tmp / "out", _ScriptedClient(replies))
+
+    def test_empty_reply_is_retried_once(self) -> None:
+        r = self.runner(["", '{"findings": []}'])
+        self.assertEqual(r._chat("4a-bola-auditor", "routes-001", [{"role": "user", "content": "audit"}]), '{"findings": []}')
+        self.assertEqual(r.retries, {"4a-bola-auditor": ["routes-001"]})
+        retry = r.client.calls[1]
+        self.assertEqual([m["role"] for m in retry], ["user", "user"])  # no empty assistant turn
+        self.assertIn("JSON object only", retry[-1]["content"])
+
+    def test_unparseable_reply_is_retried_with_its_text(self) -> None:
+        r = self.runner(["looks fine to me", '{"findings": []}'])
+        r._chat("4c-privilege-auditor", "privilege-001", [{"role": "user", "content": "audit"}])
+        self.assertEqual([m["role"] for m in r.client.calls[1]], ["user", "assistant", "user"])
+
+    def test_second_bad_reply_keeps_the_first_text(self) -> None:
+        r = self.runner(["looks fine to me", "still not json"])
+        self.assertEqual(r._chat("5-finding-verifier", "verified-1", [{"role": "user", "content": "v"}]), "looks fine to me")
+
+    def test_good_reply_is_not_retried(self) -> None:
+        r = self.runner(['{"findings": []}'])
+        r._chat("4a-bola-auditor", "routes-001", [{"role": "user", "content": "audit"}])
+        self.assertEqual((len(r.client.calls), r.retries), (1, {}))
+
+    def test_failed_call_is_not_retried(self) -> None:
+        r = self.runner([RuntimeError("timed out")])
+        self.assertIsNone(r._chat("4b-injection-auditor", "injection-001", [{"role": "user", "content": "i"}]))
+        self.assertEqual((len(r.client.calls), r.retries), (1, {}))
+        self.assertEqual(r.call_errors, {"4b-injection-auditor": {"injection-001": "timed out"}})
+
+
+if __name__ == "__main__":
+    unittest.main()

package/operate/tests/test_headless_ai.py

@@ -74,6 +74,67 @@ class TestLLMClient(unittest.TestCase):
         self.assertIn("Failed to connect to LLM endpoint", str(ctx.exception))
 
 
+class TestSanitizeFinding(unittest.TestCase):
+    """Format repairs for model output that would otherwise fail finding.schema.json (#286)."""
+
+    def test_impact_list_becomes_string(self) -> None:
+        f = HeadlessAIRunner._sanitize_finding({"category": "mass_assignment", "subtype": "whole_object_binding",
+                                                "impact": ["role escalation", "balance change"]})
+        self.assertEqual(f["impact"], "role escalation; balance change")
+
+    def test_verdict_in_subtype_moves_to_status(self) -> None:
+        f = HeadlessAIRunner._sanitize_finding({"category": "mass_assignment", "subtype": "needs_review", "status": "unverified"})
+        self.assertEqual(f["subtype"], "whole_object_binding")
+        self.assertEqual(f["status"], "needs_review")
+
+    def test_unknown_bfla_subtype_falls_back(self) -> None:
+        f = HeadlessAIRunner._sanitize_finding({"category": "bfla", "subtype": "admin_only", "status": "unverified"})
+        self.assertEqual((f["subtype"], f["status"]), ("missing_role_check", "unverified"))
+
+    def test_null_evidence_line_and_missing_status_are_repaired(self) -> None:
+        f = HeadlessAIRunner._sanitize_finding({"category": "mass_assignment", "subtype": "whole_object_binding",
+                                                "location": {"file": "a.py", "start_line": 7},
+                                                "evidence": [{"file": "a.py", "line": None, "why": "binds body"}]})
+        self.assertEqual(f["evidence"][0]["line"], 7)
+        self.assertEqual(f["status"], "unverified")
+
+    def test_valid_subtype_is_kept(self) -> None:
+        f = HeadlessAIRunner._sanitize_finding({"category": "bfla", "subtype": "role_from_request"})
+        self.assertEqual(f["subtype"], "role_from_request")
+
+
+class TestVerificationContext(unittest.TestCase):
+    """The headless verifier gets the code it judges, including server-side includes (#286)."""
+
+    def setUp(self) -> None:
+        self.tmp = Path(tempfile.mkdtemp(prefix="komagu_verify_ctx_"))
+        (self.tmp / "inc").mkdir()
+        (self.tmp / "admin").mkdir()
+        (self.tmp / "inc" / "auth.php").write_text("<?php\n// Login check only\nsession_start();\n", encoding="utf-8")
+        (self.tmp / "admin" / "delete_user.php").write_text(
+            "<?php\nrequire_once '../inc/auth.php';\n$db->prepare('DELETE FROM users WHERE id = ?');\n", encoding="utf-8")
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmp, ignore_errors=True)
+
+    def test_included_files_are_resolved(self) -> None:
+        from operate.ai.runner import _included_files
+        self.assertEqual(_included_files(self.tmp, "admin/delete_user.php"), ["inc/auth.php"])
+        self.assertEqual(_included_files(self.tmp, "missing.php"), [])
+
+    def test_context_has_site_and_included_code(self) -> None:
+        from operate.ai.runner import _verification_context
+        ctx = _verification_context(self.tmp, {"location": {"file": "admin/delete_user.php", "start_line": 3}})
+        self.assertIn("DELETE FROM users", ctx["location"])
+        self.assertEqual(ctx["included_files"][0]["file"], "inc/auth.php")
+        self.assertIn("Login check only", ctx["included_files"][0]["code"])
+
+    def test_include_outside_target_is_ignored(self) -> None:
+        from operate.ai.runner import _included_files
+        (self.tmp / "x.php").write_text("<?php include '../../etc/passwd';\n", encoding="utf-8")
+        self.assertEqual(_included_files(self.tmp, "x.php"), [])
+
+
 class TestHeadlessAIRunner(unittest.TestCase):
     """End-to-end tests for HeadlessAIRunner and multi-agent pipeline."""
 
@@ -109,7 +170,7 @@ module.exports = app;
         # Mock LLM response dispatcher
         def mock_agent_responses(messages, model):
             user_msg = messages[-1]["content"]
-            if "Auth Model" in user_msg or "Pre-scan Stack Profile" in user_msg:
+            if ("Auth Model" in user_msg or "Pre-scan Stack Profile" in user_msg) and "Findings to Verify" not in user_msg:
                 return "# Mock Application Auth Model\n\n- Auth: Express session\n- Ownership: missing on /api/orders/:id\n"
             if "dynamic query string concatenation or interpolation" in user_msg:
                 return """```json
@@ -196,6 +257,7 @@ module.exports = app;
 ```"""
             return '{"status": "ok"}'
 
+        self.mock_agent_responses = mock_agent_responses
         self.server, self.thread, self.base_url = start_mock_llm_server(
             port=0, response_generator=mock_agent_responses
         )
@@ -204,6 +266,91 @@ module.exports = app;
         self.server.server_close()
         shutil.rmtree(self.temp_dir, ignore_errors=True)
 
+    def _prescan(self) -> None:
+        import subprocess
+
+        prescan_script = PROJECT_ROOT / "assess" / "lsa" / "scripts" / "run_prescan.py"
+        res = subprocess.run(
+            [sys.executable, str(prescan_script), str(self.target_dir), str(self.out_dir)],
+            capture_output=True,
+            text=True,
+        )
+        self.assertEqual(res.returncode, 0, f"Pre-scan failed: {res.stderr}")
+
+    def test_run_stats_record_tokens_and_phases(self) -> None:
+        self._prescan()
+        client = LLMClient(endpoint=self.base_url, model="mock-model", api_key="test")
+        HeadlessAIRunner(self.target_dir, self.out_dir, client).run()
+
+        stats = json.loads((self.out_dir / "ai-run-stats.json").read_text(encoding="utf-8"))
+        self.assertEqual(stats["model"], "mock-model")
+        self.assertGreater(stats["totals"]["calls"], 0)
+        # the mock server reports 100 prompt / 50 completion tokens per call
+        self.assertEqual(stats["totals"]["prompt_tokens"], 100 * stats["totals"]["calls"])
+        self.assertEqual(stats["totals"]["completion_tokens"], 50 * stats["totals"]["calls"])
+        self.assertEqual(sum(ph["calls"] for ph in stats["phases"].values()), stats["totals"]["calls"])
+        self.assertIn("5-finding-verifier", stats["phases"])
+        # the mock answers the BOLA prompt with the auth-model markdown; the retry gets plain JSON (#319)
+        self.assertEqual(stats["retries"], {"4a-bola-auditor": ["routes-001"]})
+        self.assertEqual(stats["parse_failures"], {})
+        self.assertEqual(stats["anchoring"]["unanchored"], [])
+
+    def test_unparseable_verifier_reply_leaves_findings_unverified(self) -> None:
+        """A broken verifier reply must not confirm findings (#286)."""
+        self.server.server_close()
+
+        verifier_prompts = []
+
+        def broken_verifier(messages, model):
+            if any("Findings to Verify & Challenge" in m["content"] for m in messages):
+                verifier_prompts.append(messages[-1]["content"])
+                return "I think these all look vulnerable."
+            return self.mock_agent_responses(messages, model)
+
+        self.server, self.thread, self.base_url = start_mock_llm_server(port=0, response_generator=broken_verifier)
+        self._prescan()
+        client = LLMClient(endpoint=self.base_url, model="mock-model", api_key="test")
+        logs = []
+        HeadlessAIRunner(self.target_dir, self.out_dir, client, step_callback=lambda st, msg: logs.append((st, msg))).run()
+
+        findings = json.loads((self.out_dir / "findings.json").read_text(encoding="utf-8"))["findings"]
+        self.assertGreaterEqual(len(findings), 1)
+        self.assertNotIn("confirmed", [f.get("status") for f in findings])
+        coverage = json.loads((self.out_dir / "coverage.json").read_text(encoding="utf-8"))
+        self.assertEqual(sorted(coverage["findings_without_verification"]), sorted(f["finding_id"] for f in findings))
+        stats = json.loads((self.out_dir / "ai-run-stats.json").read_text(encoding="utf-8"))
+        self.assertEqual(stats["parse_failures"]["5-finding-verifier"], ["verified-1"])
+        self.assertEqual(stats["retries"]["5-finding-verifier"], ["verified-1"])
+        # the verifier sees the code and the auth model, and is told not to assume unseen controls
+        self.assertEqual(len(verifier_prompts), 2, "one retry after the unparseable reply")
+        self.assertIn("Answer again with the JSON object only", verifier_prompts[1])
+        self.assertIn("code_context", verifier_prompts[0])
+        self.assertIn("Mock Application Auth Model", verifier_prompts[0])
+        self.assertIn("Never assume a control", verifier_prompts[0])
+        raw = self.out_dir / "ai-raw" / "5-finding-verifier-verified-1.txt"
+        self.assertEqual(raw.read_text(encoding="utf-8"), "I think these all look vulnerable.")
+        self.assertTrue(any(st == "warn" and "verified-1" in msg for st, msg in logs))
+
+    def test_failed_batch_call_is_recorded_and_run_continues(self) -> None:
+        """A timeout on one batch must not abort the run or lose earlier phases (#286)."""
+
+        class TimeoutOnInjection(LLMClient):
+            def chat(self, messages, temperature=0.1, max_tokens=None):  # noqa: ANN001
+                if "dynamic query string concatenation or interpolation" in messages[-1]["content"]:
+                    raise RuntimeError("timed out")
+                return super().chat(messages, temperature=temperature, max_tokens=max_tokens)
+
+        self._prescan()
+        client = TimeoutOnInjection(endpoint=self.base_url, model="mock-model", api_key="test")
+        result = HeadlessAIRunner(self.target_dir, self.out_dir, client).run()
+        self.assertEqual(result["status"], "success")
+
+        stats = json.loads((self.out_dir / "ai-run-stats.json").read_text(encoding="utf-8"))
+        self.assertEqual(stats["call_errors"], {"4b-injection-auditor": {"injection-001": "timed out"}})
+        self.assertNotIn("4b-injection-auditor", stats["parse_failures"])
+        self.assertIn("5-finding-verifier", stats["phases"])
+        self.assertTrue((self.out_dir / "findings.json").is_file())
+
     def test_headless_ai_runner_full_cycle(self) -> None:
         # 1. Run deterministic pre-scan to populate OUT
         prescan_script = PROJECT_ROOT / "assess" / "lsa" / "scripts" / "run_prescan.py"

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@hoanghainh/komagu",
-  "version": "2.0.0",
+  "version": "2.0.1",
   "description": "Security-First Legacy Modernization Platform CLI (LSA Kit, IAM Perimeter, Cross-User Testing)",
   "main": "remediate/authz/nodejs/src/index.js",
   "types": "remediate/authz/nodejs/src/index.d.ts",