FastapiAdmin/backend/app/core/middlewares.py

import json
import time
import uuid

from starlette.middleware.base import (
    BaseHTTPMiddleware,
    RequestResponseEndpoint,
)
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.gzip import GZipMiddleware
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp

from app.api.v1.module_system.params.service import ParamsService
from app.common.response import ErrorResponse
from app.config.setting import settings
from app.core.exceptions import CustomException
from app.core.logger import RequestContext, logger
from app.core.security import decode_access_token
from app.core.tenant import clear_current_tenant, set_current_tenant


class CustomCORSMiddleware(CORSMiddleware):
    """CORS跨域中间件"""

    def __init__(self, app: ASGIApp) -> None:
        super().__init__(
            app,
            allow_origins=settings.ALLOW_ORIGINS,
            allow_methods=settings.ALLOW_METHODS,
            allow_headers=settings.ALLOW_HEADERS,
            allow_credentials=settings.ALLOW_CREDENTIALS,
            expose_headers=settings.CORS_EXPOSE_HEADERS,
        )


class RequestLogMiddleware(BaseHTTPMiddleware):
    """
    记录请求日志中间件: 提供一个基础的中间件类，允许你自定义请求和响应处理逻辑。
    """

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

    @staticmethod
    def _extract_session_id_from_request(request: Request) -> str | None:
        """
        从请求中提取session_id（支持从Token或已设置的scope中获取）

        参数:
        - request (Request): 请求对象

        返回:
        - str | None: 会话ID，如果无法提取则返回None
        """
        # 1. 先检查 scope 中是否已经有 session_id（登录接口会设置）
        session_id = request.scope.get("session_id")
        if session_id:
            return session_id

        # 2. 检查 TenantMiddleware 缓存到 scope 的 JWT user_info（避免重复解码）
        user_info = request.scope.get("_jwt_user_info")
        if user_info:
            session_id = user_info.get("session_id")
            if session_id:
                request.scope["session_id"] = session_id
                return session_id

        # 3. 尝试从 Authorization Header 中提取
        try:
            authorization = request.headers.get("Authorization")
            if not authorization:
                return None

            # 处理Bearer token (大小写不敏感)
            token = authorization
            if authorization.lower().startswith("bearer "):
                token = authorization[7:].strip()
            elif authorization.lower().startswith("bearer"):
                token = authorization[6:].strip()

            # 解码token
            payload = decode_access_token(token)
            if not payload or not hasattr(payload, "sub"):
                return None

            # 从payload中提取session_id
            user_info = json.loads(payload.sub)
            session_id = user_info.get("session_id")

            # 同时设置到request.scope中，避免后续重复解析
            if session_id:
                request.scope["session_id"] = session_id

            return session_id
        except Exception:
            # 解析失败静默处理，返回None（可能是未认证请求）
            return None

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        """
        记录请求日志并透传响应。

        参数:
        - request (Request): 当前请求。
        - call_next (RequestResponseEndpoint): 下一层 ASGI 可调用对象。

        返回:
        - Response: 下游中间件/路由产生的响应。
        """
        start_time = time.time()

        # 尝试提取session_id
        session_id = self._extract_session_id_from_request(request)

        # 组装请求日志字段
        log_fields = (
            f"请求来源: {request.client.host if request.client else '未知'},"
            f"请求方法: {request.method},"
            f"请求路径: {request.url.path}"
        )
        logger.info(log_fields)

        try:
            # 初始化响应变量
            response = None

            # 获取请求路径
            path = request.scope.get("path")

            # 尝试获取客户端真实IP
            request_ip = None
            request_ip = (
                x_forwarded_for.split(",")[0].strip()
                if (x_forwarded_for := request.headers.get("X-Forwarded-For"))
                else request.client.host
                if request.client
                else None
            )
            # 检查是否启用演示模式
            demo_enable = False
            ip_white_list = []
            white_api_list_path = []
            ip_black_list = []

            try:
                # 从应用实例获取Redis连接
                redis = request.app.state.redis
                if not redis:
                    raise CustomException(msg="无法获取Redis连接")

                # 使用ParamsService获取系统配置
                system_config = await ParamsService.get_system_config_for_middleware(redis)
                # 提取配置值
                demo_enable = system_config["demo_enable"]
                ip_white_list = system_config["ip_white_list"]
                white_api_list_path = system_config["white_api_list_path"]
                ip_black_list = system_config["ip_black_list"]

            except Exception as e:
                logger.error(f"获取系统配置失败: {e}")

            # 检查是否需要拦截请求
            should_block = False
            block_reason = ""

            # 1. 首先检查IP是否在黑名单中
            if request_ip and request_ip in ip_black_list:
                should_block = True
                block_reason = f"IP地址 {request_ip} 在黑名单中"

            # 2. 如果不在黑名单中，检查是否在演示模式下需要拦截
            elif demo_enable in (True, "true", "True") and request.method != "GET":
                # 在演示模式下，非GET请求需要检查白名单
                is_ip_whitelisted = request_ip in ip_white_list
                is_path_whitelisted = path in white_api_list_path

                if not is_ip_whitelisted and not is_path_whitelisted:
                    should_block = True
                    block_reason = f"演示模式下拦截非GET请求，IP: {request_ip}, 路径: {path}"

            if should_block:
                # 增强安全审计：记录详细的拦截日志
                logger.warning(
                    " | ".join([
                        f"会话ID: {session_id or '未认证'}",
                        f"请求被拦截: {block_reason}",
                        f"请求来源: {request_ip}",
                        f"请求方法: {request.method}",
                        f"请求路径: {path}",
                        f"用户代理: {request.headers.get('user-agent', '未知')}",
                        f"演示模式: {demo_enable}",
                    ])
                )
                # 拦截请求
                return ErrorResponse(msg="演示环境，禁止操作")
            # 正常处理请求
            response = await call_next(request)

            # 计算处理时间并添加到响应头
            process_time = round(time.time() - start_time, 5)
            response.headers["X-Process-Time"] = str(process_time)

            # 构建响应日志信息
            response_info = f"响应状态: {response.status_code}, 处理时间: {round(process_time * 1000, 3)}ms"
            logger.info(response_info)

            return response

        except CustomException as e:
            logger.exception(f"中间件处理异常: {e!s}")
            return ErrorResponse(msg="系统异常，请联系管理员", data=str(e))


class CustomGZipMiddleware(GZipMiddleware):
    """GZip压缩中间件"""

    def __init__(self, app: ASGIApp) -> None:
        super().__init__(
            app,
            minimum_size=settings.GZIP_MIN_SIZE,
            compresslevel=settings.GZIP_COMPRESS_LEVEL,
        )


class SecurityHeadersMiddleware(BaseHTTPMiddleware):
    """安全响应头中间件"""

    def __init__(self, app: ASGIApp) -> None:
        super().__init__(app)
        from app.config.setting import settings

        self._hsts = settings.HSTS_ENABLE

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        response = await call_next(request)

        if self._hsts:
            response.headers["Strict-Transport-Security"] = (
                "max-age=31536000; includeSubDomains; preload"
            )

        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        response.headers["Permissions-Policy"] = (
            "camera=(), microphone=(), geolocation=(), "
            "payment=(), usb=(), magnetometer=(), "
            "gyroscope=()"
        )

        return response


class CorrelationIdMiddleware(BaseHTTPMiddleware):
    """请求关联 ID 中间件"""
    
    def __init__(self, app: ASGIApp) -> None:
        self.CORRELATION_ID_HEADER = "X-Correlation-ID"
        super().__init__(app)

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        correlation_id = request.headers.get(self.CORRELATION_ID_HEADER, "")
        if not correlation_id:
            correlation_id = str(uuid.uuid4())

        token = RequestContext.set_correlation_id(correlation_id)
        RequestContext.set_request_path(request.url.path)

        try:
            response = await call_next(request)
            response.headers[self.CORRELATION_ID_HEADER] = correlation_id
            return response
        finally:
            RequestContext._correlation_id.reset(token)


# ── 租户白名单 ──

_TENANT_WHITELIST_PREFIXES: tuple[str, ...] = (
    "/docs", "/redoc", "/ljdoc", "/openapi.json", "/metrics", "/static",
)

_TENANT_WHITELIST_PATHS: tuple[str, ...] = (
    "/api/v1/system/auth/login",
    "/api/v1/system/auth/captcha",
    "/api/v1/system/auth/refresh",
    "/api/v1/health",
    "/api/v1/common/health",
)


def _tenant_is_whitelisted(path: str) -> bool:
    for wp in _TENANT_WHITELIST_PATHS:
        if path.startswith(wp):
            return True
    for prefix in _TENANT_WHITELIST_PREFIXES:
        if path.startswith(prefix):
            return True
    for exclude_path in settings.TENANT_WHITELIST_PATHS:
        if path.startswith(exclude_path):
            return True
    return False


def _extract_tenant_from_token(request: Request) -> tuple[int | None, bool]:
    authorization = request.headers.get("Authorization")
    if not authorization:
        return None, False
    token = authorization
    if authorization.lower().startswith("bearer "):
        token = authorization[7:]
    elif authorization.lower().startswith("bearer"):
        token = authorization[6:]
    if not token.strip():
        return None, False
    try:
        payload = decode_access_token(token)
        if not payload or not hasattr(payload, "sub"):
            return None, False
        user_info = json.loads(payload.sub)
        # 缓存到 request.scope，避免后续中间件/依赖重复解码 JWT
        request.scope["_jwt_payload"] = payload
        request.scope["_jwt_user_info"] = user_info
        tenant_id = user_info.get("tenant_id")
        is_super_admin = user_info.get("is_super_admin", False)
        return tenant_id, is_super_admin
    except Exception:
        return None, False


class TenantMiddleware(BaseHTTPMiddleware):
    """租户上下文中间件"""

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

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

        if request.method == "OPTIONS":
            return await call_next(request)

        if _tenant_is_whitelisted(path):
            return await call_next(request)

        try:
            tenant_id, is_super_admin = _extract_tenant_from_token(request)
            set_current_tenant(tenant_id, is_super_admin)
            if tenant_id is not None:
                logger.debug(
                    "租户上下文已设置: tenant_id={}, is_super_admin={}, path={}",
                    tenant_id, is_super_admin, path,
                )
        except Exception:
            logger.exception("租户中间件处理异常: path={}", path)
        try:
            response = await call_next(request)
            return response
        finally:
            clear_current_tenant()