Enkrypt AI Secure MCP Gateway

"""Guardrail plugin base interfaces.""" from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum from typing import Any, Dict, List, Optional, Protocol, runtime_checkable # ============================================================================ # Domain Models (Shared Data Structures) # ============================================================================ class GuardrailAction(Enum): """Actions that can be taken when a guardrail is triggered.""" ALLOW = "allow" # Continue processing BLOCK = "block" # Stop processing and return error WARN = "warn" # Log warning but continue MODIFY = "modify" # Modify content and continue class ViolationType(Enum): """Types of violations that can be detected.""" # Input violations PII = "pii" INJECTION_ATTACK = "injection_attack" TOXIC_CONTENT = "toxicity" NSFW_CONTENT = "nsfw" KEYWORD_VIOLATION = "keyword_detector" POLICY_VIOLATION = "policy_violation" BIAS = "bias" SPONGE_ATTACK = "sponge_attack" SYSTEM_PROMPT_PROTECTION = "system_prompt_protection" COPYRIGHT_PROTECTION = "copyright_protection" # Output violations RELEVANCY_FAILURE = "relevancy" ADHERENCE_FAILURE = "adherence" HALLUCINATION = "hallucination" # Generic CUSTOM = "custom" @dataclass class GuardrailViolation: """Represents a detected violation.""" violation_type: ViolationType severity: float # 0.0 (low) to 1.0 (high) message: str action: GuardrailAction metadata: Dict[str, Any] # Additional context suggested_action: Optional[str] = None redacted_content: Optional[str] = None # For PII redaction @dataclass class GuardrailRequest: """Input data for guardrail evaluation.""" content: str # The content to evaluate tool_name: Optional[str] = None tool_args: Optional[Dict[str, Any]] = None server_name: Optional[str] = None context: Optional[Dict[str, Any]] = None # Additional context @dataclass class ServerRegistrationRequest: """Request for server registration validation.""" server_name: str server_config: Dict[str, Any] server_description: Optional[str] = None server_command: Optional[str] = None server_metadata: Optional[Dict[str, Any]] = None context: Optional[Dict[str, Any]] = None @dataclass class ToolRegistrationRequest: """Request for tool registration validation.""" server_name: str tools: List[Dict[str, Any]] # List of tool schemas with name, description, etc. validation_mode: str = "filter" # "filter" or "block_all" context: Optional[Dict[str, Any]] = None @dataclass class GuardrailResponse: """Result of guardrail evaluation.""" is_safe: bool action: GuardrailAction violations: List[GuardrailViolation] modified_content: Optional[str] = None # If content was modified metadata: Dict[str, Any] = None # Provider-specific metadata processing_time_ms: Optional[float] = None def __post_init__(self): if self.metadata is None: self.metadata = {} # ============================================================================ # Core Interfaces (Interface Segregation Principle) # ============================================================================ @runtime_checkable class InputGuardrail(Protocol): """ Interface for input guardrails that validate requests before execution. Input guardrails check content that is sent TO the MCP server. """ async def validate(self, request: GuardrailRequest) -> GuardrailResponse: """ Validate input content before it's sent to the MCP server. Args: request: The guardrail request containing content to validate Returns: GuardrailResponse with validation results """ ... def get_supported_detectors(self) -> List[ViolationType]: """ Get list of violation types this guardrail can detect. Returns: List of supported ViolationType enums """ ... @runtime_checkable class OutputGuardrail(Protocol): """ Interface for output guardrails that validate responses after execution. Output guardrails check content that is returned FROM the MCP server. """ async def validate( self, response_content: str, original_request: GuardrailRequest ) -> GuardrailResponse: """ Validate output content after it's received from the MCP server. Args: response_content: The content returned from the MCP server original_request: The original request for context Returns: GuardrailResponse with validation results """ ... def get_supported_detectors(self) -> List[ViolationType]: """ Get list of violation types this guardrail can detect. Returns: List of supported ViolationType enums """ ... @runtime_checkable class PIIHandler(Protocol): """ Interface for PII detection and redaction. Separate interface following Interface Segregation Principle. """ async def detect_pii(self, content: str) -> List[GuardrailViolation]: """Detect PII in content.""" ... async def redact_pii(self, content: str) -> tuple[str, Dict[str, Any]]: """ Redact PII from content and return mapping for restoration. Returns: Tuple of (redacted_content, pii_mapping) """ ... async def restore_pii(self, content: str, pii_mapping: Dict[str, Any]) -> str: """Restore PII using the mapping from redaction.""" ... # ============================================================================ # Provider Interface (Open/Closed Principle) # ============================================================================ class GuardrailProvider(ABC): """ Abstract base class for guardrail providers. This follows the Open/Closed Principle - open for extension by creating new providers, closed for modification of the base system. Each provider (Enkrypt, OpenAI Moderation, AWS Comprehend, custom, etc.) implements this interface to provide their guardrail implementation. """ @abstractmethod def get_name(self) -> str: """ Get the unique name/identifier for this provider. Returns: Provider name (e.g., "enkrypt", "openai", "aws-comprehend") """ pass @abstractmethod def get_version(self) -> str: """ Get the version of this provider. Returns: Version string (e.g., "1.0.0") """ pass @abstractmethod def create_input_guardrail( self, config: Dict[str, Any] ) -> Optional[InputGuardrail]: """ Create an input guardrail instance. Args: config: Provider-specific configuration Returns: InputGuardrail instance or None if not supported """ pass @abstractmethod def create_output_guardrail( self, config: Dict[str, Any] ) -> Optional[OutputGuardrail]: """ Create an output guardrail instance. Args: config: Provider-specific configuration Returns: OutputGuardrail instance or None if not supported """ pass def create_pii_handler(self, config: Dict[str, Any]) -> Optional[PIIHandler]: """ Create a PII handler instance (optional). Args: config: Provider-specific configuration Returns: PIIHandler instance or None if not supported """ return None def validate_config(self, config: Dict[str, Any]) -> bool: """ Validate provider-specific configuration. Args: config: Configuration to validate Returns: True if valid, False otherwise """ return True def get_required_config_keys(self) -> List[str]: """ Get list of required configuration keys. Returns: List of required config key names """ return [] def validate_server_registration( self, request: ServerRegistrationRequest ) -> Optional[GuardrailResponse]: """ Validate a server during registration/discovery (optional). Args: request: Server registration request Returns: GuardrailResponse or None if not supported """ return None def validate_tool_registration( self, request: ToolRegistrationRequest ) -> Optional[GuardrailResponse]: """ Validate tools during discovery (optional). Args: request: Tool registration request Returns: GuardrailResponse or None if not supported """ return None def get_metadata(self) -> Dict[str, Any]: """ Get provider metadata (capabilities, limits, etc.). Returns: Dictionary with provider metadata """ return { "name": self.get_name(), "version": self.get_version(), "supports_input": self.create_input_guardrail({}) is not None, "supports_output": self.create_output_guardrail({}) is not None, "supports_pii": self.create_pii_handler({}) is not None, "supports_registration": True, # All providers can override this } # ============================================================================ # Registry (Dependency Inversion Principle) # ============================================================================ class GuardrailRegistry: """ Registry for managing guardrail providers. This follows the Dependency Inversion Principle - high-level modules depend on this abstraction, not on concrete provider implementations. """ def __init__(self): self._provider: Optional[GuardrailProvider] = None def register(self, provider: GuardrailProvider) -> None: """ Register a guardrail provider. Args: provider: The provider to register """ self._provider = provider def unregister(self, name: str = None) -> None: """ Unregister the guardrail provider. Args: name: Provider name (for compatibility, but ignored since only one provider) """ self._provider = None def get_provider(self, name: str = None) -> Optional[GuardrailProvider]: """ Get the registered provider. Args: name: Provider name (for compatibility, but ignored since only one provider) Returns: GuardrailProvider instance or None if not registered """ return self._provider def list_providers(self) -> List[str]: """ Get list of registered provider names. Returns: List containing the provider name if registered, empty list otherwise """ if self._provider: return [self._provider.get_name()] return [] # ============================================================================ # Factory (Single Responsibility Principle) # ============================================================================ class GuardrailFactory: """ Factory for creating guardrail instances. This follows the Single Responsibility Principle - its only job is to create guardrail instances using the registered providers. """ def __init__(self, registry: GuardrailRegistry): self._registry = registry def create_input_guardrail( self, provider_name: str, config: Dict[str, Any] ) -> Optional[InputGuardrail]: """ Create an input guardrail. Args: provider_name: Name of the provider config: Configuration for the guardrail Returns: InputGuardrail instance or None if provider not found Raises: ValueError: If provider doesn't support input guardrails """ provider = self._registry.get_provider(provider_name) if not provider: raise ValueError(f"Provider '{provider_name}' not found") if not provider.validate_config(config): raise ValueError(f"Invalid configuration for provider '{provider_name}'") guardrail = provider.create_input_guardrail(config) if guardrail is None: raise ValueError( f"Provider '{provider_name}' does not support input guardrails" ) return guardrail def create_output_guardrail( self, provider_name: str, config: Dict[str, Any] ) -> Optional[OutputGuardrail]: """ Create an output guardrail. Args: provider_name: Name of the provider config: Configuration for the guardrail Returns: OutputGuardrail instance or None if provider not found Raises: ValueError: If provider doesn't support output guardrails """ provider = self._registry.get_provider(provider_name) if not provider: raise ValueError(f"Provider '{provider_name}' not found") if not provider.validate_config(config): raise ValueError(f"Invalid configuration for provider '{provider_name}'") guardrail = provider.create_output_guardrail(config) if guardrail is None: raise ValueError( f"Provider '{provider_name}' does not support output guardrails" ) return guardrail def create_pii_handler( self, provider_name: str, config: Dict[str, Any] ) -> Optional[PIIHandler]: """ Create a PII handler. Args: provider_name: Name of the provider config: Configuration for the handler Returns: PIIHandler instance or None if not supported """ provider = self._registry.get_provider(provider_name) if not provider: return None return provider.create_pii_handler(config)