"""
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),
]
