api/middleware.py

"""
api/middleware.py — Rate limiting middleware and tenant context helpers.

Rate limiter: sliding-window, per-API-key, 100 req/60s default.
Tenant context: extracted from User object (set by auth dependency).
"""

from __future__ import annotations

import logging
import time
from collections import defaultdict, deque
from typing import TYPE_CHECKING

from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse

if TYPE_CHECKING:
    from api.auth import User

log = logging.getLogger(__name__)

# ─── Rate limiter ─────────────────────────────────────────────────────────────


class RateLimiter:
    """
    Sliding-window per-key rate limiter (in-memory).

    Thread-safe for CPython GIL; for multi-process deployment use Redis instead.
    """

    def __init__(self, max_requests: int = 100, window_seconds: int = 60) -> None:
        self._requests: dict[str, deque] = defaultdict(deque)
        self.max_requests = max_requests
        self.window_seconds = window_seconds

    def is_allowed(self, key: str) -> bool:
        """
        Check if key is within rate limit.

        Returns True and records the request if allowed.
        Returns False (and does NOT record) if the limit is exceeded.
        """
        now = time.time()
        window = self._requests[key]

        # Evict expired timestamps
        cutoff = now - self.window_seconds
        while window and window[0] < cutoff:
            window.popleft()

        if len(window) >= self.max_requests:
            return False

        window.append(now)
        return True

    def remaining(self, key: str) -> int:
        """Return how many requests are left in the current window."""
        now = time.time()
        window = self._requests[key]
        cutoff = now - self.window_seconds
        # Count non-expired
        active = sum(1 for t in window if t >= cutoff)
        return max(0, self.max_requests - active)

    def reset_in(self, key: str) -> float:
        """Return seconds until the oldest request expires (i.e. window resets)."""
        window = self._requests.get(key)
        if not window:
            return 0.0
        oldest = window[0]
        return max(0.0, (oldest + self.window_seconds) - time.time())

    def reset_key(self, key: str) -> None:
        """Clear the rate limit state for a key (admin action)."""
        self._requests.pop(key, None)


# Singleton used by the middleware and can be imported by tests
rate_limiter = RateLimiter(max_requests=100, window_seconds=60)


# ─── FastAPI middleware ────────────────────────────────────────────────────────


class RateLimitMiddleware(BaseHTTPMiddleware):
    """
    Starlette/FastAPI middleware that enforces per-API-key rate limits.

    Health and docs endpoints are exempt.
    """

    EXEMPT_PATHS = {"/health", "/docs", "/redoc", "/openapi.json", "/favicon.ico"}

    def __init__(self, app, limiter: RateLimiter | None = None) -> None:
        super().__init__(app)
        self._limiter = limiter or rate_limiter

    async def dispatch(self, request: Request, call_next) -> Response:
        # Exempt health/docs
        if request.url.path in self.EXEMPT_PATHS:
            return await call_next(request)

        api_key = request.headers.get("X-API-Key", "anonymous")
        allowed = self._limiter.is_allowed(api_key)

        if not allowed:
            reset_secs = self._limiter.reset_in(api_key)
            log.warning("Rate limit exceeded for key %r (resets in %.1fs)", api_key[:8], reset_secs)
            return JSONResponse(
                status_code=429,
                content={
                    "detail": "Rate limit exceeded. Max 100 requests per 60 seconds.",
                    "retry_after_seconds": round(reset_secs, 1),
                },
                headers={
                    "Retry-After": str(int(reset_secs) + 1),
                    "X-RateLimit-Limit": str(self._limiter.max_requests),
                    "X-RateLimit-Remaining": "0",
                    "X-RateLimit-Reset": str(int(time.time() + reset_secs)),
                },
            )

        response = await call_next(request)

        # Inject rate-limit headers on allowed responses
        remaining = self._limiter.remaining(api_key)
        response.headers["X-RateLimit-Limit"] = str(self._limiter.max_requests)
        response.headers["X-RateLimit-Remaining"] = str(remaining)
        response.headers["X-RateLimit-Window"] = str(self._limiter.window_seconds)

        return response


# ─── Tenant context helpers ───────────────────────────────────────────────────


def get_tenant_id(user: "User") -> str:
    """Extract tenant_id from authenticated User object."""
    return user.tenant_id


def tenant_scoped_key(user: "User", suffix: str) -> str:
    """Build a namespaced cache / rate-limit key like 'tenant-default:queries'."""
    return f"{user.tenant_id}:{suffix}"