app/middleware/rate_limit.py

"""Author RAG — Rate Limiting Middleware.

Protects auth endpoints from brute-force attacks.
Uses Redis (fakeredis in tests) for distributed rate limiting.

Limits:
  /api/auth/login          → 10 attempts / 15 minutes / IP
  /api/auth/forgot-password → 3 attempts / hour / IP
  /api/billing/webhook     → unlimited (Stripe IPs only)
  All others               → 300 req / minute / IP

Implementation: Token bucket via Redis INCR + EXPIRE.
"""

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

logger = structlog.get_logger(__name__)

# (path_prefix, max_calls, window_seconds)
RATE_RULES = [
    ("/api/auth/login",            10,  900),   # 10/15min
    ("/api/auth/forgot-password",   3, 3600),   # 3/hour
    ("/api/auth/reset-password",    5, 3600),   # 5/hour
    ("/api/auth/activate",         10,  900),   # 10/15min
    ("/api/",                     300,   60),   # 300/min default
]

EXEMPT_PREFIXES = ("/api/billing/webhook", "/health", "/docs", "/openapi")


def _get_rule(path: str):
    for prefix, limit, window in RATE_RULES:
        if path.startswith(prefix):
            return limit, window
    return 300, 60  # default


class RateLimitMiddleware(BaseHTTPMiddleware):
    """IP-based rate limiting using Redis atomic INCR."""

    def __init__(self, app, redis_client=None):
        super().__init__(app)
        self._redis = redis_client  # injected; falls back to no-op if None

    async def dispatch(self, request: Request, call_next) -> Response:
        path = request.url.path

        # Skip exempt paths
        if any(path.startswith(p) for p in EXEMPT_PREFIXES):
            return await call_next(request)

        # Only rate-limit API paths
        if not path.startswith("/api/"):
            return await call_next(request)

        if not self._redis:
            return await call_next(request)

        ip = request.client.host if request.client else "unknown"
        limit, window = _get_rule(path)

        # Bucket key: ip:path_prefix:window_slot
        slot = int(time.time() // window)
        key = f"rl:{ip}:{path[:40]}:{slot}"

        try:
            count = await self._redis.incr(key)
            if count == 1:
                await self._redis.expire(key, window)

            # Add rate limit headers
            response = await call_next(request)
            response.headers["X-RateLimit-Limit"] = str(limit)
            response.headers["X-RateLimit-Remaining"] = str(max(0, limit - count))
            response.headers["X-RateLimit-Reset"] = str((slot + 1) * window)

            if count > limit:
                logger.warning("Rate limit exceeded", ip=ip, path=path, count=count)
                return JSONResponse(
                    status_code=429,
                    content={
                        "detail": "Too many requests. Please wait before trying again.",
                        "retry_after": window,
                    },
                    headers={
                        "Retry-After": str(window),
                        "X-RateLimit-Limit": str(limit),
                        "X-RateLimit-Remaining": "0",
                    },
                )
            return response
        except Exception as e:
            # Never fail open — if Redis is down, let request through
            logger.warning("Rate limiter error — passing request through", error=str(e))
            return await call_next(request)