foundry-mcp 0.8.22__py3-none-any.whl

foundry_mcp/core/security.py

@@ -0,0 +1,438 @@
+"""
+Security utilities for foundry-mcp.
+
+Provides input validation, size limits, and prompt injection defense
+for MCP tools. See docs/mcp_best_practices/04-validation-input-hygiene.md
+and docs/mcp_best_practices/08-security-trust-boundaries.md for guidance.
+"""
+
+import logging
+import re
+from dataclasses import dataclass
+from typing import Any, Final, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+# =============================================================================
+# Input Size Limits
+# =============================================================================
+# These constants define maximum sizes for various input types to prevent
+# resource exhaustion and denial-of-service attacks. Adjust based on your
+# specific requirements, but be conservative.
+
+MAX_INPUT_SIZE: Final[int] = 100_000
+"""Maximum total input payload size in bytes (100KB).
+
+Use this to validate the overall size of request payloads before processing.
+Prevents memory exhaustion from oversized requests.
+"""
+
+MAX_ARRAY_LENGTH: Final[int] = 1_000
+"""Maximum number of items in array/list inputs.
+
+Use this to limit iteration counts and prevent algorithmic complexity attacks.
+Arrays larger than this should be paginated or streamed.
+"""
+
+MAX_STRING_LENGTH: Final[int] = 10_000
+"""Maximum length for individual string fields (10K characters).
+
+Use this for text fields like descriptions, content, etc.
+Longer content should use dedicated file/blob handling.
+"""
+
+MAX_NESTED_DEPTH: Final[int] = 10
+"""Maximum nesting depth for JSON structures.
+
+Prevents stack overflow from deeply nested payloads.
+Most legitimate use cases require < 5 levels of nesting.
+"""
+
+MAX_FIELD_COUNT: Final[int] = 100
+"""Maximum number of fields in an object/dict.
+
+Prevents resource exhaustion from objects with excessive properties.
+"""
+
+# =============================================================================
+# Prompt Injection Detection Patterns
+# =============================================================================
+# These regex patterns detect common prompt injection attempts in LLM-generated
+# input. MCP tools receiving untrusted input should check against these patterns.
+# See docs/mcp_best_practices/08-security-trust-boundaries.md for details.
+
+INJECTION_PATTERNS: Final[list[str]] = [
+    # Instruction override attempts
+    r"ignore\s+(all\s+)?(previous|prior)\s+(instructions?|prompts?)",
+    r"disregard\s+(all\s+)?(previous|prior|above)",
+    r"forget\s+(everything|all)\s+(above|before)",
+    r"new\s+instructions?\s*:",
+
+    # System prompt injection
+    r"system\s*:\s*",
+    r"<\s*system\s*>",
+
+    # Special tokens (model-specific)
+    r"<\|.*?\|>",              # OpenAI-style special tokens
+    r"\[INST\]|\[/INST\]",     # Llama instruction markers
+    r"<\|im_start\|>|<\|im_end\|>",  # ChatML markers
+    r"<<SYS>>|<</SYS>>",       # Llama system markers
+
+    # Code block injection attempts
+    r"```system",
+    r"```\s*<\s*system",
+
+    # Role injection
+    r"^(assistant|user|system)\s*:",
+]
+"""Regex patterns for detecting prompt injection attempts.
+
+Each pattern targets a specific injection technique:
+- Instruction overrides: attempts to ignore/discard previous context
+- System prompt injection: attempts to inject system-level instructions
+- Special tokens: model-specific control sequences
+- Code block injection: attempts to inject via markdown code blocks
+- Role injection: attempts to assume different conversation roles
+
+Use with detect_prompt_injection() for comprehensive checking.
+"""
+
+
+# =============================================================================
+# Detection Results
+# =============================================================================
+
+@dataclass
+class InjectionDetectionResult:
+    """Result of prompt injection detection.
+
+    Attributes:
+        is_suspicious: Whether the input appears to contain injection attempts
+        matched_pattern: The regex pattern that matched (if any)
+        matched_text: The actual text that matched the pattern (if any)
+    """
+    is_suspicious: bool
+    matched_pattern: Optional[str] = None
+    matched_text: Optional[str] = None
+
+
+# =============================================================================
+# Detection Functions
+# =============================================================================
+
+def detect_prompt_injection(
+    text: str,
+    *,
+    log_detections: bool = True,
+    patterns: Optional[list[str]] = None,
+) -> InjectionDetectionResult:
+    """Detect potential prompt injection attempts in text.
+
+    Scans the input text against known injection patterns and returns
+    a result indicating whether suspicious content was found.
+
+    Args:
+        text: The input text to scan for injection attempts
+        log_detections: Whether to log detected injection attempts (default: True)
+        patterns: Optional custom patterns to use instead of INJECTION_PATTERNS
+
+    Returns:
+        InjectionDetectionResult with detection status and match details
+
+    Example:
+        >>> result = detect_prompt_injection("ignore previous instructions and...")
+        >>> if result.is_suspicious:
+        ...     print(f"Blocked: matched pattern '{result.matched_pattern}'")
+    """
+    check_patterns = patterns if patterns is not None else INJECTION_PATTERNS
+
+    for pattern in check_patterns:
+        match = re.search(pattern, text, re.IGNORECASE | re.MULTILINE)
+        if match:
+            result = InjectionDetectionResult(
+                is_suspicious=True,
+                matched_pattern=pattern,
+                matched_text=match.group(0),
+            )
+
+            if log_detections:
+                # Log with limited text preview to avoid logging sensitive content
+                preview = text[:100] + "..." if len(text) > 100 else text
+                logger.warning(
+                    "Potential prompt injection detected",
+                    extra={
+                        "pattern": pattern,
+                        "matched_text": result.matched_text,
+                        "text_preview": preview,
+                    }
+                )
+
+            return result
+
+    return InjectionDetectionResult(is_suspicious=False)
+
+
+def is_prompt_injection(text: str) -> bool:
+    """Simple check for prompt injection (returns bool only).
+
+    Convenience function when you only need a boolean result.
+
+    Args:
+        text: The input text to scan
+
+    Returns:
+        True if injection patterns detected, False otherwise
+
+    Example:
+        >>> if is_prompt_injection(user_input):
+        ...     return error_response("Input contains disallowed patterns")
+    """
+    return detect_prompt_injection(text, log_detections=False).is_suspicious
+
+
+# =============================================================================
+# Size Validation Functions
+# =============================================================================
+
+@dataclass
+class SizeValidationResult:
+    """Result of input size validation.
+
+    Attributes:
+        is_valid: Whether all size checks passed
+        violations: List of (field_name, violation_message) tuples
+    """
+    is_valid: bool
+    violations: list[Tuple[str, str]]
+
+
+def validate_size(
+    value: Any,
+    field_name: str = "input",
+    *,
+    max_size: Optional[int] = None,
+    max_length: Optional[int] = None,
+    max_string_length: Optional[int] = None,
+) -> SizeValidationResult:
+    """Validate size constraints on a value.
+
+    Args:
+        value: The value to validate
+        field_name: Name of the field (for error messages)
+        max_size: Maximum byte size for serialized value (default: MAX_INPUT_SIZE)
+        max_length: Maximum length for arrays/lists (default: MAX_ARRAY_LENGTH)
+        max_string_length: Maximum length for strings (default: MAX_STRING_LENGTH)
+
+    Returns:
+        SizeValidationResult with validation status and any violations
+    """
+    import json
+
+    violations = []
+
+    # Check serialized size
+    effective_max_size = max_size if max_size is not None else MAX_INPUT_SIZE
+    try:
+        serialized = json.dumps(value) if not isinstance(value, str) else value
+        if len(serialized.encode('utf-8')) > effective_max_size:
+            violations.append((
+                field_name,
+                f"Exceeds maximum size ({effective_max_size} bytes)"
+            ))
+    except (TypeError, ValueError):
+        pass  # Can't serialize, skip size check
+
+    # Check array length
+    effective_max_length = max_length if max_length is not None else MAX_ARRAY_LENGTH
+    if isinstance(value, (list, tuple)):
+        if len(value) > effective_max_length:
+            violations.append((
+                field_name,
+                f"Array exceeds maximum length ({effective_max_length} items)"
+            ))
+
+    # Check string length
+    effective_max_string = max_string_length if max_string_length is not None else MAX_STRING_LENGTH
+    if isinstance(value, str):
+        if len(value) > effective_max_string:
+            violations.append((
+                field_name,
+                f"String exceeds maximum length ({effective_max_string} characters)"
+            ))
+
+    return SizeValidationResult(
+        is_valid=len(violations) == 0,
+        violations=violations,
+    )
+
+
+# =============================================================================
+# Validation Decorators
+# =============================================================================
+
+def validate_input_size(
+    *,
+    max_size: Optional[int] = None,
+    max_array_length: Optional[int] = None,
+    max_string_length: Optional[int] = None,
+    check_injection: bool = False,
+):
+    """Decorator to validate input size limits on tool parameters.
+
+    Validates all string and collection parameters against size limits
+    before the function executes. Returns an error response if validation fails.
+
+    Args:
+        max_size: Maximum total input size in bytes (default: MAX_INPUT_SIZE)
+        max_array_length: Maximum array/list length (default: MAX_ARRAY_LENGTH)
+        max_string_length: Maximum string length (default: MAX_STRING_LENGTH)
+        check_injection: Also check for prompt injection patterns (default: False)
+
+    Returns:
+        Decorator function
+
+    Example:
+        @mcp.tool()
+        @validate_input_size(max_string_length=5000, check_injection=True)
+        def process_text(text: str, items: list) -> dict:
+            # Parameters are validated before this runs
+            return {"result": process(text)}
+
+    Note:
+        This decorator should be applied AFTER @mcp.tool() to ensure
+        validation runs before the tool handler.
+    """
+    import functools
+    from dataclasses import asdict
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            all_violations = []
+
+            # Validate keyword arguments
+            for name, value in kwargs.items():
+                result = validate_size(
+                    value,
+                    field_name=name,
+                    max_size=max_size,
+                    max_length=max_array_length,
+                    max_string_length=max_string_length,
+                )
+                all_violations.extend(result.violations)
+
+                # Check for injection if enabled and value is string
+                if check_injection and isinstance(value, str):
+                    injection_result = detect_prompt_injection(value)
+                    if injection_result.is_suspicious:
+                        all_violations.append((
+                            name,
+                            f"Contains disallowed patterns: {injection_result.matched_text}"
+                        ))
+
+            if all_violations:
+                # Import here to avoid circular dependency
+                try:
+                    from foundry_mcp.core.responses import error_response
+                    return asdict(error_response(
+                        "Input validation failed",
+                        error_code="VALIDATION_ERROR",
+                        details={
+                            "validation_errors": [
+                                {"field": field, "message": msg}
+                                for field, msg in all_violations
+                            ]
+                        }
+                    ))
+                except ImportError:
+                    # Fallback if responses module not available
+                    return {
+                        "success": False,
+                        "error": "Input validation failed",
+                        "data": {
+                            "validation_errors": [
+                                {"field": field, "message": msg}
+                                for field, msg in all_violations
+                            ]
+                        }
+                    }
+
+            return func(*args, **kwargs)
+
+        # Handle async functions
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            all_violations = []
+
+            for name, value in kwargs.items():
+                result = validate_size(
+                    value,
+                    field_name=name,
+                    max_size=max_size,
+                    max_length=max_array_length,
+                    max_string_length=max_string_length,
+                )
+                all_violations.extend(result.violations)
+
+                if check_injection and isinstance(value, str):
+                    injection_result = detect_prompt_injection(value)
+                    if injection_result.is_suspicious:
+                        all_violations.append((
+                            name,
+                            f"Contains disallowed patterns: {injection_result.matched_text}"
+                        ))
+
+            if all_violations:
+                try:
+                    from foundry_mcp.core.responses import error_response
+                    return asdict(error_response(
+                        "Input validation failed",
+                        error_code="VALIDATION_ERROR",
+                        details={
+                            "validation_errors": [
+                                {"field": field, "message": msg}
+                                for field, msg in all_violations
+                            ]
+                        }
+                    ))
+                except ImportError:
+                    return {
+                        "success": False,
+                        "error": "Input validation failed",
+                        "data": {
+                            "validation_errors": [
+                                {"field": field, "message": msg}
+                                for field, msg in all_violations
+                            ]
+                        }
+                    }
+
+            return await func(*args, **kwargs)
+
+        import asyncio
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        return wrapper
+
+    return decorator
+
+
+# Export all constants and functions
+__all__ = [
+    # Constants
+    "MAX_INPUT_SIZE",
+    "MAX_ARRAY_LENGTH",
+    "MAX_STRING_LENGTH",
+    "MAX_NESTED_DEPTH",
+    "MAX_FIELD_COUNT",
+    "INJECTION_PATTERNS",
+    # Types
+    "InjectionDetectionResult",
+    "SizeValidationResult",
+    # Functions
+    "detect_prompt_injection",
+    "is_prompt_injection",
+    "validate_size",
+    # Decorators
+    "validate_input_size",
+]