from __future__ import annotations import time from typing import Any from uuid import uuid4 from fastapi import Request from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.responses import Response from app.core.config import Settings from app.core.logger import get_logger from app.security.audit import AuditService from app.security.context import AuthContext, auth_context, http_auth_applied from app.security.errors import ForbiddenError, RateLimitError, UnauthorizedError from app.security.policy import ScopePolicy from app.security.rate_limit import APIKeyRateLimiter, RateLimitLease from app.security.service import APIKeyService from app.security.studio_identity import constrain_context_to_membership, verify_studio_identity logger = get_logger(__name__) class APIKeyAuthenticationMiddleware(BaseHTTPMiddleware): """Authenticates, authorizes, rate-limits, and audits protected HTTP requests.""" def __init__( self, app: Any, *, settings: Settings, api_keys: APIKeyService, rate_limiter: APIKeyRateLimiter, audit: AuditService, ) -> None: super().__init__(app) self.settings = settings self.api_keys = api_keys self.rate_limiter = rate_limiter self.audit = audit self.policy = ScopePolicy() async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: if not getattr(request.state, "request_id", None): request.state.request_id = str(uuid4()) if not self.settings.auth_enabled or self.policy.is_public(request): return await call_next(request) started = time.monotonic() context: AuthContext | None = None lease: RateLimitLease | None = None context_token = None http_auth_token = None response_code = 500 required_scope: str | None = None bytes_uploaded = self._content_length(request.headers.get("content-length")) bytes_downloaded = 0 try: api_key = self._bearer_token(request.headers.get("authorization")) context = await self.api_keys.authenticate(api_key) studio_header_names = ( "x-mr-studio-subject", "x-mr-studio-email", "x-mr-studio-name", "x-mr-studio-timestamp", "x-mr-studio-signature", ) if any(request.headers.get(name) is not None for name in studio_header_names): secret = ( self.settings.studio_identity_secret.get_secret_value() if self.settings.studio_identity_secret else "" ) try: identity = verify_studio_identity(request.headers, secret) except ValueError as exc: logger.warning("rejected Studio identity envelope: %s", str(exc)) raise UnauthorizedError from exc principal = await self.api_keys.tenants.resolve_studio_subject(identity.subject) request.state.studio_identity_verified = True try: # The shared key authenticates transport only; persisted # membership capabilities constrain its effective scopes. context = constrain_context_to_membership(context, principal) except ValueError as exc: logger.warning("rejected Studio membership role: %s", str(exc)) raise ForbiddenError from exc request.state.auth = context context_token = auth_context.set(context) required_scope = await self.policy.required_scope(request) lease = await self.rate_limiter.acquire( context, is_job=self.policy.is_job(required_scope) or request.url.path.startswith("/v1/projects/") and "/renders" in request.url.path and request.method == "POST", is_upload=self.policy.is_upload(request, required_scope), uploaded_bytes=bytes_uploaded, ) if context.membership_role == "viewer" and required_scope not in { "templates:read", "operations:read", "jobs:read", "assets:read", "mcp:read", "system:read", "social:accounts:read", "social:posts:read", "social:schedules:read", "social:analytics:read", "analytics:read", "generation:providers:read", "generation:requests:read", "ai:read", "copilot:read", "projects:read", "members:read", "teams:read", "projects:collaborate", "comments:read", "approvals:read", }: raise ForbiddenError self.api_keys.authorize(context, required_scope) await self._apply_social_rate_limit(request, context) await self.api_keys.mark_used(context) http_auth_token = http_auth_applied.set(True) response = await call_next(request) response_code = response.status_code bytes_downloaded = self._content_length(response.headers.get("content-length")) response.headers.setdefault("X-Request-ID", request.state.request_id) return response except UnauthorizedError: response_code = 401 return self._error( 401, "Unauthorized", "Invalid or expired API key.", request, {"WWW-Authenticate": "Bearer"}, ) except ForbiddenError: response_code = 403 detail = f"Missing required scope: {required_scope}." if required_scope else "Workspace identity or membership authorization was denied." response = self._error(403, "Forbidden", detail, request) bytes_downloaded = len(response.body) return response except RateLimitError as exc: response_code = 429 response = self._error( 429, "Rate limit exceeded", "Retry later.", request, {"Retry-After": str(exc.retry_after)}, ) bytes_downloaded = len(response.body) return response finally: if lease is not None: await lease.release() if http_auth_token is not None: http_auth_applied.reset(http_auth_token) if context_token is not None: auth_context.reset(context_token) if context is not None: elapsed_ms = max(0, round((time.monotonic() - started) * 1000)) await self._audit_request( request, context, response_code, elapsed_ms, bytes_uploaded, bytes_downloaded, ) async def _apply_social_rate_limit(self, request: Request, context: AuthContext) -> None: path = request.url.path if path.startswith("/v1/analytics"): category = "analytics_sync" if path.endswith(("/sync", "/cancel")) else "analytics_read" limit = ( max(1, self.settings.social_analytics_requests_per_minute // 6) if category == "analytics_sync" else self.settings.social_analytics_requests_per_minute ) await self.rate_limiter.acquire_category( context, category, limit=limit, window_seconds=60, ) return if not path.startswith("/v1/social"): return if "/analytics" in path: await self.rate_limiter.acquire_category( context, "social_analytics", limit=self.settings.social_analytics_requests_per_minute, window_seconds=60, ) elif path.endswith("/connect") or path.endswith("/callback") or path.endswith("/refresh"): await self.rate_limiter.acquire_category( context, "social_oauth", limit=self.settings.social_oauth_requests_per_hour, window_seconds=3600, ) elif path.endswith("/publish"): await self.rate_limiter.acquire_category( context, "social_publish", limit=self.settings.social_publish_requests_per_minute, window_seconds=60, ) elif path.endswith("/schedule"): await self.rate_limiter.acquire_category( context, "social_schedule", limit=self.settings.social_schedule_requests_per_minute, window_seconds=60, ) elif "/batch" in path: await self.rate_limiter.acquire_category( context, "social_publish", limit=self.settings.social_publish_requests_per_minute, window_seconds=60, ) async def _audit_request( self, request: Request, context: AuthContext, response_code: int, elapsed_ms: int, bytes_uploaded: int, bytes_downloaded: int, ) -> None: try: await self.audit.record( request_id=str(getattr(request.state, "request_id", ""))[:64], api_key_id=str(context.api_key_id), key_name=str(getattr(context, "key_name", ""))[:128], ip_address=request.client.host if request.client else None, user_agent=request.headers.get("user-agent", "")[:512] or None, endpoint=request.url.path[:512], http_method=request.method[:16], response_code=response_code, processing_time_ms=elapsed_ms, bytes_uploaded=max(0, bytes_uploaded), bytes_downloaded=max(0, bytes_downloaded), ) except Exception: logger.exception("failed to persist API request audit record") @staticmethod def _error( status_code: int, error: str, message: str, request: Request, headers: dict[str, str] | None = None, ) -> JSONResponse: response_headers = dict(headers or {}) request_id = str(getattr(request.state, "request_id", "")) if request_id: response_headers.setdefault("X-Request-ID", request_id) return JSONResponse( status_code=status_code, headers=response_headers, content={ "error": { "code": error, "message": message, "request_id": request_id, } }, ) @staticmethod def _bearer_token(value: str | None) -> str: if not value: raise UnauthorizedError scheme, separator, token = value.partition(" ") if not separator or scheme.casefold() != "bearer" or not token.strip(): raise UnauthorizedError return token.strip() @staticmethod def _content_length(value: str | None) -> int: if not value: return 0 try: parsed = int(value) except (TypeError, ValueError): return 0 return max(0, parsed)