x

tests/test_threat_model_tool.py

@@ -0,0 +1,309 @@
+"""Tests for the target-scoped threat model cache."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING
+
+import pytest
+
+from strix.agents.factory import _BASE_TOOLS
+from strix.tools.threat_model import tools as threat_model_tools
+from strix.tools.threat_model.tools import (
+    _amend_impl,
+    _get_impl,
+    _save_impl,
+    amend_threat_model,
+    get_threat_model,
+    save_threat_model,
+)
+
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+_MODEL = """# Threat Model
+
+## Overview
+A multi-tenant billing API. Product code lives in `api/`; `scripts/` is
+developer-only tooling and is not deployed.
+
+## Trust Boundaries and Assumptions
+Requests arrive from untrusted tenants through `api/router.py`. The tenant id
+is taken from the signed session, never from the request body. Operators
+configure webhooks; developers control migrations.
+
+## Attack Surface and Attacker Stories
+The public REST surface and the webhook receiver are attacker-reachable. A
+realistic story is a tenant reading another tenant's invoices. Local CLI
+tooling is not a realistic surface.
+
+## Severity Calibration
+Critical: cross-tenant write. High: cross-tenant read. Medium: authenticated
+self-scoped information leak. Low: verbose errors.
+"""
+
+
+def _git(repo: Path, *args: str) -> None:
+    subprocess.run(["/usr/bin/env", "git", *args], cwd=repo, check=True)  # noqa: S603
+
+
+def _make_repo(tmp_path: Path, name: str = "repo") -> Path:
+    repo = tmp_path / name
+    repo.mkdir(parents=True)
+    _git(repo, "init", "-q")
+    _git(repo, "config", "user.email", "t@example.com")
+    _git(repo, "config", "user.name", "t")
+    (repo / "README.md").write_text("hi\n", encoding="utf-8")
+    _git(repo, "add", "README.md")
+    _git(repo, "commit", "-qm", "init")
+    return repo
+
+
+@pytest.fixture(autouse=True)
+def _isolated_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(threat_model_tools, "_CACHE_DIR", tmp_path / "cache")
+
+
+def test_missing_model_reports_not_found(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+
+    result = _get_impl(str(repo))
+
+    assert result["success"] is True
+    assert result["found"] is False
+    assert "save_threat_model" in result["message"]
+
+
+def test_saved_model_round_trips(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+
+    assert _save_impl(str(repo), _MODEL, "Strix")["success"] is True
+    result = _get_impl(str(repo))
+
+    assert result["found"] is True
+    assert result["stale"] is False
+    assert "multi-tenant billing API" in result["content"]
+
+
+def test_model_is_stale_after_new_revision(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    _save_impl(str(repo), _MODEL, None)
+
+    (repo / "next.py").write_text("x = 1\n", encoding="utf-8")
+    _git(repo, "add", "next.py")
+    _git(repo, "commit", "-qm", "next")
+
+    result = _get_impl(str(repo))
+
+    assert result["found"] is True
+    assert result["stale"] is True
+    assert result["content"]
+
+
+def test_cache_is_keyed_per_repository(tmp_path: Path) -> None:
+    first = _make_repo(tmp_path, "first")
+    second = _make_repo(tmp_path, "second")
+    _save_impl(str(first), _MODEL, None)
+
+    assert _get_impl(str(second))["found"] is False
+
+
+def test_rejects_model_missing_required_sections(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    thin = _MODEL.replace("## Severity Calibration", "## Notes")
+
+    result = _save_impl(str(repo), thin, None)
+
+    assert result["success"] is False
+    assert "severity calibration" in result["error"]
+
+
+def test_rejects_stub_model(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+
+    result = _save_impl(str(repo), "overview trust boundaries attack surface", None)
+
+    assert result["success"] is False
+    assert "too thin" in result["error"]
+
+
+def test_rejects_empty_target() -> None:
+    result = _get_impl("   ")
+    assert result["success"] is False
+    assert "target cannot be empty" in result["error"]
+
+
+def test_tools_are_registered() -> None:
+    assert get_threat_model in _BASE_TOOLS
+    assert save_threat_model in _BASE_TOOLS
+
+
+_ADDENDUM = (
+    "The base model calls the webhook receiver operator-controlled. It is "
+    "unauthenticated in `api/webhooks.py:31`, so treat its body as attacker-controlled."
+)
+
+
+def test_amendment_is_returned_with_the_model(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    _save_impl(str(repo), _MODEL, "root")
+
+    assert _amend_impl(str(repo), _ADDENDUM, "webhook-agent")["success"] is True
+    result = _get_impl(str(repo))
+
+    assert result["content"] == _MODEL.strip()
+    assert [a["content"] for a in result["amendments"]] == [_ADDENDUM]
+    assert result["amendments"][0]["by"] == "webhook-agent"
+
+
+def test_amendments_accumulate_without_overwriting(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    _save_impl(str(repo), _MODEL, "root")
+
+    _amend_impl(str(repo), _ADDENDUM, "agent-a")
+    second = "The `scripts/` directory ships in the container image; it is not dev-only."
+    _amend_impl(str(repo), second + " See `Dockerfile:14`.", "agent-b")
+
+    amendments = _get_impl(str(repo))["amendments"]
+    assert len(amendments) == 2
+    assert [a["by"] for a in amendments] == ["agent-a", "agent-b"]
+
+
+def test_amend_requires_an_existing_model(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+
+    result = _amend_impl(str(repo), _ADDENDUM, None)
+
+    assert result["success"] is False
+    assert "save_threat_model" in result["error"]
+
+
+def test_amend_rejects_a_stub(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    _save_impl(str(repo), _MODEL, "root")
+
+    assert _amend_impl(str(repo), "looks wrong", None)["success"] is False
+
+
+def test_save_clears_amendments_and_says_so(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    _save_impl(str(repo), _MODEL, "root")
+    _amend_impl(str(repo), _ADDENDUM, "agent-a")
+
+    result = _save_impl(str(repo), _MODEL.replace("billing API", "billing service"), "root")
+
+    assert result["amendments_cleared"] == 1
+    assert "cleared" in result["message"]
+    assert "amendments" not in _get_impl(str(repo))
+
+
+def test_amend_tool_is_registered() -> None:
+    assert amend_threat_model in _BASE_TOOLS
+
+
+_BLACKBOX_MODEL = _MODEL.replace(
+    "Product code lives in `api/`; `scripts/` is\ndeveloper-only tooling and is not deployed.",
+    "Only the deployed surface is visible; no source. Inferred from recon.",
+)
+
+
+def test_blackbox_target_round_trips() -> None:
+    target = "https://app.example.com"
+
+    assert _save_impl(target, _BLACKBOX_MODEL, "recon")["success"] is True
+    result = _get_impl(target)
+
+    assert result["found"] is True
+    assert result["stale"] is False, "a fresh model with no revision is not stale"
+    assert result["revision"] == "unversioned"
+    assert "Inferred from recon" in result["content"]
+
+
+def test_blackbox_target_spellings_share_one_model() -> None:
+    _save_impl("https://App.Example.com:443/", _BLACKBOX_MODEL, "recon")
+
+    for spelling in ("https://app.example.com", "app.example.com", "https://app.example.com/"):
+        assert _get_impl(spelling)["found"] is True, spelling
+
+    assert _get_impl("https://other.example.com")["found"] is False
+
+
+def test_blackbox_model_goes_stale_with_age() -> None:
+    target = "https://app.example.com"
+    _save_impl(target, _BLACKBOX_MODEL, "recon")
+
+    aged = (datetime.now(UTC) - timedelta(days=threat_model_tools._MAX_AGE_DAYS + 1)).isoformat()
+    path = threat_model_tools._cache_path("app.example.com:443")
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    payload["created_at"] = aged
+    path.write_text(json.dumps(payload), encoding="utf-8")
+
+    result = _get_impl(target)
+
+    assert result["stale"] is True
+    assert "re-confirm" in result["message"]
+
+
+def test_blackbox_target_can_be_amended() -> None:
+    target = "https://app.example.com"
+    _save_impl(target, _BLACKBOX_MODEL, "recon")
+
+    addendum = (
+        "The model infers /admin is IP-restricted. It is reachable with any "
+        "authenticated session; the restriction is only on /admin/settings."
+    )
+    assert _amend_impl(target, addendum, "authz-agent")["success"] is True
+    assert _get_impl(target)["amendments"][0]["content"] == addendum
+
+
+def test_checkout_and_its_remote_are_the_same_target(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    _git(repo, "remote", "add", "origin", "https://github.com/acme/billing.git")
+    _save_impl(str(repo), _MODEL, "root")
+
+    clone = _make_repo(tmp_path, "clone")
+    _git(clone, "remote", "add", "origin", "https://github.com/acme/billing.git")
+
+    assert _get_impl(str(clone))["found"] is True
+
+
+def test_path_on_a_known_host_resolves_to_the_scan_target() -> None:
+    scan_targets = ["https://app.example.com"]
+    _save_impl("https://app.example.com", _BLACKBOX_MODEL, "root", scan_targets)
+
+    # An agent testing one page names that page, not the scan's target string.
+    assert _get_impl("https://app.example.com/admin/login", scan_targets)["found"] is True
+
+
+def test_two_scan_targets_on_one_host_stay_separate() -> None:
+    scan_targets = ["https://example.com/tenant-a", "https://example.com/tenant-b"]
+    _save_impl("https://example.com/tenant-a", _BLACKBOX_MODEL, "root", scan_targets)
+
+    assert _get_impl("https://example.com/tenant-b", scan_targets)["found"] is False
+
+
+def test_unknown_host_is_not_snapped_onto_the_scan_target() -> None:
+    scan_targets = ["https://app.example.com"]
+    _save_impl("https://app.example.com", _BLACKBOX_MODEL, "root", scan_targets)
+
+    assert _get_impl("https://unrelated.test", scan_targets)["found"] is False
+
+
+def test_empty_target_falls_back_to_a_single_scan_target() -> None:
+    scan_targets = ["https://app.example.com"]
+    _save_impl("", _BLACKBOX_MODEL, "root", scan_targets)
+
+    assert _get_impl("", scan_targets)["found"] is True
+    assert _get_impl("https://app.example.com")["found"] is True
+
+
+def test_repository_subdirectory_shares_the_repository_model(tmp_path: Path) -> None:
+    repo = _make_repo(tmp_path)
+    (repo / "src").mkdir()
+    _save_impl(str(repo), _MODEL, "root")
+
+    assert _get_impl(str(repo / "src"))["found"] is True