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 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 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) 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 response = self._error(403, "Forbidden", "Missing required scope.", 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 path.endswith("/reschedule"): await self.rate_limiter.acquire_category( context, "social_schedule", limit=self.settings.social_schedule_requests_per_minute, window_seconds=60, ) elif path.endswith("/bulk"): await self.rate_limiter.acquire_category( context, "social_bulk", limit=max(1, self.settings.social_schedule_requests_per_minute // 4), window_seconds=60, ) @staticmethod def _bearer_token(header: str | None) -> str: if not header: raise UnauthorizedError parts = header.strip().split() if len(parts) != 2 or parts[0].casefold() != "bearer" or not parts[1]: raise UnauthorizedError return parts[1] @staticmethod def _content_length(value: str | None) -> int: try: return max(0, int(value or 0)) except ValueError: return 0 @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 = getattr(request.state, "request_id", None) if request_id: response_headers["X-Request-ID"] = request_id return JSONResponse( {"error": error, "message": message}, status_code=status_code, headers=response_headers, ) def _client_ip(self, request: Request) -> str | None: if self.settings.auth_trust_proxy_headers: forwarded = request.headers.get("x-forwarded-for") if forwarded: return forwarded.split(",", 1)[0].strip()[:64] return request.client.host[:64] if request.client else None async def _audit_request( self, request: Request, context: AuthContext, response_code: int, elapsed_ms: int, bytes_uploaded: int, bytes_downloaded: int, ) -> None: data = { "request_id": getattr(request.state, "request_id", "-"), "api_key_id": context.api_key_id, "key_name": context.key_name, "ip_address": self._client_ip(request), "user_agent": request.headers.get("user-agent", "")[:512] or None, "endpoint": request.url.path, "http_method": request.method, "response_code": response_code, "processing_time_ms": elapsed_ms, "bytes_uploaded": bytes_uploaded, "bytes_downloaded": bytes_downloaded, } try: await self.audit.record(**data) except Exception: logger.exception( "authentication audit persistence failed", extra={"api_key_id": context.api_key_id}, ) logger.info("authenticated request", extra=data)