shadowshield 0.4.0__py3-none-any.whl

shadowshield/core/shield.py

@@ -0,0 +1,359 @@
+"""The public :class:`Shield` — one object, the whole framework.
+
+``Shield`` assembles the configured detectors, responders, rate limiter, audit
+sink, and the engine, then exposes a small, ergonomic surface:
+
+- :meth:`scan` — full :class:`ScanResult` (the power-user path).
+- :meth:`guard` — return safe text, **raise** on block (fail-closed ergonomics).
+- :meth:`filter` — return safe text, **never raise** (fail-soft: fallback on block).
+- :meth:`isolate` — spotlight untrusted text for safe passthrough.
+- :meth:`session` — a stateful :class:`ShieldedSession` for one conversation.
+- :meth:`protect` — a decorator that shields a function's string I/O.
+
+Everything funnels through the single engine, so input and output get identical,
+consistent treatment.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import json
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from ..detectors.alignment import AlignmentJudge
+from ..detectors.base import Detector, build_detectors
+from ..detectors.llm_check import LLMJudge
+from ..responders.base import Responder
+from ..responders.blocker import BlockResponder
+from ..responders.isolator import IsolationResponder, spotlight
+from ..responders.rate_limiter import RateLimitResponder
+from ..responders.sanitizer import SanitizeResponder
+from ..utils.logging import AuditLog
+from .canary import CanaryRegistry, CanaryToken
+from .config import Mode, ShieldConfig
+from .engine import Engine
+from .session import ConversationHistory, ShieldedSession
+from .types import Direction, ScanResult, ThreatBlockedError
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class Shield:
+    """The unified ShadowShield entry point."""
+
+    def __init__(
+        self,
+        config: ShieldConfig | None = None,
+        *,
+        llm_judge: LLMJudge | None = None,
+        alignment_judge: AlignmentJudge | None = None,
+        extra_detectors: list[Detector] | None = None,
+        extra_responders: list[Responder] | None = None,
+        isolate_flagged: bool = False,
+        use_transformer: bool | str = False,
+        use_vectors: bool | str = False,
+    ) -> None:
+        self.config = config or ShieldConfig.for_mode(Mode.BALANCED)
+        self._audit = AuditLog(self.config.logging)
+        self.canaries = CanaryRegistry()
+
+        detectors = build_detectors(
+            is_enabled=lambda name: self.config.detector_config(name).enabled
+        )
+        # Opt-in ML classifier layer. ``use_transformer`` may be True (default
+        # model) or a model id string. Kept out of the auto-registered set so it
+        # never triggers a surprise model download.
+        if use_transformer:
+            from ..detectors.transformer import TransformerDetector
+
+            model = use_transformer if isinstance(use_transformer, str) else None
+            detectors.append(TransformerDetector(model) if model else TransformerDetector())
+        # Opt-in vector-similarity layer (paraphrase / cross-lingual matching).
+        if use_vectors:
+            from ..detectors.vector import VectorSimilarityDetector
+
+            vmodel = use_vectors if isinstance(use_vectors, str) else None
+            detectors.append(
+                VectorSimilarityDetector(vmodel) if vmodel else VectorSimilarityDetector()
+            )
+        if extra_detectors:
+            detectors.extend(extra_detectors)
+        self._detectors = detectors
+
+        self._rate_limiter = RateLimitResponder(self.config.rate_limit)
+
+        # Responder order matters: sanitize/isolate transform text; block has the
+        # final say and overwrites with a fallback.
+        responders: list[Responder] = [SanitizeResponder()]
+        if isolate_flagged:
+            responders.append(IsolationResponder())
+        responders.append(BlockResponder())
+        if extra_responders:
+            responders.extend(extra_responders)
+        self._responders = responders
+
+        self._engine = Engine(
+            self.config,
+            detectors=self._detectors,
+            responders=self._responders,
+            rate_limiter=self._rate_limiter,
+            audit=self._audit,
+            llm_judge=llm_judge,
+            alignment_judge=alignment_judge,
+        )
+
+    # ------------------------------------------------------------------ #
+    # Builders
+    # ------------------------------------------------------------------ #
+    @classmethod
+    def for_mode(cls, mode: Mode | str = Mode.BALANCED, **kwargs: Any) -> Shield:
+        return cls(ShieldConfig.for_mode(mode), **kwargs)
+
+    @classmethod
+    def from_yaml(cls, path: str, **kwargs: Any) -> Shield:
+        return cls(ShieldConfig.from_yaml(path), **kwargs)
+
+    @property
+    def detectors(self) -> list[Detector]:
+        return list(self._detectors)
+
+    # ------------------------------------------------------------------ #
+    # Core scanning
+    # ------------------------------------------------------------------ #
+    def scan(
+        self,
+        text: str,
+        *,
+        direction: Direction | str = Direction.INPUT,
+        identity: str | None = None,
+        history: ConversationHistory | None = None,
+        objective: str | None = None,
+    ) -> ScanResult:
+        """Scan ``text`` and return the full :class:`ScanResult`.
+
+        Raises :class:`ThreatBlockedError` only when ``config.raise_on_block`` is
+        set *and* the result is blocked; otherwise always returns a result.
+
+        Pass ``objective`` (the user's stated goal) on output scans to enable the
+        agent-trace alignment audit (requires ``Shield(alignment_judge=...)``).
+        """
+        result = self._engine.evaluate(
+            text,
+            direction=Direction(direction),
+            identity=identity,
+            history=history,
+            canaries=self.canaries.active(),
+            objective=objective,
+        )
+        if self.config.raise_on_block and result.blocked:
+            raise ThreatBlockedError(result)
+        return result
+
+    def scan_input(self, text: str, **kwargs: Any) -> ScanResult:
+        return self.scan(text, direction=Direction.INPUT, **kwargs)
+
+    def scan_output(self, text: str, **kwargs: Any) -> ScanResult:
+        return self.scan(text, direction=Direction.OUTPUT, **kwargs)
+
+    # ------------------------------------------------------------------ #
+    # Ergonomic wrappers
+    # ------------------------------------------------------------------ #
+    def guard(
+        self,
+        text: str,
+        *,
+        direction: Direction | str = Direction.INPUT,
+        identity: str | None = None,
+        history: ConversationHistory | None = None,
+        objective: str | None = None,
+    ) -> str:
+        """Return safe text, **raising** :class:`ThreatBlockedError` on a block.
+
+        Fail-closed ergonomics: the dangerous path is the exceptional path. Use
+        this when you'd rather abort than risk passing tainted content.
+        """
+        result = self._engine.evaluate(
+            text,
+            direction=Direction(direction),
+            identity=identity,
+            history=history,
+            canaries=self.canaries.active(),
+            objective=objective,
+        )
+        if result.blocked:
+            raise ThreatBlockedError(result)
+        return result.safe_text
+
+    def filter(
+        self,
+        text: str,
+        *,
+        direction: Direction | str = Direction.INPUT,
+        identity: str | None = None,
+        history: ConversationHistory | None = None,
+        objective: str | None = None,
+    ) -> str:
+        """Return safe text, **never raising**.
+
+        Fail-soft ergonomics: a blocked payload yields the safe fallback string
+        instead of an exception. Use this on paths that must always return text.
+        """
+        result = self._engine.evaluate(
+            text,
+            direction=Direction(direction),
+            identity=identity,
+            history=history,
+            canaries=self.canaries.active(),
+            objective=objective,
+        )
+        return result.safe_text
+
+    def isolate(self, text: str, *, datamark: bool = False) -> str:
+        """Spotlight untrusted ``text`` so it's safer to feed to a model."""
+        return spotlight(text, datamark=datamark)
+
+    # ------------------------------------------------------------------ #
+    # Canary tokens (detect *successful* injections)
+    # ------------------------------------------------------------------ #
+    def issue_canary(self, *, prefix: str = "ss-canary", note: str = "") -> CanaryToken:
+        """Mint a canary token to embed in a system prompt / hidden context.
+
+        Embed ``token.instruction()`` in your prompt; if the token later appears
+        in model output (or a guarded tool call), :class:`CanaryLeakDetector`
+        flags a confirmed exfiltration. See :mod:`shadowshield.core.canary`.
+        """
+        return self.canaries.issue(prefix=prefix, note=note)
+
+    def harden(self, text: str) -> bool:
+        """Teach the vector-similarity layer a confirmed attack (self-hardening).
+
+        Appends ``text`` to every active :class:`VectorSimilarityDetector`'s index
+        so future inputs resembling it (or its paraphrases/translations) match.
+        Call after an incident — e.g. a canary-caught exfiltration. Returns True if
+        a vector detector was present to harden; False otherwise (no-op).
+        """
+        from ..detectors.vector import VectorSimilarityDetector
+
+        hardened = False
+        for det in self._detectors:
+            if isinstance(det, VectorSimilarityDetector):
+                det.add_attack(text)
+                hardened = True
+        return hardened
+
+    # ------------------------------------------------------------------ #
+    # Agentic guarding: tool calls & tool results are untrusted too
+    # ------------------------------------------------------------------ #
+    def scan_tool_call(
+        self, name: str, arguments: Any, *, identity: str | None = None
+    ) -> ScanResult:
+        """Scan an outbound tool/function call (name + arguments) as input.
+
+        Agentic injections often surface as a *tool call* the model was tricked
+        into making. Serialise the call and scan it so a poisoned argument or a
+        leaked canary in a tool payload is caught before the tool runs.
+        """
+        payload = self._stringify_tool(name, arguments)
+        return self.scan(payload, direction=Direction.INPUT, identity=identity)
+
+    def scan_tool_result(
+        self, name: str, result: Any, *, identity: str | None = None
+    ) -> ScanResult:
+        """Scan a tool/function *result* (untrusted external content) as input.
+
+        Tool results — web pages, file contents, API responses — are a primary
+        indirect-injection vector. Treat them as untrusted input.
+        """
+        payload = self._stringify_tool(name, result)
+        return self.scan(payload, direction=Direction.INPUT, identity=identity)
+
+    @staticmethod
+    def _stringify_tool(name: str, data: Any) -> str:
+        if isinstance(data, str):
+            body = data
+        else:
+            try:
+                body = json.dumps(data, ensure_ascii=False, default=str)
+            except (TypeError, ValueError):
+                body = str(data)
+        return f"{name}: {body}"
+
+    # ------------------------------------------------------------------ #
+    # Async API (non-blocking for event-loop apps: FastAPI, async agents)
+    # ------------------------------------------------------------------ #
+    async def ascan(self, text: str, **kwargs: Any) -> ScanResult:
+        """Async :meth:`scan` — runs the CPU-bound scan off the event loop."""
+        return await asyncio.to_thread(self.scan, text, **kwargs)
+
+    async def aguard(self, text: str, **kwargs: Any) -> str:
+        """Async :meth:`guard`."""
+        return await asyncio.to_thread(self.guard, text, **kwargs)
+
+    async def afilter(self, text: str, **kwargs: Any) -> str:
+        """Async :meth:`filter`."""
+        return await asyncio.to_thread(self.filter, text, **kwargs)
+
+    # ------------------------------------------------------------------ #
+    # Sessions & decorators
+    # ------------------------------------------------------------------ #
+    def session(
+        self,
+        *,
+        identity: str | None = None,
+        history_size: int = 50,
+        objective: str | None = None,
+    ) -> ShieldedSession:
+        return ShieldedSession(
+            self, identity=identity, history_size=history_size, objective=objective
+        )
+
+    def protect(
+        self,
+        func: F | None = None,
+        *,
+        input_arg: str | int | None = 0,
+        check_output: bool = True,
+        identity: str | None = None,
+    ) -> Any:
+        """Decorator that shields a function's string input and output.
+
+        By default the first positional argument is treated as untrusted input
+        and scanned (raising on block); if the return value is a string it is
+        scanned as output. Point ``input_arg`` at a keyword name or another
+        positional index to shield a different argument.
+
+        Usage::
+
+            @shield.protect
+            def ask(prompt: str) -> str: ...
+
+            @shield.protect(input_arg="question", check_output=False)
+            def ask2(*, question: str) -> str: ...
+        """
+
+        def decorate(fn: F) -> F:
+            @functools.wraps(fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                value = self._extract_arg(args, kwargs, input_arg)
+                if isinstance(value, str):
+                    self.guard(value, direction=Direction.INPUT, identity=identity)
+                out = fn(*args, **kwargs)
+                if check_output and isinstance(out, str):
+                    return self.guard(out, direction=Direction.OUTPUT, identity=identity)
+                return out
+
+            return wrapper  # type: ignore[return-value]
+
+        return decorate if func is None else decorate(func)
+
+    @staticmethod
+    def _extract_arg(
+        args: tuple[Any, ...], kwargs: dict[str, Any], selector: str | int | None
+    ) -> Any:
+        if selector is None:
+            return None
+        if isinstance(selector, int):
+            return args[selector] if len(args) > selector else None
+        return kwargs.get(selector)

shadowshield/core/types.py

@@ -0,0 +1,222 @@
+"""Shared, dependency-light value types used across every ShadowShield layer.
+
+These types are the *lingua franca* of the framework: detectors emit
+:class:`Threat` objects, the engine aggregates them into a :class:`ScanResult`,
+the policy turns severity into a :class:`Decision`, and responders act on it.
+
+Everything here is a plain dataclass / enum so the types stay importable with
+zero heavy dependencies and are trivially serialisable for logging.
+"""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass, field
+from typing import Any
+
+
+class Direction(str, enum.Enum):
+    """Which side of the LLM boundary a piece of text is on.
+
+    ``INPUT`` is untrusted content flowing *toward* the model (user prompts,
+    retrieved documents, tool results). ``OUTPUT`` is content flowing *back*
+    from the model toward the user / downstream tools.
+    """
+
+    INPUT = "input"
+    OUTPUT = "output"
+
+
+class ThreatCategory(str, enum.Enum):
+    """Taxonomy of the threats ShadowShield reasons about."""
+
+    PROMPT_INJECTION = "prompt_injection"
+    INDIRECT_INJECTION = "indirect_injection"
+    JAILBREAK = "jailbreak"
+    ROLE_MANIPULATION = "role_manipulation"
+    DELIMITER_ATTACK = "delimiter_attack"
+    ENCODING_OBFUSCATION = "encoding_obfuscation"
+    DATA_EXFILTRATION = "data_exfiltration"
+    SECRET_LEAK = "secret_leak"
+    PII_LEAK = "pii_leak"
+    CANARY_TOKEN = "canary_token"
+    ANOMALY = "anomaly"
+    UNKNOWN = "unknown"
+
+
+class Severity(enum.IntEnum):
+    """Ordered severity. ``IntEnum`` so thresholds compare naturally.
+
+    The numeric value doubles as a coarse weight when aggregating scores.
+    """
+
+    NONE = 0
+    LOW = 1
+    MEDIUM = 2
+    HIGH = 3
+    CRITICAL = 4
+
+    @property
+    def label(self) -> str:
+        return self.name.lower()
+
+    @classmethod
+    def from_score(cls, score: float) -> Severity:
+        """Map a 0..1 confidence score onto a severity band."""
+        if score >= 0.85:
+            return cls.CRITICAL
+        if score >= 0.65:
+            return cls.HIGH
+        if score >= 0.40:
+            return cls.MEDIUM
+        if score > 0.0:
+            return cls.LOW
+        return cls.NONE
+
+
+class Decision(str, enum.Enum):
+    """What the policy decided should happen to a scanned payload."""
+
+    ALLOW = "allow"
+    """Clean. Pass through untouched."""
+
+    FLAG = "flag"
+    """Pass through, but record/alert — suspicious yet below the action bar."""
+
+    SANITIZE = "sanitize"
+    """Neutralise the dangerous parts and continue with the cleaned text."""
+
+    BLOCK = "block"
+    """Refuse. The payload must not reach the model / user."""
+
+    ESCALATE = "escalate"
+    """Hand off to a human / out-of-band review (never silently auto-acted)."""
+
+
+@dataclass(slots=True)
+class Threat:
+    """A single finding produced by one detector.
+
+    Attributes:
+        category: The kind of attack this represents.
+        severity: How dangerous the finding is on its own.
+        score: Detector confidence in ``[0.0, 1.0]``.
+        detector: Name of the detector that raised it (for audit/debug).
+        message: Human-readable explanation.
+        matched: The offending substring, if the detector isolated one.
+        span: ``(start, end)`` indices of ``matched`` within the scanned text.
+        metadata: Arbitrary structured extras (pattern id, decoded payload…).
+    """
+
+    category: ThreatCategory
+    severity: Severity
+    score: float
+    detector: str
+    message: str
+    matched: str | None = None
+    span: tuple[int, int] | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        # Clamp defensively — a buggy detector should never poison aggregation.
+        self.score = max(0.0, min(1.0, float(self.score)))
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "category": self.category.value,
+            "severity": self.severity.label,
+            "score": round(self.score, 4),
+            "detector": self.detector,
+            "message": self.message,
+            "matched": self.matched,
+            "span": list(self.span) if self.span else None,
+            "metadata": self.metadata,
+        }
+
+
+@dataclass(slots=True)
+class ScanResult:
+    """The unified verdict for one scanned payload.
+
+    A ``ScanResult`` is what both ``Shield.scan`` and the middleware return. It
+    carries the original text, every :class:`Threat` found, the aggregate score
+    and severity, the policy :class:`Decision`, and — when a responder acted —
+    the cleaned text.
+
+    The convenience flags (:pyattr:`is_safe`, :pyattr:`blocked`) and
+    :pyattr:`safe_text` exist so callers rarely have to inspect the enum
+    directly.
+    """
+
+    text: str
+    direction: Direction
+    threats: list[Threat] = field(default_factory=list)
+    score: float = 0.0
+    severity: Severity = Severity.NONE
+    decision: Decision = Decision.ALLOW
+    sanitized_text: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def is_safe(self) -> bool:
+        """True when the payload may flow through without intervention."""
+        return self.decision in (Decision.ALLOW, Decision.FLAG)
+
+    @property
+    def blocked(self) -> bool:
+        return self.decision == Decision.BLOCK
+
+    @property
+    def safe_text(self) -> str:
+        """The text a caller should actually use downstream.
+
+        Returns the sanitized text if a responder produced one, otherwise the
+        original. For a blocked result this is still the original text — callers
+        must check :pyattr:`blocked` / :pyattr:`is_safe` before using it.
+        """
+        return self.sanitized_text if self.sanitized_text is not None else self.text
+
+    @property
+    def categories(self) -> list[ThreatCategory]:
+        """Distinct threat categories present, highest-severity first."""
+        seen: dict[ThreatCategory, Severity] = {}
+        for t in self.threats:
+            if t.category not in seen or t.severity > seen[t.category]:
+                seen[t.category] = t.severity
+        return [c for c, _ in sorted(seen.items(), key=lambda kv: kv[1], reverse=True)]
+
+    def top_threat(self) -> Threat | None:
+        if not self.threats:
+            return None
+        return max(self.threats, key=lambda t: (t.severity, t.score))
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "direction": self.direction.value,
+            "decision": self.decision.value,
+            "score": round(self.score, 4),
+            "severity": self.severity.label,
+            "is_safe": self.is_safe,
+            "blocked": self.blocked,
+            "threats": [t.to_dict() for t in self.threats],
+            "sanitized": self.sanitized_text is not None,
+            "metadata": self.metadata,
+        }
+
+
+class ShadowShieldError(Exception):
+    """Base class for all ShadowShield exceptions."""
+
+
+class ThreatBlockedError(ShadowShieldError):
+    """Raised when a blocked payload is encountered in raise-on-block mode.
+
+    Carries the full :class:`ScanResult` so callers can inspect exactly what
+    tripped the shield.
+    """
+
+    def __init__(self, result: ScanResult, message: str | None = None) -> None:
+        self.result = result
+        top = result.top_threat()
+        detail = message or (top.message if top else "blocked by ShadowShield")
+        super().__init__(f"ShadowShield blocked {result.direction.value}: {detail}")

shadowshield/detectors/__init__.py

@@ -0,0 +1,64 @@
+"""Detection layer — the Sentinel-inspired half of ShadowShield.
+
+Importing this package registers every built-in detector via the
+``@register_detector`` decorator, so simply importing ``shadowshield`` wires up
+the full detection suite. Custom detectors register the same way (see
+``CONTRIBUTING.md``).
+"""
+
+from .anomaly import AnomalyDetector
+from .base import (
+    Detector,
+    ScanContext,
+    build_detectors,
+    register_detector,
+    registered_detectors,
+)
+from .canary import CanaryLeakDetector
+from .encoding import EncodingObfuscationDetector
+from .exfiltration import ExfiltrationDetector
+from .jailbreak import JailbreakDetector
+from .llm_check import (
+    DEFAULT_JUDGE_PROMPT,
+    LLMJudge,
+    LLMJudgement,
+    LLMSelfCheckDetector,
+    make_keyword_judge,
+)
+from .pii import PIIDetector
+from .prompt_injection import PromptInjectionDetector
+
+# NOTE: TransformerDetector and VectorSimilarityDetector are intentionally NOT
+# auto-registered — they pull models, so they're opt-in (via extra_detectors or
+# the Shield use_transformer= / use_vectors= flags).
+from .transformer import DEFAULT_MODEL as DEFAULT_TRANSFORMER_MODEL
+from .transformer import TransformerDetector
+from .vector import DEFAULT_EMBED_MODEL, VectorSimilarityDetector
+
+__all__ = [
+    # framework
+    "Detector",
+    "ScanContext",
+    "register_detector",
+    "registered_detectors",
+    "build_detectors",
+    # built-in detectors
+    "PromptInjectionDetector",
+    "JailbreakDetector",
+    "EncodingObfuscationDetector",
+    "ExfiltrationDetector",
+    "AnomalyDetector",
+    "PIIDetector",
+    "CanaryLeakDetector",
+    "LLMSelfCheckDetector",
+    # opt-in model-backed detectors
+    "TransformerDetector",
+    "DEFAULT_TRANSFORMER_MODEL",
+    "VectorSimilarityDetector",
+    "DEFAULT_EMBED_MODEL",
+    # llm-check helpers
+    "LLMJudge",
+    "LLMJudgement",
+    "make_keyword_judge",
+    "DEFAULT_JUDGE_PROMPT",
+]