app/api/middleware.py

"""
app/api/middleware.py
=====================
Custom middleware and rate-limiting setup for the Aegis-ML FastAPI app.

- RequestLoggingMiddleware: structured request/response logging
- setup_rate_limiter: slowapi limiter factory
"""

from __future__ import annotations

import logging
import time
import uuid

from fastapi import FastAPI, Request, Response
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.util import get_remote_address
from starlette.middleware.base import BaseHTTPMiddleware

from app.config import get_settings

logger = logging.getLogger(__name__)


# ── Rate Limiter ──────────────────────────────────────────────────────────────


def create_limiter() -> Limiter:
    """Create a slowapi Limiter keyed by client IP address."""
    return Limiter(key_func=get_remote_address)


def setup_rate_limiter(app: FastAPI, limiter: Limiter) -> None:
    """Attach the limiter and its error handler to the FastAPI app."""
    app.state.limiter = limiter
    app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)


# ── Request Logging Middleware ────────────────────────────────────────────────


class RequestLoggingMiddleware(BaseHTTPMiddleware):
    """
    Structured request/response logging.
    Adds an X-Request-ID header to every response for traceability.
    """

    async def dispatch(self, request: Request, call_next) -> Response:
        request_id = str(uuid.uuid4())
        start = time.perf_counter()

        # Attach request_id to the request state so route handlers can read it
        request.state.request_id = request_id

        logger.info(
            "→ %s %s  [id=%s  ip=%s]",
            request.method,
            request.url.path,
            request_id,
            request.client.host if request.client else "unknown",
        )

        response: Response = await call_next(request)

        elapsed_ms = (time.perf_counter() - start) * 1000
        logger.info(
            "← %d  %s %s  %.1f ms  [id=%s]",
            response.status_code,
            request.method,
            request.url.path,
            elapsed_ms,
            request_id,
        )

        response.headers["X-Request-ID"] = request_id
        response.headers["X-Powered-By"] = "Aegis-ML"
        return response