attackmap 0.1.0__py3-none-any.whl

attackmap/llm_review.py

@@ -0,0 +1,346 @@
+"""Anthropic-powered narrative review generator.
+
+Wraps the existing prompt scaffolding (`review_prompts.render_review_prompts`)
+in a real Claude call so users can produce a story-bearing defensive review
+without copy-pasting JSON into a chat client.
+
+Auth resolution (in order):
+  1. Explicit `client=` (tests / programmatic override)
+  2. `ANTHROPIC_API_KEY` env  → SDK with api_key  (per-token API billing)
+  3. `ANTHROPIC_AUTH_TOKEN` env → SDK with auth_token (OAuth bearer)
+  4. `claude` CLI on PATH      → shell out to `claude -p --output-format=json`
+                                 (uses whatever auth `claude login` configured —
+                                 typically the user's Pro/Max subscription)
+
+The user can force a backend via the `backend` argument or the CLI's
+`--llm-backend` flag; "auto" walks the order above and picks the first one
+that resolves.
+
+Optional dependency: `pip install attackmap[llm]` installs the anthropic SDK.
+The CLI backend only requires the `claude` binary on PATH (no SDK install).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from .models import AttackPath, AttackSurface, Finding, ScanResult
+from .review_prompts import render_review_prompts
+
+DEFAULT_MODEL = "claude-opus-4-7"
+DEFAULT_EFFORT: Literal["low", "medium", "high", "xhigh", "max"] = "high"
+DEFAULT_MAX_TOKENS = 32000
+CLAUDE_CLI_TIMEOUT_SECONDS = 600
+
+LlmBackend = Literal["auto", "api", "cli"]
+
+
+class LlmReviewError(RuntimeError):
+    """Raised when the LLM-backed review cannot be produced."""
+
+
+@dataclass(frozen=True)
+class LlmReviewResult:
+    markdown: str
+    model: str
+    stop_reason: str | None
+    usage: dict[str, int]
+    backend: Literal["api", "cli"]
+
+
+# ---------- Auth + SDK backend ----------
+
+
+def _resolve_sdk_client(api_key: str | None, client: Any | None) -> tuple[Any, str]:
+    """Return (anthropic client, auth_kind). auth_kind ∈ {'api_key', 'auth_token'}."""
+    if client is not None:
+        return client, "explicit"
+    try:
+        import anthropic
+    except ImportError as exc:
+        raise LlmReviewError(
+            "The anthropic SDK is not installed. Install with `pip install attackmap[llm]` "
+            "to use the API backend, or ensure the `claude` CLI is on PATH for the CLI backend."
+        ) from exc
+
+    resolved_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
+    if resolved_key:
+        return anthropic.Anthropic(api_key=resolved_key), "api_key"
+
+    auth_token = os.environ.get("ANTHROPIC_AUTH_TOKEN")
+    if auth_token:
+        return anthropic.Anthropic(auth_token=auth_token), "auth_token"
+
+    raise LlmReviewError(
+        "No API credentials found. Set ANTHROPIC_API_KEY (per-token API billing) or "
+        "ANTHROPIC_AUTH_TOKEN (OAuth bearer), or install the `claude` CLI and run "
+        "`claude login` to use your Claude subscription."
+    )
+
+
+def _extract_text(blocks: list[Any]) -> str:
+    parts: list[str] = []
+    for block in blocks:
+        if getattr(block, "type", None) == "text":
+            text = getattr(block, "text", "")
+            if text:
+                parts.append(text)
+    return "\n".join(parts).strip()
+
+
+def _usage_dict(usage: Any) -> dict[str, int]:
+    out: dict[str, int] = {}
+    for field in (
+        "input_tokens",
+        "output_tokens",
+        "cache_creation_input_tokens",
+        "cache_read_input_tokens",
+    ):
+        value = getattr(usage, field, None)
+        if isinstance(value, int):
+            out[field] = value
+    return out
+
+
+def _run_via_sdk(
+    rendered_system: str,
+    rendered_user: str,
+    *,
+    model: str,
+    effort: str,
+    max_tokens: int,
+    api_key: str | None,
+    client: Any | None,
+) -> LlmReviewResult:
+    sdk_client, _ = _resolve_sdk_client(api_key, client)
+
+    request_kwargs: dict[str, Any] = {
+        "model": model,
+        "max_tokens": max_tokens,
+        "system": [
+            {"type": "text", "text": rendered_system, "cache_control": {"type": "ephemeral"}}
+        ],
+        "messages": [{"role": "user", "content": rendered_user}],
+        "thinking": {"type": "adaptive"},
+        "output_config": {"effort": effort},
+    }
+
+    try:
+        with sdk_client.messages.stream(**request_kwargs) as stream:
+            final_message = stream.get_final_message()
+    except LlmReviewError:
+        raise
+    except Exception as exc:  # pragma: no cover - surface SDK errors verbatim
+        raise LlmReviewError(f"Claude API call failed: {exc}") from exc
+
+    markdown = _extract_text(getattr(final_message, "content", []) or [])
+    if not markdown:
+        raise LlmReviewError("Claude returned no text content for the review.")
+
+    stop_reason = getattr(final_message, "stop_reason", None)
+    return LlmReviewResult(
+        markdown=markdown,
+        model=model,
+        stop_reason=stop_reason if isinstance(stop_reason, str) else None,
+        usage=_usage_dict(getattr(final_message, "usage", None)),
+        backend="api",
+    )
+
+
+# ---------- Claude CLI backend ----------
+
+
+def _claude_cli_available() -> bool:
+    return shutil.which("claude") is not None
+
+
+def _run_via_claude_cli(
+    rendered_system: str,
+    rendered_user: str,
+    *,
+    model: str,
+    runner: Any | None = None,
+) -> LlmReviewResult:
+    """Invoke `claude -p --output-format=json` and return the result.
+
+    The `runner` argument is only used by tests to inject a fake subprocess
+    runner; in production we go straight to subprocess.run.
+    """
+    if runner is None and not _claude_cli_available():
+        raise LlmReviewError(
+            "`claude` CLI was not found on PATH. Install Claude Code (https://claude.com/claude-code) "
+            "and run `claude login`, or set ANTHROPIC_API_KEY / ANTHROPIC_AUTH_TOKEN to use the SDK backend."
+        )
+
+    cmd = [
+        "claude",
+        "-p",
+        "--output-format=json",
+        "--model",
+        model,
+        "--system-prompt",
+        rendered_system,
+    ]
+
+    try:
+        if runner is None:
+            completed = subprocess.run(
+                cmd,
+                input=rendered_user,
+                text=True,
+                capture_output=True,
+                check=False,
+                timeout=CLAUDE_CLI_TIMEOUT_SECONDS,
+            )
+        else:
+            completed = runner(cmd, rendered_user)
+    except FileNotFoundError as exc:
+        raise LlmReviewError("`claude` CLI was not found on PATH.") from exc
+    except subprocess.TimeoutExpired as exc:
+        raise LlmReviewError(
+            f"`claude` CLI did not respond within {CLAUDE_CLI_TIMEOUT_SECONDS}s."
+        ) from exc
+
+    stdout = (completed.stdout or "").strip()
+    stderr = (completed.stderr or "").strip()
+    if completed.returncode != 0 and not stdout:
+        raise LlmReviewError(
+            f"`claude` CLI exited with status {completed.returncode}: {stderr or '<no stderr>'}"
+        )
+    if not stdout:
+        raise LlmReviewError("`claude` CLI produced no output.")
+
+    try:
+        payload = json.loads(stdout)
+    except json.JSONDecodeError as exc:
+        raise LlmReviewError(
+            f"`claude` CLI returned non-JSON output: {stdout[:200]}..."
+        ) from exc
+
+    if payload.get("is_error"):
+        errors = payload.get("errors") or [payload.get("subtype") or "unknown error"]
+        raise LlmReviewError(f"`claude` CLI reported an error: {'; '.join(str(e) for e in errors)}")
+
+    markdown = payload.get("result")
+    if not isinstance(markdown, str) or not markdown.strip():
+        raise LlmReviewError("`claude` CLI returned no `result` text.")
+
+    usage_payload = payload.get("usage") or {}
+    usage: dict[str, int] = {}
+    for field in (
+        "input_tokens",
+        "output_tokens",
+        "cache_creation_input_tokens",
+        "cache_read_input_tokens",
+    ):
+        value = usage_payload.get(field)
+        if isinstance(value, int):
+            usage[field] = value
+
+    cli_model = model
+    model_usage = payload.get("modelUsage")
+    if isinstance(model_usage, dict) and model_usage:
+        cli_model = next(iter(model_usage.keys()), model)
+
+    stop_reason = payload.get("stop_reason")
+    return LlmReviewResult(
+        markdown=markdown.strip(),
+        model=cli_model,
+        stop_reason=stop_reason if isinstance(stop_reason, str) else None,
+        usage=usage,
+        backend="cli",
+    )
+
+
+# ---------- Public entry point ----------
+
+
+def _resolve_backend(
+    backend: LlmBackend,
+    *,
+    api_key: str | None,
+    client: Any | None,
+) -> Literal["api", "cli"]:
+    if backend == "api":
+        return "api"
+    if backend == "cli":
+        if not _claude_cli_available():
+            raise LlmReviewError(
+                "`--llm-backend cli` was requested but the `claude` CLI is not on PATH."
+            )
+        return "cli"
+    # auto
+    if client is not None or api_key:
+        return "api"
+    if os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("ANTHROPIC_AUTH_TOKEN"):
+        return "api"
+    if _claude_cli_available():
+        return "cli"
+    raise LlmReviewError(
+        "No LLM backend available. Set ANTHROPIC_API_KEY, set ANTHROPIC_AUTH_TOKEN, "
+        "or install the `claude` CLI and run `claude login`."
+    )
+
+
+def generate_llm_review(
+    scan: ScanResult,
+    attack_surfaces: list[AttackSurface],
+    findings: list[Finding],
+    attack_paths: list[AttackPath],
+    *,
+    model: str | None = None,
+    effort: Literal["low", "medium", "high", "xhigh", "max"] | None = None,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    api_key: str | None = None,
+    client: Any | None = None,
+    backend: LlmBackend = "auto",
+    cli_runner: Any | None = None,
+) -> LlmReviewResult:
+    """Produce a narrative defensive review by calling Claude.
+
+    Resolves which backend to use (API SDK or `claude` CLI) based on available
+    auth, then runs the existing prompt pack through it. Streams the SDK path
+    so we never hit HTTP timeouts on long reviews. The CLI path runs synchronously
+    via `claude -p --output-format=json`.
+    """
+    resolved_model = model or os.environ.get("ATTACKMAP_LLM_MODEL") or DEFAULT_MODEL
+    resolved_effort = effort or DEFAULT_EFFORT
+
+    rendered = render_review_prompts(scan, attack_surfaces, findings, attack_paths)
+    if backend == "cli" and cli_runner is not None:
+        chosen_backend: Literal["api", "cli"] = "cli"
+    else:
+        chosen_backend = _resolve_backend(backend, api_key=api_key, client=client)
+
+    if chosen_backend == "api":
+        return _run_via_sdk(
+            rendered.system,
+            rendered.user,
+            model=resolved_model,
+            effort=resolved_effort,
+            max_tokens=max_tokens,
+            api_key=api_key,
+            client=client,
+        )
+    return _run_via_claude_cli(
+        rendered.system,
+        rendered.user,
+        model=resolved_model,
+        runner=cli_runner,
+    )
+
+
+__all__ = [
+    "generate_llm_review",
+    "LlmReviewError",
+    "LlmReviewResult",
+    "LlmBackend",
+    "DEFAULT_MODEL",
+    "DEFAULT_EFFORT",
+    "DEFAULT_MAX_TOKENS",
+]

attackmap/models.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class Route(BaseModel):
+    path: str
+    method: str = "ANY"
+    file: str
+    line: int | None = None
+
+
+class ExternalCall(BaseModel):
+    target: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+
+
+class DatabaseHint(BaseModel):
+    kind: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+
+
+class AuthHint(BaseModel):
+    hint: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.7
+
+
+class ServiceHint(BaseModel):
+    hint: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.7
+
+
+class EdgeHint(BaseModel):
+    hint: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.7
+
+
+class EntrypointHint(BaseModel):
+    hint: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.7
+
+
+class ProtocolHint(BaseModel):
+    hint: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.7
+
+
+class FrameworkHint(BaseModel):
+    hint: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.7
+
+
+class SecretHint(BaseModel):
+    name: str
+    file: str
+    line: int | None = None
+    evidence_text: str | None = None
+    confidence: float = 0.85
+
+
+SignalKind = Literal[
+    "route",
+    "external_call",
+    "database",
+    "auth",
+    "service",
+    "edge",
+    "entrypoint",
+    "protocol",
+    "framework",
+    "secret",
+]
+
+
+class Signal(BaseModel):
+    """Unified view of a single static-analysis signal.
+
+    Synthesized from the typed hint lists on `ScanResult` via `all_signals()`.
+    Existing analyzer plugins keep populating the typed hint lists; downstream
+    consumers (insights, controls, asset detection, prompts) can iterate over
+    the unified Signal stream to reason uniformly about location, confidence,
+    and evidence.
+    """
+
+    kind: SignalKind
+    label: str
+    file: str
+    line: int | None = None
+    confidence: float = 0.7
+    evidence_text: str | None = None
+    properties: dict[str, str] = Field(default_factory=dict)
+
+    def location(self) -> str:
+        """Stable `file:line` reference if line known, else just file."""
+        return f"{self.file}:{self.line}" if self.line is not None else self.file
+
+
+class AttackSurface(BaseModel):
+    route: str
+    method: str
+    file: str
+    category: Literal["webhook", "admin", "auth", "upload", "internal", "health", "public_api"]
+    exposure: Literal["public", "internal", "unknown"] = "public"
+    risk: Literal["low", "medium", "high"]
+    auth_signals: list[str] = Field(default_factory=list)
+    data_store_interaction: bool = False
+    outbound_integration: bool = False
+    rationale: list[str] = Field(default_factory=list)
+    line: int | None = None
+
+    def location(self) -> str:
+        return f"{self.file}:{self.line}" if self.line is not None else self.file
+
+
+class AttackTechnique(BaseModel):
+    """Reference to a MITRE ATT&CK technique (Enterprise matrix)."""
+
+    technique_id: str  # e.g., "T1190", "T1078.004"
+    name: str  # e.g., "Exploit Public-Facing Application"
+    tactic: str  # e.g., "Initial Access"
+    url: str | None = None  # https://attack.mitre.org/techniques/T1190/
+
+
+class Finding(BaseModel):
+    title: str
+    severity: Literal["low", "medium", "high"]
+    evidence: list[str] = Field(default_factory=list)
+    mitigation: str
+    confidence: Literal["low", "medium", "high"] = "medium"
+    attack_techniques: list[AttackTechnique] = Field(default_factory=list)
+
+
+class AttackPath(BaseModel):
+    name: str
+    steps: list[str]
+    impact: str
+
+
+AssetKind = Literal[
+    "credentials",
+    "session",
+    "user_pii",
+    "payment",
+    "internal_secret",
+    "audit_log",
+    "business_data",
+    "configuration",
+]
+
+
+class Asset(BaseModel):
+    id: str
+    kind: AssetKind
+    name: str
+    criticality: Literal["critical", "high", "medium", "low"]
+    locations: list[str] = Field(default_factory=list)
+    evidence: list[str] = Field(default_factory=list)
+
+
+ControlKind = Literal[
+    "authentication",
+    "authorization",
+    "input_validation",
+    "output_encoding",
+    "rate_limiting",
+    "csrf_protection",
+    "encryption_at_rest",
+    "encryption_in_transit",
+    "audit_logging",
+    "rbac",
+    "mfa",
+    "secret_management",
+    "security_headers",
+]
+
+
+ControlStrength = Literal["strong", "moderate", "weak", "absent"]
+
+
+class Control(BaseModel):
+    id: str
+    kind: ControlKind
+    name: str
+    strength: ControlStrength
+    scope: Literal["global", "module", "route", "service", "asset"]
+    placements: list[str] = Field(default_factory=list)
+    evidence: list[str] = Field(default_factory=list)
+    notes: str | None = None
+
+
+InsightKind = Literal[
+    "shared_secret_blast_radius",
+    "sensitive_asset_reachability",
+    "control_bypass",
+    "defense_gap_in_chain",
+    "asymmetric_protection",
+    "trust_boundary_violation",
+    "audit_gap",
+    "control_strength_mismatch",
+    "single_point_of_failure",
+    "stale_or_contradictory_signal",
+    "admin_action_without_auth",
+]
+
+
+class Insight(BaseModel):
+    id: str
+    kind: InsightKind
+    title: str
+    narrative: str
+    severity: Literal["critical", "high", "medium", "low", "informational"]
+    confidence: Literal["high", "medium", "low"]
+    evidence: list[str] = Field(default_factory=list)
+    related_assets: list[str] = Field(default_factory=list)
+    related_controls: list[str] = Field(default_factory=list)
+    related_routes: list[str] = Field(default_factory=list)
+    suggested_action: str | None = None
+    attack_techniques: list[AttackTechnique] = Field(default_factory=list)
+
+
+class DetectionOpportunity(BaseModel):
+    """A defender-facing detection-engineering hint produced from an insight or finding.
+
+    Designed to answer: "given this static-analysis finding, what runtime
+    detection signal would catch the same condition in production?"
+    """
+
+    id: str
+    title: str
+    rationale: str
+    signal_kind: Literal["log", "metric", "trace", "network", "config_audit"]
+    suggested_rule: str  # human-readable rule sketch (Sigma/KQL/Splunk-style)
+    related_insight_ids: list[str] = Field(default_factory=list)
+    related_finding_titles: list[str] = Field(default_factory=list)
+    attack_techniques: list[AttackTechnique] = Field(default_factory=list)
+
+
+class ScanResult(BaseModel):
+    root: str
+    languages: list[str] = Field(default_factory=list)
+    routes: list[Route] = Field(default_factory=list)
+    external_calls: list[ExternalCall] = Field(default_factory=list)
+    databases: list[DatabaseHint] = Field(default_factory=list)
+    auth_hints: list[AuthHint] = Field(default_factory=list)
+    service_hints: list[ServiceHint] = Field(default_factory=list)
+    edge_hints: list[EdgeHint] = Field(default_factory=list)
+    entrypoint_hints: list[EntrypointHint] = Field(default_factory=list)
+    protocol_hints: list[ProtocolHint] = Field(default_factory=list)
+    framework_hints: list[FrameworkHint] = Field(default_factory=list)
+    secret_hints: list[SecretHint] = Field(default_factory=list)
+    files_scanned: int = 0
+
+    @property
+    def root_path(self) -> Path:
+        return Path(self.root)
+
+    def all_signals(self) -> list[Signal]:
+        """Synthesize the unified Signal stream from typed hint lists.
+
+        Cheap O(n) view — does not mutate `self`. Order: routes, external calls,
+        databases, then the hint families (auth, service, edge, entrypoint,
+        protocol, framework, secret) in that order.
+        """
+        signals: list[Signal] = []
+        for r in self.routes:
+            signals.append(
+                Signal(
+                    kind="route",
+                    label=f"{r.method} {r.path}",
+                    file=r.file,
+                    line=r.line,
+                    properties={"method": r.method, "path": r.path},
+                )
+            )
+        for ext in self.external_calls:
+            signals.append(
+                Signal(
+                    kind="external_call",
+                    label=ext.target,
+                    file=ext.file,
+                    line=ext.line,
+                    evidence_text=ext.evidence_text,
+                )
+            )
+        for db in self.databases:
+            signals.append(
+                Signal(
+                    kind="database",
+                    label=db.kind,
+                    file=db.file,
+                    line=db.line,
+                    evidence_text=db.evidence_text,
+                    properties={"kind": db.kind},
+                )
+            )
+        _hint_kind_pairs: tuple[tuple[SignalKind, list], ...] = (
+            ("auth", self.auth_hints),
+            ("service", self.service_hints),
+            ("edge", self.edge_hints),
+            ("entrypoint", self.entrypoint_hints),
+            ("protocol", self.protocol_hints),
+            ("framework", self.framework_hints),
+        )
+        for kind, hints in _hint_kind_pairs:
+            for h in hints:
+                signals.append(
+                    Signal(
+                        kind=kind,
+                        label=h.hint,
+                        file=h.file,
+                        line=getattr(h, "line", None),
+                        evidence_text=getattr(h, "evidence_text", None),
+                        confidence=getattr(h, "confidence", 0.7),
+                    )
+                )
+        for s in self.secret_hints:
+            signals.append(
+                Signal(
+                    kind="secret",
+                    label=s.name,
+                    file=s.file,
+                    line=s.line,
+                    evidence_text=s.evidence_text,
+                    confidence=getattr(s, "confidence", 0.85),
+                )
+            )
+        return signals