app/middleware/rate_limit_middleware.py

"""Author RAG Chatbot SaaS — Rate Limiting Middleware (Pure ASGI).

Implements sliding window rate limiting using Redis.
- Chat endpoint:    60 requests/minute per IP (main POST only)
- Auth endpoint:    10 requests/minute per IP
- Visitor chat:     60 messages/hour per visitor fingerprint (configurable)

Uses pure ASGI to avoid BaseHTTPMiddleware request body consumption issues.

Fixes:
  B10 — Visitor limit targets only POST /api/chat/{slug} (not session/init etc.).
  B11 — Blocked 429 responses do not consume quota (decr after reject).
  B12 — IP chat limit applies to main chat POST only, not init/track-click.
"""

import hashlib
import json
import re

import structlog
from starlette.types import ASGIApp, Receive, Scope, Send

from app.config import get_settings

logger = structlog.get_logger(__name__)
cfg = get_settings()

# Path-to-limit mapping: (path_prefix, limit_per_minute)
_RATE_LIMIT_RULES: list[tuple[str, int]] = [
    ("/api/auth", cfg.RATE_LIMIT_AUTH_PER_MINUTE),
    ("/api/admin", 120),  # Admin API: generous but bounded
]

# Regex for the MAIN chat endpoint only (not /session/init, /track-click etc.)
# Pattern: /api/chat/{author_slug}  — exactly two path segments after /api/chat
_MAIN_CHAT_PATTERN = re.compile(r"^/api/chat/[^/]+$")

# Visitor rate limit: messages/hour per fingerprint
_VISITOR_LIMIT = cfg.RATE_LIMIT_VISITOR_PER_HOUR
_VISITOR_WINDOW = 3600  # 1 hour in seconds
_CHAT_IP_LIMIT = cfg.RATE_LIMIT_CHAT_PER_MINUTE


class RateLimitMiddleware:
    """Sliding window rate limiter backed by Redis (pure ASGI).

    Two independent limits on the main chat POST only:
    1. IP-based limit: per-minute (general abuse prevention)
    2. Visitor fingerprint limit: per-hour (budget protection)
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")
        method = scope.get("method", "GET")

        try:
            from app.dependencies import get_redis
            redis = await get_redis()
        except Exception as e:
            logger.error("Rate limit check failed (Redis error)", error=str(e))
            await self.app(scope, receive, send)
            return

        client_ip = _get_client_ip(scope)

        # ── Main chat POST: IP + visitor limits ───────────────────────────────
        if method == "POST" and _MAIN_CHAT_PATTERN.match(path):
            ip_key = f"ratelimit:chat:{client_ip}"
            ip_count = await redis.incr(ip_key)
            if ip_count == 1:
                await redis.expire(ip_key, 60)
            if ip_count > _CHAT_IP_LIMIT:
                await redis.decr(ip_key)
                logger.warning("IP rate limit exceeded", ip=client_ip, path=path, count=ip_count)
                await _send_429(send, f"Too many requests. Limit: {_CHAT_IP_LIMIT}/minute.")
                return

            headers = dict(scope.get("headers", []))
            ua = headers.get(b"user-agent", b"").decode()
            accept_lang = headers.get(b"accept-language", b"").decode()
            fp = _make_fingerprint(client_ip, ua, accept_lang)
            visitor_key = f"vrl:{fp}"

            v_count = await redis.incr(visitor_key)
            if v_count == 1:
                await redis.expire(visitor_key, _VISITOR_WINDOW)
            if v_count > _VISITOR_LIMIT:
                await redis.decr(visitor_key)
                logger.warning(
                    "Visitor rate limit exceeded",
                    fp=fp[:8], path=path, count=v_count,
                )
                await _send_429(
                    send,
                    "You've sent a lot of messages! Take a short break and come back in a bit.",
                )
                return

        else:
            # ── Other routes: prefix-based IP limits (auth, admin) ──────────
            ip_limit = _get_limit_for_path(path)
            if ip_limit is not None:
                key_segment = path.split("/")[2] if len(path.split("/")) > 2 else "unknown"
                ip_key = f"ratelimit:{key_segment}:{client_ip}"
                count = await redis.incr(ip_key)
                if count == 1:
                    await redis.expire(ip_key, 60)
                if count > ip_limit:
                    await redis.decr(ip_key)
                    logger.warning("IP rate limit exceeded", ip=client_ip, path=path, count=count)
                    await _send_429(send, f"Too many requests. Limit: {ip_limit}/minute.")
                    return

        await self.app(scope, receive, send)


def _get_limit_for_path(path: str) -> int | None:
    """Find the IP-based rate limit for a given path prefix."""
    for prefix, limit in _RATE_LIMIT_RULES:
        if path.startswith(prefix):
            return limit
    return None


def _get_client_ip(scope: Scope) -> str:
    """Extract client IP — checks X-Forwarded-For first (proxy/Docker aware)."""
    headers = dict(scope.get("headers", []))
    forwarded = headers.get(b"x-forwarded-for", b"").decode()
    if forwarded:
        return forwarded.split(",")[0].strip()
    client = scope.get("client")
    return client[0] if client else "unknown"


def _make_fingerprint(ip: str, ua: str, accept_lang: str) -> str:
    """Build a short anonymized fingerprint from non-PII signals."""
    raw = f"{ip}|{ua}|{accept_lang}"
    return hashlib.sha256(raw.encode()).hexdigest()[:32]


async def _send_429(send: Send, message: str) -> None:
    """Send a 429 Too Many Requests response with a human-readable message."""
    body = json.dumps({"detail": message}).encode()
    await send({
        "type": "http.response.start",
        "status": 429,
        "headers": [
            (b"content-type", b"application/json"),
            (b"retry-after", b"60"),
            (b"content-length", str(len(body)).encode()),
        ],
    })
    await send({
        "type": "http.response.body",
        "body": body,
    })