""" Production Middleware Security, rate limiting, caching, request ID """ from fastapi import Request, Response from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.trustedhost import TrustedHostMiddleware from starlette.middleware.base import BaseHTTPMiddleware import time import uuid import hashlib from typing import Optional import redis class SecurityHeadersMiddleware(BaseHTTPMiddleware): """Add security headers to all responses""" async def dispatch(self, request: Request, call_next): response = await call_next(request) # Security headers response.headers["X-Content-Type-Options"] = "nosniff" response.headers["X-Frame-Options"] = "DENY" response.headers["X-XSS-Protection"] = "1; mode=block" response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains" response.headers["Content-Security-Policy"] = "default-src 'self'" response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin" response.headers["Permissions-Policy"] = "geolocation=(self), camera=(self)" return response class RequestIDMiddleware(BaseHTTPMiddleware): """Add request ID for tracing""" async def dispatch(self, request: Request, call_next): # Generate or extract request ID request_id = request.headers.get("X-Request-ID", str(uuid.uuid4())) request.state.request_id = request_id response = await call_next(request) response.headers["X-Request-ID"] = request_id return response class RateLimitMiddleware(BaseHTTPMiddleware): """Rate limiting per IP and user""" def __init__(self, app, redis_client: Optional[redis.Redis] = None): super().__init__(app) self.redis = redis_client or redis.Redis(host='localhost', port=6379, db=0) self.default_limit = 100 # requests per minute self.authenticated_limit = 1000 # requests per minute async def dispatch(self, request: Request, call_next): # Get client identifier client_id = self._get_client_id(request) # Check rate limit if not await self._check_rate_limit(client_id, request): return Response( content='{"error": "Rate limit exceeded"}', status_code=429, media_type="application/json", headers={"Retry-After": "60"} ) response = await call_next(request) # Add rate limit headers remaining = await self._get_remaining(client_id) response.headers["X-RateLimit-Limit"] = str(self.default_limit) response.headers["X-RateLimit-Remaining"] = str(remaining) return response def _get_client_id(self, request: Request) -> str: """Get client identifier""" # Try to get user ID from auth user_id = getattr(request.state, "user_id", None) if user_id: return f"user:{user_id}" # Fall back to IP forwarded = request.headers.get("X-Forwarded-For") if forwarded: return f"ip:{forwarded.split(',')[0].strip()}" return f"ip:{request.client.host}" async def _check_rate_limit(self, client_id: str, request: Request) -> bool: """Check if request is within rate limit""" key = f"ratelimit:{client_id}" # Use Redis for rate limiting try: pipe = self.redis.pipeline() pipe.incr(key) pipe.expire(key, 60) results = pipe.execute() count = results[0] # Check if authenticated is_authenticated = hasattr(request.state, "user_id") limit = self.authenticated_limit if is_authenticated else self.default_limit return count <= limit except redis.ConnectionError: # If Redis is down, allow request return True async def _get_remaining(self, client_id: str) -> int: """Get remaining requests""" key = f"ratelimit:{client_id}" try: count = int(self.redis.get(key) or 0) return max(0, self.default_limit - count) except redis.ConnectionError: return self.default_limit class CacheMiddleware(BaseHTTPMiddleware): """Response caching with Redis""" def __init__(self, app, redis_client: Optional[redis.Redis] = None, ttl: int = 300): super().__init__(app) self.redis = redis_client or redis.Redis(host='localhost', port=6379, db=1) self.ttl = ttl self.cacheable_paths = [ "/taxonomy/domains", "/taxonomy/domains/", "/defects", "/health" ] async def dispatch(self, request: Request, call_next): # Only cache GET requests if request.method != "GET": return await call_next(request) # Check if path is cacheable if not any(request.url.path.startswith(path) for path in self.cacheable_paths): return await call_next(request) # Generate cache key cache_key = self._generate_cache_key(request) # Try to get from cache try: cached = self.redis.get(cache_key) if cached: return Response( content=cached, media_type="application/json", headers={"X-Cache": "HIT"} ) except redis.ConnectionError: pass # Get from origin response = await call_next(request) # Cache response if response.status_code == 200: try: body = b"" async for chunk in response.body_iterator: body += chunk self.redis.setex(cache_key, self.ttl, body) return Response( content=body, status_code=response.status_code, headers=dict(response.headers), media_type=response.media_type ) except redis.ConnectionError: pass return response def _generate_cache_key(self, request: Request) -> str: """Generate cache key from request""" key = f"cache:{request.method}:{request.url.path}" if request.query_params: key += f"?{request.query_params}" return hashlib.md5(key.encode()).hexdigest() class LoggingMiddleware(BaseHTTPMiddleware): """Request/response logging""" async def dispatch(self, request: Request, call_next): start_time = time.time() # Log request request_id = getattr(request.state, "request_id", "unknown") print(f"[{request_id}] {request.method} {request.url.path} - Started") try: response = await call_next(request) # Log response duration = time.time() - start_time print(f"[{request_id}] {request.method} {request.url.path} - {response.status_code} ({duration:.3f}s)") return response except Exception as e: duration = time.time() - start_time print(f"[{request_id}] {request.method} {request.url.path} - ERROR ({duration:.3f}s): {e}") raise class GracefulShutdownMiddleware(BaseHTTPMiddleware): """Handle graceful shutdown""" def __init__(self, app): super().__init__(app) self.shutdown = False async def dispatch(self, request: Request, call_next): if self.shutdown: return Response( content='{"error": "Server is shutting down"}', status_code=503, media_type="application/json", headers={"Retry-After": "30"} ) return await call_next(request) def initiate_shutdown(self): """Initiate graceful shutdown""" self.shutdown = True def setup_middleware(app): """Setup all middleware""" # Security app.add_middleware(SecurityHeadersMiddleware) app.add_middleware(TrustedHostMiddleware, allowed_hosts=["*.landvex.com", "*.quixzoom.com", "localhost"]) # CORS app.add_middleware( CORSMiddleware, allow_origins=["https://landvex.com", "https://quixzoom.com", "https://admin.landvex.com"], allow_credentials=True, allow_methods=["GET", "POST", "PUT", "DELETE"], allow_headers=["*"], max_age=3600 ) # Request tracing app.add_middleware(RequestIDMiddleware) # Rate limiting app.add_middleware(RateLimitMiddleware) # Caching app.add_middleware(CacheMiddleware) # Logging app.add_middleware(LoggingMiddleware) # Graceful shutdown app.add_middleware(GracefulShutdownMiddleware)