xiaoxia-saas/apps/api/app/middleware/logging.py

"""请求日志和限流中间件"""

import logging
import re
import time
from typing import Optional

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware

logger = logging.getLogger(__name__)

# 敏感参数名称模式（不区分大小写）
SENSITIVE_PARAM_PATTERNS = re.compile(
    r"^(password|passwd|pwd|token|secret|key|authorization|auth|api_key|"  # noqa: E501
    r"apikey|access_token|refresh_token|accesstoken|refreshtoken|"  # noqa: E501
    r"session_id|sessionid|sid|cookie|csrf|xsrf|bearer)$",  # noqa: E501
    re.IGNORECASE,
)


def filter_sensitive_params(query_string: Optional[str]) -> Optional[str]:
    """
    过滤 query string 中的敏感参数

    Args:
        query_string: 原始 query string，例如 "name=xxx&password=secret&token=abc"

    Returns:
        过滤后的 query string，敏感参数的值被替换为 "***"
        如果 query_string 为空或 None，返回原始值
    """
    if not query_string:
        return query_string

    if query_string.startswith("?"):
        query_string = query_string[1:]

    if not query_string:
        return query_string

    parts = query_string.split("&")
    filtered_parts = []

    for part in parts:
        if "=" in part:
            key, value = part.split("=", 1)
            if SENSITIVE_PARAM_PATTERNS.match(key):
                filtered_parts.append(f"{key}=***")
            else:
                filtered_parts.append(part)
        else:
            # 没有 = 的参数，保留原样
            filtered_parts.append(part)

    return "&".join(filtered_parts)


class RequestLoggingMiddleware(BaseHTTPMiddleware):
    """请求日志中间件"""

    async def dispatch(self, request: Request, call_next):
        # 记录请求开始时间
        start_time = time.time()

        # 过滤 query string 中的敏感参数
        raw_query = str(request.url.query) if request.url.query else ""
        safe_query = filter_sensitive_params(raw_query)

        # 记录请求信息（不包含敏感参数）
        if safe_query:
            logger.info(f"Request: {request.method} {request.url.path}?{safe_query}")  # noqa: E501
        else:
            logger.info(f"Request: {request.method} {request.url.path}")

        # 处理请求
        response = await call_next(request)

        # 记录请求结束时间
        process_time = time.time() - start_time

        # 记录响应信息
        logger.info(
            f"Response: {request.method} {request.url.path} " f"status={response.status_code} time={process_time:.3f}s"
        )

        # 添加处理时间到响应头
        response.headers["X-Process-Time"] = str(process_time)

        return response


class RateLimitMiddleware(BaseHTTPMiddleware):
    """基于 IP 的简单限流中间件

    Args:
        app: ASGI 应用
        max_requests: 窗口期内最大请求数
        window_seconds: 时间窗口（秒）
        paths: 限流的路径列表，None 表示所有路径
    """

    def __init__(self, app, max_requests: int = 100, window_seconds: int = 60, paths: list[str] | None = None):
        super().__init__(app)
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.paths = set(paths) if paths else None
        self.requests: dict[str, list[float]] = {}

    async def dispatch(self, request: Request, call_next):
        # 如果配置了路径过滤，只对指定路径限流
        if self.paths is not None and request.url.path not in self.paths:
            return await call_next(request)

        # 获取客户端 IP
        client_ip = request.client.host

        current_time = time.time()

        # 清理过期记录
        if client_ip in self.requests:
            self.requests[client_ip] = [
                ts for ts in self.requests[client_ip] if current_time - ts < self.window_seconds
            ]

        # 计算请求次数
        request_count = len(self.requests.get(client_ip, []))

        if request_count >= self.max_requests:
            from fastapi.responses import JSONResponse

            return JSONResponse(
                status_code=429,
                content={
                    "error": {
                        "code": "RATE_LIMIT_EXCEEDED",
                        "message": (  # noqa: E501
                            f"Too many requests. Limit: " f"{self.max_requests} per {self.window_seconds}s"
                        ),
                    }
                },
            )

        # 记录请求
        if client_ip not in self.requests:
            self.requests[client_ip] = []
        self.requests[client_ip].append(current_time)

        # 处理请求
        response = await call_next(request)

        # 添加限流信息到响应头
        response.headers["X-RateLimit-Limit"] = str(self.max_requests)
        response.headers["X-RateLimit-Remaining"] = str(self.max_requests - len(self.requests[client_ip]))

        return response