mirror of
https://github.com/usestrix/strix.git
synced 2026-08-28 05:25:00 +00:00
Scope threat models to the current run instead of caching them on disk
This commit is contained in:
parent
a5856108a7
commit
7d8d71beea
7 changed files with 129 additions and 256 deletions
|
|
@ -219,7 +219,7 @@ VALIDATION REQUIREMENTS:
|
|||
- CLOSURE DISCIPLINE: every candidate you open ends in exactly one explicit state — `confirmed` (working PoC, or a complete source→control→sink→impact trace that is reachable), `ruled_out` (you can name the SPECIFIC control, at a location, that runs on every attacker-reachable path before the sink), or `open_proof_gap` (plausible, unconfirmed, and you could NOT name such a control). "I moved on" is not a closure state. Silently dropping an uncertain candidate is mislabelling an `open_proof_gap` as `ruled_out` and is how real bugs get missed.
|
||||
- Missing information is NOT proof of safety: no caller found, can't tell if deployed/exposed, couldn't stand up the service, build failed — each is an `open_proof_gap`, never a reason to mark a candidate clean. Difficulty is a reason to defer, not to suppress.
|
||||
- COVERAGE: record every surface you assess with `record_coverage` (surface + risk area + outcome + evidence), including the ones that came back clean — a report that only lists findings cannot say what was reviewed and cleared. Use the `needs_follow_up` outcome for anything left in an `open_proof_gap` state, and carry the same items up in `agent_finish(open_items=[...])`. The ledger is shared and mutable: when you resolve a surface another agent left open — or find that a closed one is not — move that entry with `update_coverage` instead of recording a second one for the same surface. The root agent reconciles all of it via `list_coverage` before `finish_scan`.
|
||||
- THREAT MODEL: before you start testing, call `get_threat_model` on the target you were pointed at — it is the scan's shared answer to who the attacker is, where the trust boundaries sit, and what counts as critical here, and it is cached per target rather than per scan. Read it instead of re-deriving trust boundaries yourself; where your testing disproves it — a boundary it calls trusted turns out to be attacker-reachable, a role it did not know about, a host or endpoint it never listed — record that with `amend_threat_model` so the agents after you inherit the correction. Amending is not optional politeness: a model nobody corrects turns the first agent's guesses into everyone's assumptions.
|
||||
- THREAT MODEL: before you start testing, call `get_threat_model` on the target you were pointed at — it is the scan's shared answer to who the attacker is, where the trust boundaries sit, and what counts as critical here. It is scoped to this scan and nothing carries over from an earlier run, so `found: false` means no agent on this run has derived one yet. Read it instead of re-deriving trust boundaries yourself; where your testing disproves it — a boundary it calls trusted turns out to be attacker-reachable, a role it did not know about, a host or endpoint it never listed — record that with `amend_threat_model` so the agents after you inherit the correction. Amending is not optional politeness: a model nobody corrects turns the first agent's guesses into everyone's assumptions.
|
||||
- Before filing any report, run the counterevidence pass: argue the strongest case AGAINST the finding, record what you found in the `counterevidence` field, set `confidence` honestly (a static-only trace you couldn't execute is at best `medium`), and state what evidence would change the severity. See the counterevidence and severity-calibration knowledge above.
|
||||
- A vulnerability is ONLY considered reported when a reporting agent uses create_vulnerability_report (or create_dependency_report for known-CVE dependency/supply-chain findings) with full details. Mentions in agent_finish, finish_scan, or generic messages are NOT sufficient
|
||||
- Reporting and fixing are ONE step, not two: when source is available, the reporting agent derives the concrete fix and files it INLINE via create_vulnerability_report (`code_locations` with `fix_before`/`fix_after` + `fix_pr_body`) — the report is not complete without it. Do NOT report first and then spawn a separate downstream agent to re-derive and re-apply the same patch; that just re-does the analysis and wastes tokens. (Do not silently patch a finding WITHOUT filing a report — the report, with its embedded fix, is the deliverable.)
|
||||
|
|
|
|||
|
|
@ -96,14 +96,12 @@ func TestCoverageDuplicateRejectionSurfacesTheError(t *testing.T) {
|
|||
requireContains(t, out, "/login", "already has coverage entry a1b2c3")
|
||||
}
|
||||
|
||||
func TestGetThreatModelRendersStalenessAndAmendments(t *testing.T) {
|
||||
func TestGetThreatModelRendersAmendments(t *testing.T) {
|
||||
out := ansi.Strip(Tool(tool("get_threat_model",
|
||||
map[string]any{"target": "https://app.example.com"},
|
||||
map[string]any{
|
||||
"success": true,
|
||||
"found": true,
|
||||
"stale": true,
|
||||
"cached_revision": "0123456789abcdef",
|
||||
"success": true,
|
||||
"found": true,
|
||||
"content": "# Overview\nMulti-tenant billing app.\n\n" +
|
||||
"## Trust Boundaries and Assumptions\n\n## Attack Surface\n",
|
||||
"amendments": []any{
|
||||
|
|
@ -116,7 +114,6 @@ func TestGetThreatModelRendersStalenessAndAmendments(t *testing.T) {
|
|||
"completed")))
|
||||
requireContains(t, out,
|
||||
"Threat Model", "https://app.example.com",
|
||||
"stale", "01234567",
|
||||
"1 amendment(s)", "ReconAgent", "staging host shares the production database",
|
||||
"Multi-tenant billing app.", "Overview", "Trust Boundaries and Assumptions",
|
||||
)
|
||||
|
|
@ -126,7 +123,7 @@ func TestGetThreatModelMissingModelIsExplicit(t *testing.T) {
|
|||
out := ansi.Strip(Tool(tool("get_threat_model",
|
||||
map[string]any{"target": "10.0.0.5"},
|
||||
map[string]any{"success": true, "found": false}, "completed")))
|
||||
requireContains(t, out, "No model cached for this target yet")
|
||||
requireContains(t, out, "No model derived for this target yet")
|
||||
}
|
||||
|
||||
func TestSaveThreatModelWarnsWhenAmendmentsAreCleared(t *testing.T) {
|
||||
|
|
@ -134,15 +131,10 @@ func TestSaveThreatModelWarnsWhenAmendmentsAreCleared(t *testing.T) {
|
|||
map[string]any{"target": "app.example.com", "content": "# Overview\nA thing.\n"},
|
||||
map[string]any{
|
||||
"success": true,
|
||||
"revision": "unversioned",
|
||||
"amendments_cleared": 2,
|
||||
},
|
||||
"completed")))
|
||||
requireContains(t, out, "Threat Model Saved", "saved", "cleared 2 amendment(s)")
|
||||
// An unversioned target has no revision worth printing.
|
||||
if strings.Contains(out, "unversioned") {
|
||||
t.Fatalf("unversioned revision should not be rendered:\n%s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmendThreatModelRendersAddendum(t *testing.T) {
|
||||
|
|
|
|||
|
|
@ -56,9 +56,6 @@ func renderThreatModel(name string, args map[string]any, result any) string {
|
|||
threatModelBody(&b, StringValue(args["addendum"]))
|
||||
default:
|
||||
b.WriteString("\n " + Col(Green).Render("✓ saved"))
|
||||
if revision := shortRevision(StringValue(m["revision"])); revision != "" {
|
||||
b.WriteString(Dim().Render(" at " + revision))
|
||||
}
|
||||
// Saving folds amendments away, so the count that vanished is worth
|
||||
// stating: it is the one destructive thing this tool does.
|
||||
if cleared, ok := NumericValue(m["amendments_cleared"]); ok && cleared > 0 {
|
||||
|
|
@ -72,15 +69,9 @@ func renderThreatModel(name string, args map[string]any, result any) string {
|
|||
|
||||
func threatModelReadBody(b *strings.Builder, result map[string]any) {
|
||||
if !truthy(result["found"]) {
|
||||
b.WriteString("\n " + Dim().Render("No model cached for this target yet"))
|
||||
b.WriteString("\n " + Dim().Render("No model derived for this target yet"))
|
||||
return
|
||||
}
|
||||
if truthy(result["stale"]) {
|
||||
b.WriteString("\n " + Col(AmberY).Render("⚠ stale"))
|
||||
if cached := shortRevision(StringValue(result["cached_revision"])); cached != "" {
|
||||
b.WriteString(Dim().Render(" (written at " + cached + ")"))
|
||||
}
|
||||
}
|
||||
if amendments, ok := result["amendments"].([]any); ok && len(amendments) > 0 {
|
||||
b.WriteString("\n " + Col(Gold).Render("+ "+strconv.Itoa(len(amendments))+
|
||||
" amendment(s)") + Dim().Render(" — later statements win"))
|
||||
|
|
@ -126,13 +117,3 @@ func threatModelBody(b *strings.Builder, content string) {
|
|||
b.WriteString("\n " + Dim().Render(strings.Join(headings, " · ")))
|
||||
}
|
||||
}
|
||||
|
||||
// shortRevision abbreviates a git sha; "unversioned" targets have no revision
|
||||
// worth showing.
|
||||
func shortRevision(revision string) string {
|
||||
revision = strings.TrimSpace(revision)
|
||||
if revision == "" || revision == "unversioned" {
|
||||
return ""
|
||||
}
|
||||
return firstN(revision, 8)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,13 +16,6 @@ const ACTION_LABELS: Record<string, { label: string; Icon: typeof Crosshair }> =
|
|||
amend_threat_model: { label: "Threat model amended", Icon: Plus },
|
||||
};
|
||||
|
||||
/** A git sha is noise past its first bytes, and "unversioned" is not a revision. */
|
||||
function shortRevision(revision: unknown): string {
|
||||
const value = typeof revision === "string" ? revision.trim() : "";
|
||||
if (!value || value === "unversioned") return "";
|
||||
return value.slice(0, 8);
|
||||
}
|
||||
|
||||
export default function ThreatModelRenderer({ toolName, args, result }: ToolRendererProps) {
|
||||
const action = ACTION_LABELS[toolName] ?? { label: "Threat model", Icon: Crosshair };
|
||||
const ActionIcon = action.Icon;
|
||||
|
|
@ -59,22 +52,15 @@ export default function ThreatModelRenderer({ toolName, args, result }: ToolRend
|
|||
return (
|
||||
<div>
|
||||
{header}
|
||||
<div className="mt-1.5 text-[#555] text-xs">No model cached for this target yet</div>
|
||||
<div className="mt-1.5 text-[#555] text-xs">No model derived for this target yet</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
const rawAmendments = structured?.amendments;
|
||||
const amendments: Amendment[] = Array.isArray(rawAmendments) ? (rawAmendments as Amendment[]) : [];
|
||||
const cachedRevision = shortRevision(structured?.cached_revision);
|
||||
return (
|
||||
<div>
|
||||
{header}
|
||||
{structured?.stale === true && (
|
||||
<div className="mt-1.5 flex items-center gap-1.5 text-yellow-400/80 text-xs">
|
||||
<AlertTriangle className="w-3 h-3 shrink-0" />
|
||||
<span>stale{cachedRevision ? ` — written at ${cachedRevision}` : ""}</span>
|
||||
</div>
|
||||
)}
|
||||
{amendments.length > 0 && (
|
||||
<div className="mt-2">
|
||||
<span className="text-amber-400/70 text-xs font-semibold">
|
||||
|
|
@ -119,12 +105,10 @@ export default function ThreatModelRenderer({ toolName, args, result }: ToolRend
|
|||
}
|
||||
|
||||
const cleared = (structured?.amendments_cleared as number | undefined) ?? 0;
|
||||
const revision = shortRevision(structured?.revision);
|
||||
const content = (args.content as string) ?? "";
|
||||
return (
|
||||
<div>
|
||||
{header}
|
||||
{revision && <div className="mt-1.5 text-[#666] font-mono text-xs">at {revision}</div>}
|
||||
{/* Saving folds amendments away — the one destructive thing this tool does. */}
|
||||
{cleared > 0 && (
|
||||
<div className="mt-1.5 flex items-center gap-1.5 text-yellow-400/80 text-xs">
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ Before spawning agents, analyze the target from the scan config/scope and any pr
|
|||
|
||||
## Establish the Threat Model
|
||||
|
||||
Every scan needs one shared answer to "who is the attacker here, and what are they attacking" — black-box or white-box. Without it, five agents derive five different answers and their findings cannot be reconciled. Call `get_threat_model` on the target (a host, a URL, or a repository path) before you spawn hunters; if nothing is cached, derive one and persist it with `save_threat_model`. It is cached per target, so a later scan of the same host or tree reads it back instead of paying for it twice, and a model written from source is read back by an agent testing the deployment.
|
||||
Every scan needs one shared answer to "who is the attacker here, and what are they attacking" — black-box or white-box. Without it, five agents derive five different answers and their findings cannot be reconciled. Call `get_threat_model` on the target (a host, a URL, or a repository path) before you spawn hunters; if no model exists yet, derive one and share it with `save_threat_model`. It lives for this scan only — nothing carries over from an earlier run, so every scan derives its own — but within the run every agent reads the same document, and a model written from source is read back by an agent testing the deployment.
|
||||
|
||||
**When the target includes a repository**, derive it up front: the code tells you the boundaries, entrypoints, and controls before you send a single request.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,28 +1,27 @@
|
|||
"""Target-scoped threat models — cached under ``~/.strix/threat-models``.
|
||||
"""Run-scoped threat models — held in memory for the duration of one scan.
|
||||
|
||||
A threat model describes the target, not the scan: a host, an application, an
|
||||
API, a repository, or whatever else the engagement is pointed at. It stays
|
||||
valid across unrelated runs against the same target, so it is keyed by target
|
||||
identity rather than by run id — one agent derives it, every later agent in
|
||||
this run and in future runs against the same target reads it back instead of
|
||||
A threat model is the scan's shared answer to who the attacker is, where the
|
||||
trust boundaries sit, and what counts as critical for the target. One agent
|
||||
derives it and every other agent on the same run reads it back instead of
|
||||
re-deriving trust boundaries from scratch.
|
||||
|
||||
Where the target is a checkout, the model is additionally pinned to the git
|
||||
revision, so a moved ``HEAD`` marks it stale. Black-box targets have no
|
||||
revision to pin to; those age out instead.
|
||||
It never outlives the run. Nothing is written to disk and nothing is shared
|
||||
between scans: a new scan against the same host or checkout starts with no
|
||||
model and derives its own. Agents do spell one target several ways within a
|
||||
run — the URL they were handed, the page they happen to be testing, a checkout
|
||||
path — so a model is keyed by a normalized target identity to keep them
|
||||
converging on one document instead of each starting a fresh one.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
|
|
@ -35,16 +34,18 @@ from strix.core.agents import AgentCoordinator
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
_CACHE_DIR = Path.home() / ".strix" / "threat-models"
|
||||
_MAX_MODEL_BYTES = 512 * 1024
|
||||
_MIN_MODEL_CHARS = 400
|
||||
_MIN_AMENDMENT_CHARS = 80
|
||||
_MAX_AMENDMENTS = 40
|
||||
_GIT_TIMEOUT_SECONDS = 10
|
||||
_UNVERSIONED = "unversioned"
|
||||
_MAX_AGE_DAYS = 14
|
||||
_DEFAULT_PORTS = {"http": "80", "https": "443"}
|
||||
_cache_lock = threading.RLock()
|
||||
|
||||
_store_lock = threading.RLock()
|
||||
|
||||
# The whole store: target identity -> model. Process-local and never persisted,
|
||||
# so it holds exactly the models this run derived and dies with it.
|
||||
_MODELS: dict[str, dict[str, Any]] = {}
|
||||
|
||||
_REQUIRED_SECTIONS = (
|
||||
"overview",
|
||||
|
|
@ -95,7 +96,7 @@ def _remote_authority(target: str) -> str:
|
|||
|
||||
|
||||
def _normalize_remote_target(target: str) -> str:
|
||||
"""Collapse the spellings of one remote target onto a single cache key."""
|
||||
"""Collapse the spellings of one remote target onto a single key."""
|
||||
authority = _remote_authority(target)
|
||||
if not authority:
|
||||
return re.sub(r"\s+", " ", target.lower()).strip()
|
||||
|
|
@ -132,30 +133,23 @@ def _normalize_git_remote(remote: str) -> str:
|
|||
return normalized.removesuffix(".git")
|
||||
|
||||
|
||||
def _target_identity(target: str) -> tuple[str, str]:
|
||||
"""Return the (stable identity, revision) pair a cached model is keyed on.
|
||||
def _target_identity(target: str) -> str:
|
||||
"""Return the stable identity a model is stored under.
|
||||
|
||||
A checkout is keyed on its remote (so the same repository cloned to two
|
||||
paths shares one model, and a subdirectory resolves to the whole tree) and
|
||||
pinned to ``HEAD``. Everything else — a host, a URL, an API base, a named
|
||||
scope — is keyed on its normalized form and carries no revision. Both
|
||||
routes run through the same normalization, so a checkout and the URL it
|
||||
was cloned from land on one key.
|
||||
A checkout is keyed on its remote, so the same repository checked out at
|
||||
two paths shares one model and a subdirectory resolves to the whole tree.
|
||||
Everything else — a host, a URL, an API base, a named scope — is keyed on
|
||||
its normalized form. Both routes run through the same normalization, so a
|
||||
checkout and the URL it was cloned from land on one key.
|
||||
"""
|
||||
directory = _local_directory(target)
|
||||
if directory is None:
|
||||
return _normalize_remote_target(target).removesuffix(".git"), _UNVERSIONED
|
||||
return _normalize_remote_target(target).removesuffix(".git")
|
||||
remote = _git(directory, ["config", "--get", "remote.origin.url"])
|
||||
revision = _git(directory, ["rev-parse", "HEAD"]) or _UNVERSIONED
|
||||
if remote:
|
||||
return _normalize_git_remote(remote), revision
|
||||
return _normalize_git_remote(remote)
|
||||
toplevel = _git(directory, ["rev-parse", "--show-toplevel"])
|
||||
return toplevel or str(directory), revision
|
||||
|
||||
|
||||
def _cache_path(identity: str) -> Path:
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:16]
|
||||
return _CACHE_DIR / f"{digest}.json"
|
||||
return toplevel or str(directory)
|
||||
|
||||
|
||||
def _snap_to_scan_target(raw: str, scan_targets: list[str]) -> str:
|
||||
|
|
@ -163,13 +157,13 @@ def _snap_to_scan_target(raw: str, scan_targets: list[str]) -> str:
|
|||
|
||||
Agents name the same target differently — one passes the URL it was given,
|
||||
the next the page it happens to be testing, a third the checkout path. Left
|
||||
alone those become separate cache keys, every lookup misses, and each agent
|
||||
alone those become separate keys, every lookup misses, and each agent
|
||||
quietly derives its own model, which is the exact failure the shared model
|
||||
exists to prevent. So a target that is recognisably one of the scan's own
|
||||
targets is resolved to that target instead.
|
||||
"""
|
||||
identity, _ = _target_identity(raw)
|
||||
scoped = [(target, _target_identity(target)[0]) for target in scan_targets]
|
||||
identity = _target_identity(raw)
|
||||
scoped = [(target, _target_identity(target)) for target in scan_targets]
|
||||
if any(known == identity for _, known in scoped):
|
||||
return raw
|
||||
|
||||
|
|
@ -208,126 +202,54 @@ def _resolve_target(
|
|||
return (_snap_to_scan_target(raw, known) if known else raw), None
|
||||
|
||||
|
||||
def _is_expired(created_at: str | None) -> bool:
|
||||
if not created_at:
|
||||
return True
|
||||
try:
|
||||
created = datetime.fromisoformat(created_at)
|
||||
except ValueError:
|
||||
return True
|
||||
if created.tzinfo is None:
|
||||
created = created.replace(tzinfo=UTC)
|
||||
return datetime.now(UTC) - created > timedelta(days=_MAX_AGE_DAYS)
|
||||
|
||||
|
||||
def _missing_sections(content: str) -> list[str]:
|
||||
lowered = content.lower()
|
||||
return [section for section in _REQUIRED_SECTIONS if section not in lowered]
|
||||
|
||||
|
||||
def _read_cache(path: Path) -> dict[str, Any] | None:
|
||||
"""Load a cached model. Callers must already hold ``_cache_lock``."""
|
||||
if not path.is_file():
|
||||
return None
|
||||
try:
|
||||
cached = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
logger.exception("threat model cache at %s is unreadable", path)
|
||||
return None
|
||||
return cached if isinstance(cached, dict) else None
|
||||
|
||||
|
||||
def _write_cache(path: Path, payload: dict[str, Any]) -> str | None:
|
||||
"""Atomically persist a model. Callers must already hold ``_cache_lock``."""
|
||||
try:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
dir=str(path.parent),
|
||||
prefix=f".{path.name}.",
|
||||
suffix=".tmp",
|
||||
delete=False,
|
||||
) as tmp:
|
||||
tmp.write(json.dumps(payload, ensure_ascii=False))
|
||||
tmp_path = Path(tmp.name)
|
||||
tmp_path.replace(path)
|
||||
except OSError as exc:
|
||||
logger.exception("threat model persist to %s failed", path)
|
||||
return f"Failed to persist threat model: {exc}"
|
||||
return None
|
||||
|
||||
|
||||
def _amendments_of(cached: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
raw = cached.get("amendments")
|
||||
def _amendments_of(model: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
raw = model.get("amendments")
|
||||
if not isinstance(raw, list):
|
||||
return []
|
||||
return [item for item in raw if isinstance(item, dict)]
|
||||
|
||||
|
||||
def _not_found(identity: str, revision: str) -> dict[str, Any]:
|
||||
def _not_found(identity: str) -> dict[str, Any]:
|
||||
return {
|
||||
"success": True,
|
||||
"found": False,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"message": (
|
||||
"No threat model cached for this target. Derive one — from the code if "
|
||||
"you have it, from recon output if you do not — and persist it with "
|
||||
"save_threat_model, so every agent on this scan shares one view of the "
|
||||
"trust boundaries instead of each inventing their own."
|
||||
"No threat model for this target on this scan. Nothing carries over "
|
||||
"from earlier runs, so derive one — from the code if you have it, from "
|
||||
"recon output if you do not — and share it with save_threat_model, so "
|
||||
"every agent on this scan works from one view of the trust boundaries "
|
||||
"instead of each inventing their own."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _staleness(cached: dict[str, Any], revision: str) -> tuple[bool, str | None]:
|
||||
"""Decide whether a cached model can still be trusted, and why not."""
|
||||
if revision != _UNVERSIONED:
|
||||
if cached.get("revision") == revision:
|
||||
return False, None
|
||||
return True, (
|
||||
"This model was derived against a different revision. Use it as a "
|
||||
"starting point, re-check the boundaries it names against the current "
|
||||
"tree, and save the corrected version."
|
||||
)
|
||||
created_at = cached.get("created_at")
|
||||
if not _is_expired(created_at if isinstance(created_at, str) else None):
|
||||
return False, None
|
||||
return True, (
|
||||
f"This model is more than {_MAX_AGE_DAYS} days old and there is no revision "
|
||||
"to pin it to, so the target may have moved under it. Treat its surface "
|
||||
"inventory as a lead list to re-confirm during recon, not as fact, and save "
|
||||
"the corrected version."
|
||||
)
|
||||
|
||||
|
||||
def _get_impl(target: str, scan_targets: list[str] | None = None) -> dict[str, Any]:
|
||||
resolved, error = _resolve_target(target, scan_targets)
|
||||
if resolved is None:
|
||||
return {"success": False, "error": error}
|
||||
|
||||
identity, revision = _target_identity(resolved)
|
||||
path = _cache_path(identity)
|
||||
with _cache_lock:
|
||||
cached = _read_cache(path)
|
||||
if cached is None:
|
||||
return _not_found(identity, revision)
|
||||
content = cached.get("content")
|
||||
identity = _target_identity(resolved)
|
||||
with _store_lock:
|
||||
model = _MODELS.get(identity)
|
||||
if model is None:
|
||||
return _not_found(identity)
|
||||
content = model.get("content")
|
||||
amendments = list(_amendments_of(model))
|
||||
if not isinstance(content, str) or not content.strip():
|
||||
return _not_found(identity, revision)
|
||||
return _not_found(identity)
|
||||
|
||||
stale, stale_message = _staleness(cached, revision)
|
||||
result: dict[str, Any] = {
|
||||
"success": True,
|
||||
"found": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"cached_revision": cached.get("revision"),
|
||||
"created_at": cached.get("created_at"),
|
||||
"stale": stale,
|
||||
"content": content,
|
||||
}
|
||||
amendments = _amendments_of(cached)
|
||||
if amendments:
|
||||
result["amendments"] = amendments
|
||||
result["amendments_note"] = (
|
||||
|
|
@ -335,8 +257,6 @@ def _get_impl(target: str, scan_targets: list[str] | None = None) -> dict[str, A
|
|||
"correct or extend it and have not been folded in yet - read them as "
|
||||
"part of the model, and prefer the later one where they conflict."
|
||||
)
|
||||
if stale_message:
|
||||
result["message"] = stale_message
|
||||
return result
|
||||
|
||||
|
||||
|
|
@ -376,25 +296,20 @@ def _save_impl(
|
|||
),
|
||||
}
|
||||
|
||||
identity, revision = _target_identity(resolved)
|
||||
path = _cache_path(identity)
|
||||
payload: dict[str, Any] = {
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"created_at": datetime.now(UTC).isoformat(),
|
||||
"created_by": agent_name,
|
||||
"content": body,
|
||||
}
|
||||
with _cache_lock:
|
||||
existing = _read_cache(path)
|
||||
identity = _target_identity(resolved)
|
||||
with _store_lock:
|
||||
existing = _MODELS.get(identity)
|
||||
folded = len(_amendments_of(existing)) if existing else 0
|
||||
error = _write_cache(path, payload)
|
||||
if error:
|
||||
return {"success": False, "error": error}
|
||||
_MODELS[identity] = {
|
||||
"target": identity,
|
||||
"written_at": datetime.now(UTC).isoformat(),
|
||||
"written_by": agent_name,
|
||||
"content": body,
|
||||
}
|
||||
|
||||
message = (
|
||||
"Threat model saved. Subagents should call get_threat_model before they "
|
||||
"start, and treat its trust boundaries as the shared baseline."
|
||||
"Threat model shared with this scan. Subagents should call get_threat_model "
|
||||
"before they start, and treat its trust boundaries as the shared baseline."
|
||||
)
|
||||
if folded:
|
||||
message += (
|
||||
|
|
@ -404,34 +319,34 @@ def _save_impl(
|
|||
return {
|
||||
"success": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"amendments_cleared": folded,
|
||||
"message": message,
|
||||
}
|
||||
|
||||
|
||||
def _append_amendment(
|
||||
path: Path, amendment: dict[str, Any]
|
||||
identity: str, amendment: dict[str, Any]
|
||||
) -> tuple[list[dict[str, Any]] | None, str | None]:
|
||||
"""Add an amendment to the cached model. Returns (amendments, error)."""
|
||||
with _cache_lock:
|
||||
cached = _read_cache(path)
|
||||
if cached is None or not str(cached.get("content", "")).strip():
|
||||
"""Add an amendment to the stored model. Returns (amendments, error)."""
|
||||
with _store_lock:
|
||||
model = _MODELS.get(identity)
|
||||
if model is None or not str(model.get("content", "")).strip():
|
||||
return None, (
|
||||
"No threat model exists for this target yet, so there is nothing to "
|
||||
"amend. Derive the base model and call save_threat_model instead."
|
||||
)
|
||||
amendments = _amendments_of(cached)
|
||||
amendments = _amendments_of(model)
|
||||
if len(amendments) >= _MAX_AMENDMENTS:
|
||||
return None, (
|
||||
f"This model already carries {len(amendments)} amendments. Fold them "
|
||||
"into the base model with save_threat_model before adding more."
|
||||
)
|
||||
amendments.append(amendment)
|
||||
cached["amendments"] = amendments
|
||||
if len(json.dumps(cached, ensure_ascii=False).encode("utf-8")) > _MAX_MODEL_BYTES:
|
||||
candidate = [*amendments, amendment]
|
||||
sized = {**model, "amendments": candidate}
|
||||
if len(json.dumps(sized, ensure_ascii=False).encode("utf-8")) > _MAX_MODEL_BYTES:
|
||||
return None, "Threat model with this amendment exceeds 512KB; tighten it."
|
||||
return amendments, _write_cache(path, cached)
|
||||
model["amendments"] = candidate
|
||||
return candidate, None
|
||||
|
||||
|
||||
def _amend_impl(
|
||||
|
|
@ -455,13 +370,12 @@ def _amend_impl(
|
|||
),
|
||||
}
|
||||
|
||||
identity, revision = _target_identity(resolved)
|
||||
identity = _target_identity(resolved)
|
||||
amendments, amend_error = _append_amendment(
|
||||
_cache_path(identity),
|
||||
identity,
|
||||
{
|
||||
"at": datetime.now(UTC).isoformat(),
|
||||
"by": agent_name,
|
||||
"revision": revision,
|
||||
"content": body,
|
||||
},
|
||||
)
|
||||
|
|
@ -471,7 +385,6 @@ def _amend_impl(
|
|||
return {
|
||||
"success": True,
|
||||
"target": identity,
|
||||
"revision": revision,
|
||||
"amendment_count": len(amendments),
|
||||
"message": (
|
||||
"Amendment recorded. Agents calling get_threat_model will now see it "
|
||||
|
|
@ -500,25 +413,25 @@ def _scan_targets(ctx: RunContextWrapper) -> list[str]:
|
|||
|
||||
@function_tool(timeout=30)
|
||||
async def get_threat_model(ctx: RunContextWrapper, target: str) -> str:
|
||||
"""Read the cached threat model for a target, if one exists.
|
||||
"""Read this scan's threat model for a target, if an agent has derived one.
|
||||
|
||||
A threat model belongs to the target, not to this scan — the same
|
||||
trust boundaries hold across unrelated runs against the same host
|
||||
or application. Call this before you start hunting so you inherit
|
||||
the shared view instead of re-deriving it, and so every agent on
|
||||
this run agrees on what "attacker-controlled" means here.
|
||||
The threat model is this run's shared answer to who the attacker
|
||||
is, where the trust boundaries sit, and what counts as critical
|
||||
here. Call it before you start hunting so you inherit the shared
|
||||
view instead of re-deriving it, and so every agent on this run
|
||||
agrees on what "attacker-controlled" means.
|
||||
|
||||
It is scoped to this scan and nothing is carried over from an
|
||||
earlier run, so an empty result means no agent has derived one yet.
|
||||
|
||||
Works black-box or white-box. The target can be a host, a URL, an
|
||||
API base, or a repository path; equivalent spellings of the same
|
||||
host resolve to the same model, and a checkout resolves to its
|
||||
remote, so a model derived white-box is read back by a black-box
|
||||
agent testing the deployment.
|
||||
remote, so a model derived white-box by one agent is read back by
|
||||
another testing the deployment.
|
||||
|
||||
Returns ``found: false`` when nothing is cached — derive one and
|
||||
persist it with ``save_threat_model``. ``stale: true`` means the
|
||||
checkout moved to a different revision, or that a model with no
|
||||
revision to pin to has aged out: use it as a starting point,
|
||||
re-confirm what it claims, and save the corrected version.
|
||||
Returns ``found: false`` when nothing has been derived yet — derive
|
||||
one and share it with ``save_threat_model``.
|
||||
|
||||
Any ``amendments`` in the response are corrections other agents
|
||||
recorded after the base model was written. They are part of the
|
||||
|
|
@ -540,10 +453,10 @@ async def get_threat_model(ctx: RunContextWrapper, target: str) -> str:
|
|||
|
||||
@function_tool(timeout=30)
|
||||
async def save_threat_model(ctx: RunContextWrapper, target: str, content: str) -> str:
|
||||
"""Persist a target-scoped threat model for reuse by other agents.
|
||||
"""Share a target-scoped threat model with the other agents on this scan.
|
||||
|
||||
Keyed by target identity, so a later scan of the same host or tree
|
||||
reads it back instead of paying to derive it again.
|
||||
The model lives for this run only — it is not written to disk and a
|
||||
later scan of the same host or tree starts without it.
|
||||
|
||||
**This replaces the whole document, and clears any amendments** —
|
||||
it is for the agent establishing the baseline (normally root,
|
||||
|
|
@ -561,9 +474,9 @@ async def save_threat_model(ctx: RunContextWrapper, target: str, content: str) -
|
|||
necessarily provisional — say which parts are inferred rather than
|
||||
observed, and let later agents amend it as the picture fills in.
|
||||
|
||||
**Scope it to the target, not to this scan.** Do not centre it on
|
||||
the diff you were handed, the subsystem you were assigned, or the
|
||||
one host that happened to answer first. With source, distinguish
|
||||
**Scope it to the target, not to your slice of it.** Do not centre
|
||||
it on the diff you were handed, the subsystem you were assigned, or
|
||||
the one host that happened to answer first. With source, distinguish
|
||||
real product and runtime surfaces from test, docs, example, and
|
||||
developer-tooling paths — in a monorepo, do not let ``tests/`` or
|
||||
one-off scripts become the centre of gravity unless the code shows
|
||||
|
|
|
|||
|
|
@ -1,10 +1,8 @@
|
|||
"""Tests for the target-scoped threat model cache."""
|
||||
"""Tests for the run-scoped threat model store."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
|
@ -64,8 +62,9 @@ def _make_repo(tmp_path: Path, name: str = "repo") -> Path:
|
|||
|
||||
|
||||
@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 _empty_store() -> None:
|
||||
"""Each test is its own run, so it starts with an empty store."""
|
||||
threat_model_tools._MODELS.clear()
|
||||
|
||||
|
||||
def test_missing_model_reports_not_found(tmp_path: Path) -> None:
|
||||
|
|
@ -85,11 +84,34 @@ def test_saved_model_round_trips(tmp_path: Path) -> None:
|
|||
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:
|
||||
def test_nothing_is_written_to_disk(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""The model must not outlive the run, so no file may be left behind."""
|
||||
home = tmp_path / "home"
|
||||
home.mkdir()
|
||||
monkeypatch.setenv("HOME", str(home))
|
||||
repo = _make_repo(tmp_path)
|
||||
|
||||
_save_impl(str(repo), _MODEL, "root")
|
||||
_amend_impl(str(repo), _ADDENDUM, "agent-a")
|
||||
|
||||
assert list(home.rglob("*")) == []
|
||||
|
||||
|
||||
def test_a_new_run_starts_without_the_model(tmp_path: Path) -> None:
|
||||
"""A later scan of the same target inherits nothing from this one."""
|
||||
repo = _make_repo(tmp_path)
|
||||
_save_impl(str(repo), _MODEL, "root")
|
||||
|
||||
threat_model_tools._MODELS.clear() # what a fresh process starts from
|
||||
|
||||
assert _get_impl(str(repo))["found"] is False
|
||||
|
||||
|
||||
def test_model_survives_a_new_revision_within_the_run(tmp_path: Path) -> None:
|
||||
"""The model is not pinned to a revision; a commit mid-run does not drop it."""
|
||||
repo = _make_repo(tmp_path)
|
||||
_save_impl(str(repo), _MODEL, None)
|
||||
|
||||
|
|
@ -100,11 +122,10 @@ def test_model_is_stale_after_new_revision(tmp_path: Path) -> None:
|
|||
result = _get_impl(str(repo))
|
||||
|
||||
assert result["found"] is True
|
||||
assert result["stale"] is True
|
||||
assert result["content"]
|
||||
assert "multi-tenant billing API" in result["content"]
|
||||
|
||||
|
||||
def test_cache_is_keyed_per_repository(tmp_path: Path) -> None:
|
||||
def test_store_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)
|
||||
|
|
@ -218,8 +239,6 @@ def test_blackbox_target_round_trips() -> None:
|
|||
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"]
|
||||
|
||||
|
||||
|
|
@ -232,22 +251,6 @@ def test_blackbox_target_spellings_share_one_model() -> None:
|
|||
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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue