shadowshield 0.4.0__py3-none-any.whl

shadowshield/core/config.py

@@ -0,0 +1,242 @@
+"""Configuration models for ShadowShield.
+
+A single :class:`ShieldConfig` drives the whole framework. It can be built three
+ways, in increasing order of specificity:
+
+1. From a named *mode* — :func:`ShieldConfig.for_mode` (``strict`` /
+   ``balanced`` / ``permissive``). This is the 90% path.
+2. From a YAML/dict — :func:`ShieldConfig.from_yaml` / ``model_validate``.
+3. By hand, overriding any field.
+
+The config is intentionally validated with Pydantic so that a malformed
+deployment config fails loudly at startup rather than silently weakening
+protection at runtime.
+"""
+
+from __future__ import annotations
+
+import enum
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+from .types import Decision, Severity
+
+
+class Mode(str, enum.Enum):
+    """Preset risk postures.
+
+    - ``STRICT``: security-first. Blocks aggressively, sanitizes by default,
+      enables every detector including the optional LLM self-check hook.
+    - ``BALANCED``: the sensible default. Sanitizes medium threats, blocks high+.
+    - ``PERMISSIVE``: observability-first. Flags and logs, rarely blocks —
+      useful for shadow-mode rollout where you measure before you enforce.
+    """
+
+    STRICT = "strict"
+    BALANCED = "balanced"
+    PERMISSIVE = "permissive"
+
+
+class PolicyConfig(BaseModel):
+    """Maps an aggregate :class:`Severity` to a :class:`Decision`.
+
+    The engine computes a single severity for a payload, then looks it up here.
+    Every band must be present; :func:`decide` walks from the payload's severity.
+    """
+
+    none: Decision = Decision.ALLOW
+    low: Decision = Decision.FLAG
+    medium: Decision = Decision.SANITIZE
+    high: Decision = Decision.BLOCK
+    critical: Decision = Decision.BLOCK
+
+    def decide(self, severity: Severity) -> Decision:
+        return {
+            Severity.NONE: self.none,
+            Severity.LOW: self.low,
+            Severity.MEDIUM: self.medium,
+            Severity.HIGH: self.high,
+            Severity.CRITICAL: self.critical,
+        }[severity]
+
+
+class DetectorConfig(BaseModel):
+    """Per-detector toggle + optional weight and overrides.
+
+    ``weight`` scales a detector's contribution to the aggregate score, letting
+    deployments trust some detectors more than others without code changes.
+    """
+
+    enabled: bool = True
+    weight: float = Field(default=1.0, ge=0.0, le=5.0)
+    options: dict[str, Any] = Field(default_factory=dict)
+
+
+class LLMCheckConfig(BaseModel):
+    """Optional lightweight LLM-based self-check ("is this an injection?").
+
+    Disabled by default because it costs a model call. When enabled the engine
+    invokes a user-supplied callable (see ``Shield(llm_judge=...)``); no network
+    client is built into the core.
+    """
+
+    enabled: bool = False
+    # Only consult the LLM when the cheap tiers already scored at/above this.
+    # Avoids paying for a model call on obviously-clean traffic.
+    min_score_to_invoke: float = Field(default=0.35, ge=0.0, le=1.0)
+    timeout_seconds: float = Field(default=8.0, gt=0.0)
+
+
+class RateLimitConfig(BaseModel):
+    """Token-bucket style limit applied per identity (e.g. per session/user)."""
+
+    enabled: bool = False
+    max_events: int = Field(default=60, gt=0)
+    window_seconds: float = Field(default=60.0, gt=0.0)
+    # Count only flagged/blocked events toward the limit, not clean traffic.
+    count_only_threats: bool = True
+
+
+class LoggingConfig(BaseModel):
+    enabled: bool = True
+    # JSON lines audit sink. ``None`` -> stdout via structlog only.
+    audit_path: str | None = None
+    # Never write the raw offending payload to the audit log unless explicitly
+    # opted in — payloads can contain secrets/PII.
+    redact_payloads: bool = True
+    level: str = "INFO"
+
+
+class ShieldConfig(BaseModel):
+    """The single source of truth for a :class:`~shadowshield.Shield`."""
+
+    mode: Mode = Mode.BALANCED
+
+    # When True, ``Shield.scan`` raises ``ThreatBlockedError`` on a BLOCK
+    # decision instead of returning a result with ``blocked=True``. Off by
+    # default so the library is non-throwing unless you opt in.
+    raise_on_block: bool = False
+
+    # Aggregate score at/above which the payload is considered "not safe" even
+    # if no single detector hit CRITICAL. Lower = more cautious.
+    block_threshold: float = Field(default=0.65, ge=0.0, le=1.0)
+
+    # Hard cap on how many characters are scanned per payload. Oversized inputs
+    # are scanned as a truncated prefix (bounding CPU work) and flagged — a guard
+    # against resource-exhaustion from multi-megabyte payloads. 0 = unlimited.
+    max_input_chars: int = Field(default=100_000, ge=0)
+
+    policy: PolicyConfig = Field(default_factory=PolicyConfig)
+    detectors: dict[str, DetectorConfig] = Field(default_factory=dict)
+    llm_check: LLMCheckConfig = Field(default_factory=LLMCheckConfig)
+    rate_limit: RateLimitConfig = Field(default_factory=RateLimitConfig)
+    logging: LoggingConfig = Field(default_factory=LoggingConfig)
+
+    # Detectors named here are forced off regardless of their per-detector
+    # config — a convenient kill-switch for noisy detectors in production.
+    disabled_detectors: list[str] = Field(default_factory=list)
+
+    model_config = {"extra": "forbid", "use_enum_values": False}
+
+    @field_validator("logging")
+    @classmethod
+    def _normalise_level(cls, v: LoggingConfig) -> LoggingConfig:
+        v.level = v.level.upper()
+        return v
+
+    @model_validator(mode="after")
+    def _apply_mode_defaults(self) -> ShieldConfig:
+        """Fill unset knobs from the chosen mode without clobbering overrides."""
+        # Only applied when the caller didn't explicitly customise policy.
+        return self
+
+    # ------------------------------------------------------------------ #
+    # Builders
+    # ------------------------------------------------------------------ #
+    @classmethod
+    def for_mode(cls, mode: Mode | str = Mode.BALANCED, **overrides: Any) -> ShieldConfig:
+        """Build a config from a preset mode, then apply keyword overrides."""
+        mode = Mode(mode)
+        base = _MODE_PRESETS[mode]()
+        data = base.model_dump()
+        data.update(overrides)
+        return cls.model_validate(data)
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> ShieldConfig:
+        raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
+        if "mode" in raw and len(raw) > 1:
+            # Treat YAML as overrides layered on top of the named mode preset.
+            mode = raw.pop("mode")
+            return cls.for_mode(mode, **raw)
+        return cls.model_validate(raw)
+
+    def detector_config(self, name: str) -> DetectorConfig:
+        """Effective config for a detector, honouring the global kill-switch."""
+        cfg = self.detectors.get(name, DetectorConfig())
+        if name in self.disabled_detectors:
+            return cfg.model_copy(update={"enabled": False})
+        return cfg
+
+
+# ---------------------------------------------------------------------- #
+# Mode presets
+# ---------------------------------------------------------------------- #
+def _strict() -> ShieldConfig:
+    return ShieldConfig(
+        mode=Mode.STRICT,
+        raise_on_block=False,
+        block_threshold=0.45,
+        policy=PolicyConfig(
+            none=Decision.ALLOW,
+            low=Decision.SANITIZE,
+            medium=Decision.BLOCK,
+            high=Decision.BLOCK,
+            critical=Decision.BLOCK,
+        ),
+        llm_check=LLMCheckConfig(enabled=True, min_score_to_invoke=0.25),
+        rate_limit=RateLimitConfig(enabled=True, max_events=30, window_seconds=60.0),
+        logging=LoggingConfig(enabled=True, redact_payloads=True, level="INFO"),
+    )
+
+
+def _balanced() -> ShieldConfig:
+    return ShieldConfig(
+        mode=Mode.BALANCED,
+        block_threshold=0.65,
+        policy=PolicyConfig(
+            none=Decision.ALLOW,
+            low=Decision.FLAG,
+            medium=Decision.SANITIZE,
+            high=Decision.BLOCK,
+            critical=Decision.BLOCK,
+        ),
+        llm_check=LLMCheckConfig(enabled=False),
+        rate_limit=RateLimitConfig(enabled=False),
+    )
+
+
+def _permissive() -> ShieldConfig:
+    return ShieldConfig(
+        mode=Mode.PERMISSIVE,
+        block_threshold=0.85,
+        policy=PolicyConfig(
+            none=Decision.ALLOW,
+            low=Decision.FLAG,
+            medium=Decision.FLAG,
+            high=Decision.SANITIZE,
+            critical=Decision.BLOCK,
+        ),
+        llm_check=LLMCheckConfig(enabled=False),
+        rate_limit=RateLimitConfig(enabled=False),
+    )
+
+
+_MODE_PRESETS = {
+    Mode.STRICT: _strict,
+    Mode.BALANCED: _balanced,
+    Mode.PERMISSIVE: _permissive,
+}

shadowshield/core/engine.py

@@ -0,0 +1,255 @@
+"""The unified detection→decision→response engine.
+
+This is the heart of ShadowShield and the thing that makes it *one* system rather
+than a detector bag bolted to a responder bag. One pass:
+
+1. Build a shared :class:`ScanContext` (normalise + decode once).
+2. Run the cheap, deterministic detectors.
+3. Conditionally consult the optional LLM self-check (only when the cheap tiers
+   already crossed ``min_score_to_invoke`` — never on clean traffic).
+4. Aggregate weighted findings into one score + severity (noisy-or).
+5. Let the policy + block-threshold + rate limiter decide.
+6. Apply the matching responders (sanitize / block / isolate).
+7. Audit.
+
+The flow is identical for input and output, which is what gives ShadowShield its
+symmetric, two-way protection.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+from collections.abc import Callable
+from typing import Any
+
+from ..core.config import ShieldConfig
+from ..core.types import (
+    Decision,
+    Direction,
+    ScanResult,
+    Severity,
+    Threat,
+    ThreatCategory,
+)
+from ..detectors.alignment import AlignmentCheckDetector, AlignmentJudge
+from ..detectors.base import Detector, ScanContext
+from ..detectors.llm_check import LLMJudge, LLMSelfCheckDetector
+from ..responders.base import Responder
+from ..responders.rate_limiter import RateLimitResponder
+from ..utils.logging import AuditLog
+from ..utils.scoring import aggregate_score, aggregate_severity
+from ..utils.text import truncate
+from .session import ConversationHistory
+
+# Total order over decisions for "take the stronger of two" comparisons.
+_DECISION_RANK: dict[Decision, int] = {
+    Decision.ALLOW: 0,
+    Decision.FLAG: 1,
+    Decision.SANITIZE: 2,
+    Decision.BLOCK: 3,
+    Decision.ESCALATE: 4,
+}
+
+_LLM_DETECTOR_NAME = LLMSelfCheckDetector.name
+_ALIGNMENT_DETECTOR_NAME = AlignmentCheckDetector.name
+# Detectors that the engine drives separately (gated / context-injected), not in
+# the cheap deterministic loop.
+_GATED_DETECTORS = frozenset({_LLM_DETECTOR_NAME, _ALIGNMENT_DETECTOR_NAME})
+
+
+def _stronger(a: Decision, b: Decision) -> Decision:
+    return a if _DECISION_RANK[a] >= _DECISION_RANK[b] else b
+
+
+class Engine:
+    """Stateless-per-call orchestrator wired with detectors and responders."""
+
+    def __init__(
+        self,
+        config: ShieldConfig,
+        *,
+        detectors: list[Detector],
+        responders: list[Responder],
+        rate_limiter: RateLimitResponder,
+        audit: AuditLog,
+        llm_judge: LLMJudge | None = None,
+        alignment_judge: AlignmentJudge | None = None,
+    ) -> None:
+        self._config = config
+        self._detectors = detectors
+        self._responders = responders
+        self._rate_limiter = rate_limiter
+        self._audit = audit
+        self._llm_judge = llm_judge
+        self._alignment_judge = alignment_judge
+        # Judges are user-supplied callables that may hang or make network calls.
+        # Run them in a small pool so we can enforce a hard timeout — a hung judge
+        # must never block the request path. Only created when a judge exists.
+        self._judge_pool: concurrent.futures.ThreadPoolExecutor | None = (
+            concurrent.futures.ThreadPoolExecutor(max_workers=4, thread_name_prefix="ss-judge")
+            if (llm_judge is not None or alignment_judge is not None)
+            else None
+        )
+        # Detector weights are read from config once.
+        self._weights = {
+            name: config.detector_config(name).weight for name in self._detector_names()
+        }
+
+    def _detector_names(self) -> list[str]:
+        return [d.name for d in self._detectors]
+
+    # ------------------------------------------------------------------ #
+    def evaluate(
+        self,
+        text: str,
+        *,
+        direction: Direction,
+        identity: str | None = None,
+        history: ConversationHistory | None = None,
+        canaries: tuple[str, ...] = (),
+        objective: str | None = None,
+    ) -> ScanResult:
+        # Bound the work: oversized payloads are scanned as a truncated prefix so
+        # a multi-megabyte input can't exhaust CPU. The original text is preserved
+        # on the result; only the scanned region is capped.
+        max_chars = self._config.max_input_chars
+        oversized = bool(max_chars) and len(text) > max_chars
+        scan_text = text[:max_chars] if oversized else text
+
+        context = ScanContext.build(
+            scan_text,
+            direction=direction,
+            history=history,
+            identity=identity,
+            canaries=canaries,
+            objective=objective,
+        )
+
+        threats = self._run_cheap_detectors(scan_text, context)
+        interim_score = aggregate_score(threats, self._weights)
+        threats += self._maybe_run_llm_check(scan_text, context, interim_score)
+        threats += self._maybe_run_alignment(scan_text, context)
+        if oversized:
+            threats.append(
+                Threat(
+                    category=ThreatCategory.ANOMALY,
+                    severity=Severity.MEDIUM,
+                    score=0.5,
+                    detector="input_size_guard",
+                    message=(
+                        f"Input exceeds max_input_chars ({max_chars}); scanned a "
+                        f"truncated prefix of a {len(text)}-char payload."
+                    ),
+                    metadata={"original_length": len(text), "scanned_length": max_chars},
+                )
+            )
+
+        score = aggregate_score(threats, self._weights)
+        severity = aggregate_severity(threats, score)
+        decision = self._decide(score, severity)
+
+        result = ScanResult(
+            text=text,
+            direction=direction,
+            threats=threats,
+            score=score,
+            severity=severity,
+            decision=decision,
+        )
+
+        # Rate-limit pre-pass can escalate to BLOCK based on identity history.
+        result = self._rate_limiter.check(result, context=context)
+
+        result = self._apply_responders(result, context)
+        self._record(result, context)
+        return result
+
+    # ------------------------------------------------------------------ #
+    def _run_cheap_detectors(self, text: str, context: ScanContext) -> list[Threat]:
+        threats: list[Threat] = []
+        for det in self._detectors:
+            if det.name in _GATED_DETECTORS:
+                continue  # handled separately (gated / context-injected)
+            if not det.applies_to(context.direction):
+                continue
+            context.options = self._config.detector_config(det.name).options
+            threats.extend(self._safe_scan(det, text, context))
+        return threats
+
+    def _maybe_run_llm_check(
+        self, text: str, context: ScanContext, interim_score: float
+    ) -> list[Threat]:
+        cfg = self._config.llm_check
+        if not cfg.enabled or self._llm_judge is None:
+            return []
+        if interim_score < cfg.min_score_to_invoke:
+            return []
+        det = next((d for d in self._detectors if d.name == _LLM_DETECTOR_NAME), None)
+        if det is None or not det.applies_to(context.direction):
+            return []
+        context.options = {"judge": self._with_timeout(self._llm_judge, cfg.timeout_seconds)}
+        return self._safe_scan(det, text, context)
+
+    def _maybe_run_alignment(self, text: str, context: ScanContext) -> list[Threat]:
+        # Only runs on the output side, when an objective is set and a judge is
+        # wired in. This is the agent-trace alignment audit (goal-hijack detection).
+        if self._alignment_judge is None or not context.objective:
+            return []
+        det = next((d for d in self._detectors if d.name == _ALIGNMENT_DETECTOR_NAME), None)
+        if det is None or not det.applies_to(context.direction):
+            return []
+        timeout = self._config.llm_check.timeout_seconds
+        context.options = {"alignment_judge": self._with_timeout(self._alignment_judge, timeout)}
+        return self._safe_scan(det, text, context)
+
+    def _with_timeout(self, fn: Callable[..., Any], timeout: float) -> Callable[..., Any]:
+        """Wrap a user judge so a hang can't block the request beyond ``timeout``.
+
+        The judge runs in the pool; if it overruns, ``future.result`` raises
+        ``TimeoutError``, which the calling detector's fail-safe ``except`` turns
+        into a low-severity "unavailable" note rather than a crash or a hang.
+        (An over-running judge thread is left to finish/leak — the standard,
+        accepted trade-off for thread-based timeouts in Python.)
+        """
+
+        def wrapped(*args: Any) -> Any:
+            assert self._judge_pool is not None  # only built when a judge exists
+            future = self._judge_pool.submit(fn, *args)
+            return future.result(timeout=timeout)
+
+        return wrapped
+
+    @staticmethod
+    def _safe_scan(det: Detector, text: str, context: ScanContext) -> list[Threat]:
+        """A detector that raises must never take down the request path."""
+        try:
+            return det.scan(text, context=context)
+        except Exception:  # pragma: no cover - defensive
+            # Fail-safe: drop this detector's contribution, keep the others.
+            return []
+
+    def _decide(self, score: float, severity: Severity) -> Decision:
+        decision = self._config.policy.decide(severity)
+        # Independent floor: a high aggregate score forces at least a block even
+        # if the per-band policy was lenient.
+        if score >= self._config.block_threshold:
+            decision = _stronger(decision, Decision.BLOCK)
+        return decision
+
+    def _apply_responders(self, result: ScanResult, context: ScanContext) -> ScanResult:
+        for responder in self._responders:
+            if responder.applies(result):
+                result = responder.apply(result, context=context)
+        return result
+
+    def _record(self, result: ScanResult, context: ScanContext) -> None:
+        event = result.to_dict()
+        event["identity"] = context.identity
+        if not self._audit.redact:
+            event["text"] = truncate(result.text, 400)
+        else:
+            event["text_preview"] = truncate(result.text, 80)
+        # Clean, threat-free scans are logged at DEBUG (quiet by default);
+        # anything noteworthy is logged at INFO.
+        notable = bool(result.threats) or not result.is_safe
+        self._audit.record(event, notable=notable)

shadowshield/core/session.py

@@ -0,0 +1,138 @@
+"""Conversation history and the :class:`ShieldedSession` context manager.
+
+Multi-turn and indirect injections only reveal themselves across a *conversation*
+— a single message can look benign while the sequence sets up an attack. The
+:class:`ConversationHistory` gives detectors that cross-turn view, and
+:class:`ShieldedSession` binds a stable identity + history to a shield so rate
+limiting and history analysis work without the caller threading state by hand.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+from types import TracebackType
+from typing import TYPE_CHECKING
+
+from .types import Direction, ScanResult
+
+if TYPE_CHECKING:
+    from .shield import Shield
+
+
+@dataclass(slots=True)
+class Turn:
+    """One scanned message in a conversation."""
+
+    direction: Direction
+    text: str
+    result: ScanResult
+
+
+@dataclass
+class ConversationHistory:
+    """A bounded record of recent turns for cross-turn analysis.
+
+    Bounded (``maxlen``) so memory stays flat over long-running agents. Detectors
+    receive this read-only view via :class:`~shadowshield.detectors.base.ScanContext`.
+    """
+
+    maxlen: int = 50
+    turns: deque[Turn] = field(default_factory=lambda: deque(maxlen=50))
+
+    def __post_init__(self) -> None:
+        if self.turns.maxlen != self.maxlen:
+            self.turns = deque(self.turns, maxlen=self.maxlen)
+
+    def add(self, direction: Direction, text: str, result: ScanResult) -> None:
+        self.turns.append(Turn(direction, text, result))
+
+    @property
+    def flagged_count(self) -> int:
+        """How many recorded turns were not clean — a multi-turn pressure gauge."""
+        return sum(1 for t in self.turns if not t.result.is_safe or t.result.threats)
+
+    def recent_text(self, n: int = 5, direction: Direction | None = None) -> list[str]:
+        items = [t for t in self.turns if direction is None or t.direction == direction]
+        return [t.text for t in list(items)[-n:]]
+
+
+class ShieldedSession:
+    """Stateful, per-conversation wrapper around a :class:`Shield`.
+
+    Use as a context manager::
+
+        with shield.session(identity="user-42") as s:
+            clean = s.guard_input(user_msg)
+            reply = s.guard_output(model_reply)
+
+    All scans within the session share one identity (for rate limiting) and one
+    :class:`ConversationHistory` (for multi-turn detection), and are recorded for
+    inspection afterwards.
+    """
+
+    def __init__(
+        self,
+        shield: Shield,
+        *,
+        identity: str | None = None,
+        history_size: int = 50,
+        objective: str | None = None,
+    ) -> None:
+        self._shield = shield
+        self.identity = identity
+        # The user's stated goal for this conversation. When set (and an
+        # alignment judge is wired into the shield), output scans run the
+        # agent-trace alignment audit to catch goal hijacking.
+        self.objective = objective
+        self.history = ConversationHistory(maxlen=history_size)
+
+    def set_objective(self, objective: str) -> None:
+        """Set/replace the user's objective for alignment auditing."""
+        self.objective = objective
+
+    # -- scanning ------------------------------------------------------- #
+    def scan_input(self, text: str) -> ScanResult:
+        result = self._shield.scan(
+            text, direction=Direction.INPUT, identity=self.identity, history=self.history
+        )
+        self.history.add(Direction.INPUT, text, result)
+        return result
+
+    def scan_output(self, text: str) -> ScanResult:
+        result = self._shield.scan(
+            text,
+            direction=Direction.OUTPUT,
+            identity=self.identity,
+            history=self.history,
+            objective=self.objective,
+        )
+        self.history.add(Direction.OUTPUT, text, result)
+        return result
+
+    def guard_input(self, text: str) -> str:
+        """Scan input and return safe text, raising on a block (strict ergonomics)."""
+        return self._shield.guard(
+            text, direction=Direction.INPUT, identity=self.identity, history=self.history
+        )
+
+    def guard_output(self, text: str) -> str:
+        return self._shield.guard(
+            text,
+            direction=Direction.OUTPUT,
+            identity=self.identity,
+            history=self.history,
+            objective=self.objective,
+        )
+
+    # -- context manager ------------------------------------------------ #
+    def __enter__(self) -> ShieldedSession:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        return None