secure mcp protocol

src/mcp/server/fastmcp/secure/__init__.py

@@ -0,0 +1,54 @@
+"""
+Secure annotations and decorators for MCP tools, resources, and prompts.
+
+This module provides enhanced security features including:
+- Bidirectional authentication (client ↔ tool)
+- End-to-end encryption
+- Tool attestation and signing
+- Rate limiting and audit logging
+"""
+
+from .annotations import (
+    AuthMethod,
+    SecureAnnotations,
+    SecureToolAnnotations,
+    SecureResourceAnnotations,
+    SecurePromptAnnotations,
+)
+from .tool import SecureTool, secure_tool
+from .resource import SecureResource, secure_resource
+from .prompt import SecurePrompt, secure_prompt
+from .identity import ToolIdentity, ClientIdentity, create_tool_identity
+from .session import SecureSession, SessionManager
+from .utils import SecureAnnotationProcessor, encrypt_data, decrypt_data
+
+__all__ = [
+    # Annotations
+    "AuthMethod",
+    "SecureAnnotations",
+    "SecureToolAnnotations",
+    "SecureResourceAnnotations",
+    "SecurePromptAnnotations",
+
+    # Secure wrappers
+    "SecureTool",
+    "SecureResource",
+    "SecurePrompt",
+
+    # Decorators
+    "secure_tool",
+    "secure_resource",
+    "secure_prompt",
+
+    # Identity & Session
+    "ToolIdentity",
+    "ClientIdentity",
+    "SecureSession",
+    "SessionManager",
+    "create_tool_identity",
+
+    # Utils
+    "SecureAnnotationProcessor",
+    "encrypt_data",
+    "decrypt_data",
+]

src/mcp/server/fastmcp/secure/annotations.py

@@ -0,0 +1,326 @@
+"""
+Secure annotations for MCP tools, resources, and prompts.
+
+These annotations extend the standard MCP annotations with security features.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+from mcp.types import ToolAnnotations, ResourceAnnotations, PromptAnnotations
+
+
+class AuthMethod(Enum):
+    """Supported authentication methods."""
+    JWT = "jwt"
+    CERTIFICATE = "certificate"
+    TEE_ATTESTATION = "tee"
+    OAUTH = "oauth"
+    API_KEY = "api_key"
+    MTLS = "mtls"  # Mutual TLS
+
+
+@dataclass
+class SecureAnnotations:
+    """
+    Base security annotations that can be attached to tools, resources, or prompts.
+
+    These annotations enable security features like authentication, encryption,
+    and attestation for MCP operations.
+    """
+
+    # Authentication settings
+    require_auth: bool = False
+    auth_methods: list[AuthMethod] = field(default_factory=lambda: [AuthMethod.JWT])
+    required_permissions: set[str] = field(default_factory=set)
+    require_mutual_auth: bool = False  # Bidirectional authentication
+
+    # Encryption settings
+    encrypt_input: bool = False
+    encrypt_output: bool = False
+    encryption_algorithm: str = "AES-256-GCM"
+    key_exchange_method: str = "ECDH"  # ECDH, RSA, Pre-shared
+
+    # Tool/Server attestation
+    require_tool_attestation: bool = False
+    tool_certificate_fingerprint: Optional[str] = None
+    attestation_type: Optional[str] = None  # "software", "sgx", "sev", "trustzone"
+    tool_signature_required: bool = False
+
+    # Client verification
+    verify_client_certificate: bool = False
+    trusted_client_issuers: list[str] = field(default_factory=list)
+    client_attestation_required: bool = False
+
+    # Audit and compliance
+    audit_log: bool = True
+    audit_include_inputs: bool = False
+    audit_include_outputs: bool = False
+    audit_retention_days: int = 90
+
+    # Rate limiting
+    rate_limit: Optional[int] = None  # requests per minute
+    rate_limit_per_client: bool = True
+    burst_limit: Optional[int] = None
+
+    # Data handling
+    security_level: str = "standard"  # "standard", "high", "critical"
+    data_classification: str = "public"  # "public", "internal", "confidential", "secret"
+    compliance_tags: list[str] = field(default_factory=list)  # ["HIPAA", "PCI-DSS", "GDPR", "SOC2"]
+
+    # Session management
+    session_timeout_minutes: int = 60
+    require_session_binding: bool = False  # Bind session to client IP/fingerprint
+    max_concurrent_sessions: Optional[int] = None
+
+    # Advanced security
+    require_replay_protection: bool = False
+    max_request_age_seconds: int = 300  # For replay protection
+    require_integrity_check: bool = True  # Verify message integrity
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "require_auth": self.require_auth,
+            "auth_methods": [method.value for method in self.auth_methods],
+            "required_permissions": list(self.required_permissions),
+            "require_mutual_auth": self.require_mutual_auth,
+            "encrypt_input": self.encrypt_input,
+            "encrypt_output": self.encrypt_output,
+            "encryption_algorithm": self.encryption_algorithm,
+            "key_exchange_method": self.key_exchange_method,
+            "require_tool_attestation": self.require_tool_attestation,
+            "tool_certificate_fingerprint": self.tool_certificate_fingerprint,
+            "attestation_type": self.attestation_type,
+            "tool_signature_required": self.tool_signature_required,
+            "verify_client_certificate": self.verify_client_certificate,
+            "trusted_client_issuers": self.trusted_client_issuers,
+            "client_attestation_required": self.client_attestation_required,
+            "audit_log": self.audit_log,
+            "audit_include_inputs": self.audit_include_inputs,
+            "audit_include_outputs": self.audit_include_outputs,
+            "security_level": self.security_level,
+            "data_classification": self.data_classification,
+            "compliance_tags": self.compliance_tags,
+            "session_timeout_minutes": self.session_timeout_minutes,
+            "require_session_binding": self.require_session_binding,
+            "require_replay_protection": self.require_replay_protection,
+            "require_integrity_check": self.require_integrity_check,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SecureAnnotations:
+        """Create from dictionary."""
+        auth_methods = [AuthMethod(m) for m in data.get("auth_methods", ["jwt"])]
+        return cls(
+            require_auth=data.get("require_auth", False),
+            auth_methods=auth_methods,
+            required_permissions=set(data.get("required_permissions", [])),
+            require_mutual_auth=data.get("require_mutual_auth", False),
+            encrypt_input=data.get("encrypt_input", False),
+            encrypt_output=data.get("encrypt_output", False),
+            encryption_algorithm=data.get("encryption_algorithm", "AES-256-GCM"),
+            key_exchange_method=data.get("key_exchange_method", "ECDH"),
+            require_tool_attestation=data.get("require_tool_attestation", False),
+            tool_certificate_fingerprint=data.get("tool_certificate_fingerprint"),
+            attestation_type=data.get("attestation_type"),
+            tool_signature_required=data.get("tool_signature_required", False),
+            verify_client_certificate=data.get("verify_client_certificate", False),
+            trusted_client_issuers=data.get("trusted_client_issuers", []),
+            client_attestation_required=data.get("client_attestation_required", False),
+            audit_log=data.get("audit_log", True),
+            security_level=data.get("security_level", "standard"),
+            data_classification=data.get("data_classification", "public"),
+            compliance_tags=data.get("compliance_tags", []),
+        )
+
+
+class SecureToolAnnotations(ToolAnnotations):
+    """
+    Tool annotations with integrated security features.
+
+    This extends the standard ToolAnnotations with security metadata.
+    """
+
+    def __init__(
+        self,
+        secure: Optional[SecureAnnotations] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.secure = secure or SecureAnnotations()
+
+        # Store security annotations in extensions
+        if not hasattr(self, 'extensions'):
+            self.extensions = {}
+        self.extensions["security"] = self.secure.to_dict()
+
+    @classmethod
+    def create(
+        cls,
+        # Security parameters
+        require_auth: bool = False,
+        auth_methods: Optional[list[AuthMethod]] = None,
+        required_permissions: Optional[set[str]] = None,
+        encrypt_io: bool = False,
+        require_mutual_auth: bool = False,
+        security_level: str = "standard",
+
+        # Standard tool annotation parameters
+        audience: Optional[list[str]] = None,
+        capabilities: Optional[dict[str, Any]] = None,
+        **kwargs
+    ) -> SecureToolAnnotations:
+        """
+        Factory method to create secure tool annotations.
+
+        Args:
+            require_auth: Whether to require authentication
+            auth_methods: List of accepted authentication methods
+            required_permissions: Set of required permissions
+            encrypt_io: Whether to encrypt input/output
+            require_mutual_auth: Whether to require bidirectional authentication
+            security_level: Security level (standard/high/critical)
+            audience: Target audience for the tool
+            capabilities: Tool capabilities
+            **kwargs: Additional security parameters
+        """
+        secure_annotations = SecureAnnotations(
+            require_auth=require_auth,
+            auth_methods=auth_methods or [AuthMethod.JWT],
+            required_permissions=required_permissions or set(),
+            encrypt_input=encrypt_io,
+            encrypt_output=encrypt_io,
+            require_mutual_auth=require_mutual_auth,
+            security_level=security_level,
+            **kwargs
+        )
+
+        return cls(
+            secure=secure_annotations,
+            audience=audience,
+            capabilities=capabilities
+        )
+
+
+class SecureResourceAnnotations(ResourceAnnotations):
+    """
+    Resource annotations with integrated security features.
+
+    This extends the standard ResourceAnnotations with security metadata.
+    """
+
+    def __init__(
+        self,
+        secure: Optional[SecureAnnotations] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.secure = secure or SecureAnnotations()
+
+        # Store security annotations in extensions
+        if not hasattr(self, 'extensions'):
+            self.extensions = {}
+        self.extensions["security"] = self.secure.to_dict()
+
+    @classmethod
+    def create(
+        cls,
+        # Security parameters
+        require_auth: bool = False,
+        data_classification: str = "public",
+        encrypt_io: bool = False,
+        audit_access: bool = True,
+
+        # Standard resource annotation parameters
+        content_type: Optional[str] = None,
+        cache_control: Optional[str] = None,
+        **kwargs
+    ) -> SecureResourceAnnotations:
+        """
+        Factory method to create secure resource annotations.
+
+        Args:
+            require_auth: Whether to require authentication
+            data_classification: Data classification level
+            encrypt_io: Whether to encrypt input/output
+            audit_access: Whether to audit resource access
+            content_type: Resource content type
+            cache_control: Cache control headers
+            **kwargs: Additional security parameters
+        """
+        secure_annotations = SecureAnnotations(
+            require_auth=require_auth,
+            data_classification=data_classification,
+            encrypt_input=encrypt_io,
+            encrypt_output=encrypt_io,
+            audit_log=audit_access,
+            **kwargs
+        )
+
+        return cls(
+            secure=secure_annotations,
+            content_type=content_type,
+            cache_control=cache_control
+        )
+
+
+class SecurePromptAnnotations(PromptAnnotations):
+    """
+    Prompt annotations with integrated security features.
+
+    This extends the standard PromptAnnotations with security metadata.
+    """
+
+    def __init__(
+        self,
+        secure: Optional[SecureAnnotations] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.secure = secure or SecureAnnotations()
+
+        # Store security annotations in extensions
+        if not hasattr(self, 'extensions'):
+            self.extensions = {}
+        self.extensions["security"] = self.secure.to_dict()
+
+    @classmethod
+    def create(
+        cls,
+        # Security parameters
+        require_auth: bool = False,
+        audit_usage: bool = True,
+        compliance_tags: Optional[list[str]] = None,
+
+        # Standard prompt annotation parameters
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        **kwargs
+    ) -> SecurePromptAnnotations:
+        """
+        Factory method to create secure prompt annotations.
+
+        Args:
+            require_auth: Whether to require authentication
+            audit_usage: Whether to audit prompt usage
+            compliance_tags: Compliance tags (e.g., ["GDPR", "HIPAA"])
+            max_tokens: Maximum tokens for prompt
+            temperature: Temperature for prompt generation
+            **kwargs: Additional security parameters
+        """
+        secure_annotations = SecureAnnotations(
+            require_auth=require_auth,
+            audit_log=audit_usage,
+            compliance_tags=compliance_tags or [],
+            **kwargs
+        )
+
+        return cls(
+            secure=secure_annotations,
+            max_tokens=max_tokens,
+            temperature=temperature
+        )