boc/iom/api/middleware.py

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