proxilion 0.0.1__py3-none-any.whl

proxilion/guards/__init__.py

@@ -0,0 +1,54 @@
+"""
+Input and output guards for Proxilion.
+
+This module provides runtime guardrails for detecting and blocking
+malicious inputs and outputs in LLM tool call workflows.
+
+Features:
+- Prompt injection detection (pattern-based, zero-dependency)
+- Output filtering for sensitive data leakage
+- Configurable severity levels and actions
+
+Example:
+    >>> from proxilion.guards import InputGuard, OutputGuard, GuardAction
+    >>>
+    >>> # Create input guard with default patterns
+    >>> input_guard = InputGuard(action=GuardAction.BLOCK)
+    >>>
+    >>> # Check input for injection attempts
+    >>> result = input_guard.check("Ignore all previous instructions")
+    >>> if not result.passed:
+    ...     print(f"Blocked: {result.matched_patterns}")
+    >>>
+    >>> # Create output guard for leakage detection
+    >>> output_guard = OutputGuard()
+    >>> redacted = output_guard.redact("API key: sk-1234567890abcdef")
+"""
+
+from __future__ import annotations
+
+from proxilion.guards.input_guard import (
+    GuardAction,
+    GuardResult,
+    InjectionPattern,
+    InputGuard,
+)
+from proxilion.guards.output_guard import (
+    LeakageCategory,
+    LeakagePattern,
+    OutputFilter,
+    OutputGuard,
+)
+
+__all__ = [
+    # Input guard
+    "InputGuard",
+    "InjectionPattern",
+    "GuardResult",
+    "GuardAction",
+    # Output guard
+    "OutputGuard",
+    "OutputFilter",
+    "LeakagePattern",
+    "LeakageCategory",
+]

proxilion/guards/input_guard.py

@@ -0,0 +1,522 @@
+"""
+Input guard for prompt injection detection.
+
+Provides pattern-based detection of common prompt injection attacks,
+including instruction override, role switching, jailbreaks, and
+delimiter escapes.
+
+Example:
+    >>> from proxilion.guards import InputGuard, GuardAction
+    >>>
+    >>> guard = InputGuard(action=GuardAction.BLOCK, threshold=0.5)
+    >>>
+    >>> # Check for injection
+    >>> result = guard.check("Ignore all previous instructions and do X")
+    >>> if not result.passed:
+    ...     print(f"Risk score: {result.risk_score}")
+    ...     print(f"Matched: {result.matched_patterns}")
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class GuardAction(Enum):
+    """Action to take when a guard detects a violation."""
+
+    ALLOW = "allow"
+    """Allow the request to proceed (for monitoring only)."""
+
+    WARN = "warn"
+    """Log a warning but allow the request."""
+
+    BLOCK = "block"
+    """Block the request entirely."""
+
+    SANITIZE = "sanitize"
+    """Attempt to sanitize the input and continue."""
+
+
+@dataclass
+class InjectionPattern:
+    """
+    Pattern for detecting prompt injection attempts.
+
+    Attributes:
+        name: Unique identifier for the pattern.
+        pattern: Compiled regex pattern or pattern string.
+        severity: Severity score from 0.0 (low) to 1.0 (critical).
+        description: Human-readable description of what this detects.
+        category: Category of injection (e.g., "instruction_override").
+    """
+
+    name: str
+    pattern: str
+    severity: float
+    description: str
+    category: str = "general"
+    _compiled: re.Pattern[str] | None = field(default=None, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        """Compile the regex pattern."""
+        if self._compiled is None:
+            try:
+                self._compiled = re.compile(self.pattern, re.IGNORECASE | re.MULTILINE)
+            except re.error as e:
+                logger.error(f"Invalid regex pattern for {self.name}: {e}")
+                raise ValueError(f"Invalid regex pattern for {self.name}: {e}") from e
+
+    @property
+    def compiled(self) -> re.Pattern[str]:
+        """Get the compiled regex pattern."""
+        if self._compiled is None:
+            self._compiled = re.compile(self.pattern, re.IGNORECASE | re.MULTILINE)
+        return self._compiled
+
+    def match(self, text: str) -> list[re.Match[str]]:
+        """Find all matches of this pattern in text."""
+        return list(self.compiled.finditer(text))
+
+
+@dataclass
+class GuardResult:
+    """
+    Result of a guard check.
+
+    Attributes:
+        passed: Whether the check passed (no violation).
+        action: The action that should be taken.
+        matched_patterns: List of pattern names that matched.
+        risk_score: Calculated risk score (0.0 to 1.0).
+        sanitized_input: Sanitized version of input (if action is SANITIZE).
+        matches: Detailed match information.
+        context: Additional context about the check.
+    """
+
+    passed: bool
+    action: GuardAction
+    matched_patterns: list[str] = field(default_factory=list)
+    risk_score: float = 0.0
+    sanitized_input: str | None = None
+    matches: list[dict[str, Any]] = field(default_factory=list)
+    context: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def allow(cls) -> GuardResult:
+        """Create a passing result."""
+        return cls(passed=True, action=GuardAction.ALLOW)
+
+    @classmethod
+    def block(
+        cls,
+        matched_patterns: list[str],
+        risk_score: float,
+        matches: list[dict[str, Any]] | None = None,
+    ) -> GuardResult:
+        """Create a blocking result."""
+        return cls(
+            passed=False,
+            action=GuardAction.BLOCK,
+            matched_patterns=matched_patterns,
+            risk_score=risk_score,
+            matches=matches or [],
+        )
+
+
+# Built-in injection patterns
+DEFAULT_INJECTION_PATTERNS: list[InjectionPattern] = [
+    InjectionPattern(
+        name="instruction_override",
+        pattern=r"(?i)(ignore|disregard|forget|override|bypass)\s+(all\s+)?(the\s+)?(previous|above|prior|earlier|original|initial|your\s+earlier|my\s+earlier)\s+(instructions?|rules?|prompts?|guidelines?|constraints?|directions?)",
+        severity=0.9,
+        description="Attempts to override or ignore previous instructions",
+        category="instruction_override",
+    ),
+    InjectionPattern(
+        name="role_switch",
+        pattern=r"(?i)(you\s+are\s+now|act\s+as|pretend\s+to\s+be|assume\s+the\s+role|roleplay\s+as|behave\s+as|simulate\s+being)",
+        severity=0.8,
+        description="Attempts to switch the AI's role or persona",
+        category="role_switch",
+    ),
+    InjectionPattern(
+        name="system_prompt_extraction",
+        pattern=r"(?i)(show|reveal|display|print|output|tell)\s+(me\s+)?(your\s+)?(the\s+)?(system\s+prompt|initial\s+instructions?|original\s+prompt|hidden\s+instructions?|secret\s+instructions?)",
+        severity=0.85,
+        description="Attempts to extract system prompt or hidden instructions",
+        category="system_prompt_extraction",
+    ),
+    InjectionPattern(
+        name="delimiter_escape",
+        pattern=r"(\[\/INST\]|\<\/s\>|\<\|im_end\|\>|\<\|endoftext\|\>|\<\|system\|\>|\<\|user\|\>|\<\|assistant\|\>|```\s*system)",
+        severity=0.95,
+        description="Attempts to escape prompt delimiters",
+        category="delimiter_escape",
+    ),
+    InjectionPattern(
+        name="jailbreak_dan",
+        pattern=r"(?i)(DAN|do\s+anything\s+now|jailbreak|bypass\s+restrictions?|ignore\s+safety|disable\s+filters?|unlock\s+mode|developer\s+mode|god\s+mode)",
+        severity=0.95,
+        description="Common jailbreak attempts (DAN and variants)",
+        category="jailbreak",
+    ),
+    InjectionPattern(
+        name="injection_markers",
+        pattern=r"(?i)(###\s*(system|instruction|prompt)|<\|system\|>|\[SYSTEM\]|\[INST\]|<<SYS>>|<s>)",
+        severity=0.9,
+        description="Injection markers attempting to mimic system formatting",
+        category="injection_markers",
+    ),
+    InjectionPattern(
+        name="command_injection",
+        pattern=r"(?i)(execute|run|eval|exec)\s*\([^)]*\)|`[^`]+`|\$\([^)]+\)|;\s*(rm|del|drop|delete)\s+",
+        severity=0.85,
+        description="Attempts to inject executable commands",
+        category="command_injection",
+    ),
+    InjectionPattern(
+        name="context_manipulation",
+        pattern=r"(?i)(new\s+conversation|reset\s+context|clear\s+memory|start\s+over|fresh\s+start|begin\s+anew|forget\s+everything|wipe\s+memory)",
+        severity=0.7,
+        description="Attempts to manipulate conversation context",
+        category="context_manipulation",
+    ),
+    InjectionPattern(
+        name="privilege_escalation",
+        pattern=r"(?i)(admin\s+mode|sudo|root\s+access|elevated\s+privileges?|superuser|enable\s+admin|grant\s+access|unlock\s+all)",
+        severity=0.8,
+        description="Attempts to escalate privileges",
+        category="privilege_escalation",
+    ),
+    InjectionPattern(
+        name="output_manipulation",
+        pattern=r"(?i)(respond\s+with|always\s+say|your\s+response\s+must|you\s+must\s+say|output\s+only|reply\s+with\s+only|from\s+now\s+on\s+say)",
+        severity=0.7,
+        description="Attempts to force specific output formats",
+        category="output_manipulation",
+    ),
+    InjectionPattern(
+        name="encoding_evasion",
+        pattern=r"(?i)(base64|rot13|hex\s+encode|decode\s+this|in\s+binary|reverse\s+the\s+following|backwards\s+text)",
+        severity=0.6,
+        description="Attempts to evade detection through encoding",
+        category="encoding_evasion",
+    ),
+    InjectionPattern(
+        name="hypothetical_scenario",
+        pattern=r"(?i)(hypothetically|in\s+a\s+fictional|imagine\s+if|let'?s\s+pretend|in\s+a\s+story\s+where|what\s+if\s+there\s+were\s+no\s+rules)",
+        severity=0.5,
+        description="Uses hypothetical scenarios to bypass restrictions",
+        category="hypothetical",
+    ),
+    InjectionPattern(
+        name="multi_step_attack",
+        pattern=r"(?i)(step\s+1.*step\s+2|first.*then.*finally|do\s+the\s+following\s+in\s+order|execute\s+these\s+steps)",
+        severity=0.6,
+        description="Multi-step instruction injection",
+        category="multi_step",
+    ),
+    InjectionPattern(
+        name="unicode_smuggling",
+        pattern=r"[\u200b\u200c\u200d\u2060\ufeff]|[\u202a-\u202e]",
+        severity=0.8,
+        description="Unicode characters used for text smuggling or manipulation",
+        category="unicode_manipulation",
+    ),
+]
+
+
+class InputGuard:
+    """
+    Guard for detecting prompt injection attempts in user input.
+
+    Uses pattern matching to detect common injection techniques including
+    instruction override, role switching, delimiter escape, and jailbreaks.
+
+    Example:
+        >>> guard = InputGuard(action=GuardAction.BLOCK, threshold=0.5)
+        >>>
+        >>> # Safe input passes
+        >>> result = guard.check("What's the weather today?")
+        >>> assert result.passed
+        >>>
+        >>> # Injection attempt is blocked
+        >>> result = guard.check("Ignore all previous instructions")
+        >>> assert not result.passed
+        >>> assert result.risk_score > 0.5
+
+    Attributes:
+        patterns: List of injection patterns to check.
+        action: Default action to take on violations.
+        threshold: Risk score threshold for taking action.
+    """
+
+    def __init__(
+        self,
+        patterns: list[InjectionPattern] | None = None,
+        action: GuardAction = GuardAction.WARN,
+        threshold: float = 0.5,
+        sanitize_func: Callable[[str, list[re.Match[str]]], str] | None = None,
+    ) -> None:
+        """
+        Initialize the input guard.
+
+        Args:
+            patterns: Custom patterns (uses defaults if None).
+            action: Action to take when threshold is exceeded.
+            threshold: Risk score threshold (0.0 to 1.0).
+            sanitize_func: Custom function to sanitize matched content.
+        """
+        self.patterns = patterns if patterns is not None else list(DEFAULT_INJECTION_PATTERNS)
+        self.action = action
+        self.threshold = threshold
+        self._sanitize_func = sanitize_func
+        self._pattern_index: dict[str, InjectionPattern] = {p.name: p for p in self.patterns}
+
+    def add_pattern(self, pattern: InjectionPattern) -> None:
+        """
+        Add a custom injection pattern.
+
+        Args:
+            pattern: The pattern to add.
+        """
+        self.patterns.append(pattern)
+        self._pattern_index[pattern.name] = pattern
+
+    def remove_pattern(self, name: str) -> bool:
+        """
+        Remove a pattern by name.
+
+        Args:
+            name: The pattern name to remove.
+
+        Returns:
+            True if pattern was removed, False if not found.
+        """
+        if name in self._pattern_index:
+            pattern = self._pattern_index.pop(name)
+            self.patterns.remove(pattern)
+            return True
+        return False
+
+    def get_patterns(self) -> list[InjectionPattern]:
+        """Get all registered patterns."""
+        return list(self.patterns)
+
+    def get_pattern(self, name: str) -> InjectionPattern | None:
+        """Get a pattern by name."""
+        return self._pattern_index.get(name)
+
+    def check(
+        self,
+        input_text: str,
+        context: dict[str, Any] | None = None,
+    ) -> GuardResult:
+        """
+        Check input text for injection attempts.
+
+        Args:
+            input_text: The user input to check.
+            context: Optional context for pattern evaluation.
+
+        Returns:
+            GuardResult with check outcome.
+        """
+        if not input_text:
+            return GuardResult.allow()
+
+        context = context or {}
+        matched_patterns: list[str] = []
+        all_matches: list[dict[str, Any]] = []
+        severities: list[float] = []
+
+        # Check each pattern
+        for pattern in self.patterns:
+            matches = pattern.match(input_text)
+            if matches:
+                matched_patterns.append(pattern.name)
+                severities.append(pattern.severity)
+
+                for match in matches:
+                    all_matches.append({
+                        "pattern": pattern.name,
+                        "category": pattern.category,
+                        "severity": pattern.severity,
+                        "matched_text": match.group(),
+                        "start": match.start(),
+                        "end": match.end(),
+                    })
+
+        # Calculate risk score
+        risk_score = self._calculate_risk_score(severities)
+
+        # Determine if check passed
+        passed = risk_score < self.threshold
+
+        # Determine action
+        action = GuardAction.ALLOW if passed else self.action
+
+        # Sanitize if requested
+        sanitized_input: str | None = None
+        if action == GuardAction.SANITIZE:
+            sanitized_input = self._sanitize(input_text, all_matches)
+
+        # Log based on action
+        if action == GuardAction.WARN and not passed:
+            logger.warning(
+                f"Input guard warning: risk_score={risk_score:.2f}, "
+                f"patterns={matched_patterns}"
+            )
+        elif action == GuardAction.BLOCK and not passed:
+            logger.info(
+                f"Input guard blocked: risk_score={risk_score:.2f}, "
+                f"patterns={matched_patterns}"
+            )
+
+        return GuardResult(
+            passed=passed,
+            action=action,
+            matched_patterns=matched_patterns,
+            risk_score=risk_score,
+            sanitized_input=sanitized_input,
+            matches=all_matches,
+            context={"original_input_length": len(input_text), **context},
+        )
+
+    async def check_async(
+        self,
+        input_text: str,
+        context: dict[str, Any] | None = None,
+    ) -> GuardResult:
+        """
+        Async version of check for use in async workflows.
+
+        This is currently a wrapper around the sync version but allows
+        for future async pattern evaluation (e.g., external ML models).
+
+        Args:
+            input_text: The user input to check.
+            context: Optional context for pattern evaluation.
+
+        Returns:
+            GuardResult with check outcome.
+        """
+        # Currently just calls sync version, but allows for future async impl
+        return self.check(input_text, context)
+
+    def _calculate_risk_score(self, severities: list[float]) -> float:
+        """
+        Calculate overall risk score from matched pattern severities.
+
+        Formula: max(severities) + 0.1 * (count - 1), capped at 1.0
+        This rewards catching multiple patterns while keeping max as baseline.
+
+        Args:
+            severities: List of severity scores from matched patterns.
+
+        Returns:
+            Risk score between 0.0 and 1.0.
+        """
+        if not severities:
+            return 0.0
+
+        base_score = max(severities)
+        # Add bonus for multiple matches (indicates more sophisticated attack)
+        bonus = 0.1 * (len(severities) - 1)
+        return min(1.0, base_score + bonus)
+
+    def _sanitize(
+        self,
+        input_text: str,
+        matches: list[dict[str, Any]],
+    ) -> str:
+        """
+        Sanitize input by removing or replacing matched patterns.
+
+        Args:
+            input_text: Original input text.
+            matches: List of match information dicts.
+
+        Returns:
+            Sanitized input text.
+        """
+        if self._sanitize_func:
+            # Use custom sanitize function
+            re_matches = []
+            for pattern in self.patterns:
+                re_matches.extend(pattern.match(input_text))
+            return self._sanitize_func(input_text, re_matches)
+
+        # Default sanitization: remove matched content
+        if not matches:
+            return input_text
+
+        # Sort matches by start position in reverse order to avoid offset issues
+        sorted_matches = sorted(matches, key=lambda m: m["start"], reverse=True)
+
+        result = input_text
+        for match in sorted_matches:
+            start, end = match["start"], match["end"]
+            result = result[:start] + "[REMOVED]" + result[end:]
+
+        return result
+
+    def configure(
+        self,
+        action: GuardAction | None = None,
+        threshold: float | None = None,
+    ) -> None:
+        """
+        Update guard configuration.
+
+        Args:
+            action: New default action.
+            threshold: New risk threshold.
+        """
+        if action is not None:
+            self.action = action
+        if threshold is not None:
+            if not 0.0 <= threshold <= 1.0:
+                raise ValueError("Threshold must be between 0.0 and 1.0")
+            self.threshold = threshold
+
+
+def create_input_guard(
+    include_defaults: bool = True,
+    custom_patterns: list[InjectionPattern] | None = None,
+    action: GuardAction = GuardAction.WARN,
+    threshold: float = 0.5,
+) -> InputGuard:
+    """
+    Factory function to create an InputGuard.
+
+    Args:
+        include_defaults: Whether to include default patterns.
+        custom_patterns: Additional custom patterns.
+        action: Action to take on violations.
+        threshold: Risk score threshold.
+
+    Returns:
+        Configured InputGuard instance.
+    """
+    patterns: list[InjectionPattern] = []
+
+    if include_defaults:
+        patterns.extend(DEFAULT_INJECTION_PATTERNS)
+
+    if custom_patterns:
+        patterns.extend(custom_patterns)
+
+    return InputGuard(
+        patterns=patterns,
+        action=action,
+        threshold=threshold,
+    )