mcphub/core/endpoints/middleware.py

"""
Middleware for Multi-Endpoint Architecture

Provides authentication, rate limiting, and audit logging
that works with the multi-endpoint architecture.
"""

import logging
import time
from collections.abc import Callable
from dataclasses import dataclass

from fastmcp.exceptions import ToolError
from fastmcp.server.dependencies import get_http_headers
from fastmcp.server.middleware import Middleware, MiddlewareContext

from core.api_keys import get_api_key_manager
from core.audit_log import EventType, LogLevel, get_audit_logger
from core.auth import get_auth_manager
from core.context import clear_api_key_context, set_api_key_context
from core.rate_limiter import get_rate_limiter

from .config import EndpointConfig

logger = logging.getLogger(__name__)


@dataclass
class AuthContext:
    """Authentication context for a request"""

    key_id: str | None = None
    project_id: str | None = None
    scope: str = "read"
    is_master_key: bool = False
    is_oauth_token: bool = False
    client_ip: str | None = None


class EndpointAuthMiddleware(Middleware):
    """
    Authentication middleware for multi-endpoint architecture.

    Validates API keys/tokens and enforces endpoint-specific access rules.
    """

    def __init__(self, endpoint_config: EndpointConfig):
        """
        Initialize middleware with endpoint configuration.

        Args:
            endpoint_config: Configuration for this endpoint
        """
        self.config = endpoint_config
        self.auth_manager = get_auth_manager()
        self.api_key_manager = get_api_key_manager()

    async def on_call_tool(self, context: MiddlewareContext, call_next: Callable):
        """
        Handle tool call with authentication and authorization.

        Args:
            context: Middleware context
            call_next: Next middleware in chain
        """
        tool_name = getattr(context.message, "name", "unknown")
        start_time = time.time()

        try:
            # Extract and validate authentication
            auth_context = await self._authenticate(context)

            # Check endpoint access
            self._check_endpoint_access(auth_context)

            # Check tool access
            self._check_tool_access(tool_name, auth_context)

            # Set context for downstream handlers
            if auth_context.key_id:
                set_api_key_context(
                    key_id=auth_context.key_id,
                    project_id=auth_context.project_id or "*",
                    scope=auth_context.scope,
                    is_global=auth_context.project_id == "*",
                )

            # Call the actual tool
            result = await call_next(context)

            # Log success
            self._log_success(tool_name, auth_context, start_time)

            return result

        except ToolError:
            raise
        except Exception as e:
            self._log_error(tool_name, str(e), start_time)
            raise ToolError(f"Authentication error: {str(e)}")
        finally:
            clear_api_key_context()

    async def _authenticate(self, context: MiddlewareContext) -> AuthContext:
        """
        Extract and validate authentication from request.

        Args:
            context: Middleware context

        Returns:
            AuthContext with authentication details
        """
        auth_context = AuthContext()

        # Get headers
        try:
            headers = get_http_headers()
        except Exception:
            headers = {}

        # Extract client IP
        auth_context.client_ip = headers.get("x-forwarded-for", "unknown")

        # Get authorization header
        auth_header = headers.get("authorization", "")

        if not auth_header:
            # No auth provided
            if self.config.require_master_key:
                raise ToolError("Master API key required for this endpoint")
            return auth_context

        # Parse authorization
        if auth_header.startswith("Bearer "):
            token = auth_header[7:]
        else:
            token = auth_header

        # Check token type
        if token.startswith("sk-"):
            # Master API key
            if self.auth_manager.validate_master_key(token):
                auth_context.is_master_key = True
                auth_context.project_id = "*"
                auth_context.scope = "admin"
                auth_context.key_id = "master"
                return auth_context
            else:
                raise ToolError("Invalid master API key")

        elif token.startswith("cmp_"):
            # Project API key
            key = self.api_key_manager.get_key_by_token(token)
            if not key:
                raise ToolError("Invalid API key")

            if key.revoked:
                raise ToolError("API key has been revoked")

            if key.is_expired():
                raise ToolError("API key has expired")

            auth_context.key_id = key.key_id
            auth_context.project_id = key.project_id
            auth_context.scope = key.scope
            return auth_context

        else:
            # Possibly OAuth token (JWT)
            try:
                from core.oauth import get_token_manager

                token_manager = get_token_manager()
                payload = token_manager.validate_access_token(token)

                if payload:
                    auth_context.is_oauth_token = True
                    auth_context.project_id = payload.get("project_id", "*")
                    auth_context.scope = payload.get("scope", "read")
                    auth_context.key_id = f"oauth_{payload.get('sub', 'unknown')}"
                    return auth_context
            except Exception:
                pass

            raise ToolError("Invalid authentication token")

    def _check_endpoint_access(self, auth_context: AuthContext):
        """
        Check if auth context allows access to this endpoint.

        Args:
            auth_context: Authentication context
        """
        # Master key always has access
        if auth_context.is_master_key:
            return

        # Check if endpoint requires master key
        if self.config.require_master_key:
            raise ToolError(f"Endpoint {self.config.path} requires master API key")

        # Check scope requirements
        if self.config.allowed_scopes:
            # Check if any of the user's scopes are allowed
            user_scopes = set(auth_context.scope.split())
            if not user_scopes & self.config.allowed_scopes:
                raise ToolError(
                    f"Insufficient scope. Required: {self.config.allowed_scopes}, "
                    f"Got: {user_scopes}"
                )

        # Check plugin type access
        if auth_context.project_id and auth_context.project_id != "*":
            # Extract plugin type from project_id (e.g., "wordpress_site4" -> "wordpress")
            if "_" in auth_context.project_id:
                key_plugin_type = auth_context.project_id.split("_")[0]

                # Check if endpoint allows this plugin type
                if self.config.plugin_types and key_plugin_type not in self.config.plugin_types:
                    raise ToolError(
                        f"API key for {key_plugin_type} cannot access "
                        f"{self.config.endpoint_type.value} endpoint"
                    )

    def _check_tool_access(self, tool_name: str, auth_context: AuthContext):
        """
        Check if auth context allows access to specific tool.

        Args:
            tool_name: Name of the tool
            auth_context: Authentication context
        """
        # Master key has access to all tools
        if auth_context.is_master_key:
            return

        # Check tool blacklist
        if not self.config.allows_tool(tool_name):
            raise ToolError(f"Access denied to tool: {tool_name}")

        # Check site filter for project endpoints
        if self.config.site_filter:
            # Tool must be for the configured site
            # This is handled by parameter injection in the wrapper
            pass

    def _log_success(self, tool_name: str, auth_context: AuthContext, start_time: float):
        """Log successful tool execution"""
        duration_ms = int((time.time() - start_time) * 1000)
        logger.debug(
            f"Tool {tool_name} executed successfully "
            f"(key={auth_context.key_id}, duration={duration_ms}ms)"
        )

    def _log_error(self, tool_name: str, error: str, start_time: float):
        """Log tool execution error"""
        duration_ms = int((time.time() - start_time) * 1000)
        logger.warning(f"Tool {tool_name} failed: {error} (duration={duration_ms}ms)")


class EndpointRateLimitMiddleware(Middleware):
    """
    Rate limiting middleware for multi-endpoint architecture.
    """

    def __init__(self, endpoint_config: EndpointConfig):
        self.config = endpoint_config
        self.rate_limiter = get_rate_limiter()

    async def on_call_tool(self, context: MiddlewareContext, call_next: Callable):
        """Apply rate limiting before tool execution"""
        # Get client identifier
        try:
            headers = get_http_headers()
            client_id = headers.get("authorization", "anonymous")[:50]
        except Exception:
            client_id = "unknown"

        # Check rate limit
        allowed, info = self.rate_limiter.check_rate_limit(client_id)

        if not allowed:
            raise ToolError(
                f"Rate limit exceeded. Retry after {info.get('retry_after', 60)} seconds"
            )

        # Proceed with request
        return await call_next(context)


class EndpointAuditMiddleware(Middleware):
    """
    Audit logging middleware for multi-endpoint architecture.
    """

    def __init__(self, endpoint_config: EndpointConfig):
        self.config = endpoint_config
        self.audit_logger = get_audit_logger()

    async def on_call_tool(self, context: MiddlewareContext, call_next: Callable):
        """Log tool execution to audit log"""
        tool_name = getattr(context.message, "name", "unknown")
        start_time = time.time()

        try:
            result = await call_next(context)

            # Log success
            self.audit_logger.log(
                level=LogLevel.INFO,
                event_type=EventType.TOOL_CALL,
                message=f"Tool executed: {tool_name}",
                details={
                    "tool": tool_name,
                    "endpoint": self.config.path,
                    "duration_ms": int((time.time() - start_time) * 1000),
                    "success": True,
                },
            )

            return result

        except Exception as e:
            # Log failure
            self.audit_logger.log(
                level=LogLevel.WARNING,
                event_type=EventType.TOOL_CALL,
                message=f"Tool failed: {tool_name}",
                details={
                    "tool": tool_name,
                    "endpoint": self.config.path,
                    "duration_ms": int((time.time() - start_time) * 1000),
                    "success": False,
                    "error": str(e),
                },
            )
            raise


def create_endpoint_middleware(endpoint_config: EndpointConfig) -> list:
    """
    Create middleware stack for an endpoint.

    Args:
        endpoint_config: Endpoint configuration

    Returns:
        List of middleware instances
    """
    return [
        EndpointAuthMiddleware(endpoint_config),
        EndpointRateLimitMiddleware(endpoint_config),
        EndpointAuditMiddleware(endpoint_config),
    ]