dpb/backend/app/core/middlewares.py

import json
import uuid
from typing import Any

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.gzip import GZipMiddleware
from starlette.middleware.trustedhost import TrustedHostMiddleware
from starlette.requests import Request
from starlette.responses import RedirectResponse, Response
from starlette.types import ASGIApp

from app.common.enums import RedisInitKeyConfig, SysParamKey
from app.common.response import ErrorResponse
from app.config.setting import settings
from app.core.exceptions import CustomException
from app.core.logger import logger, reset_correlation_id, set_correlation_id
from app.core.redis_crud import RedisCURD
from app.utils.ip_local_util import get_client_ip


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):
    """演示模式 & IP黑名单拦截"""

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

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

        try:
            path = request.url.path
            config = await self._load_config(request)
            is_blacklisted = bool(client_ip and client_ip in config[SysParamKey.IP_BLACK_LIST])
            in_demo = (
                config.get(SysParamKey.DEMO_ENABLE, False)
                and request.method != "GET"
                and (client_ip is None or client_ip not in config.get(SysParamKey.IP_WHITE_LIST, ()))
                and not any(
                    path.startswith(item.rstrip("*")) if item.endswith("*") else path == item
                    for item in settings.WHITE_API_LIST_PATH
                    if isinstance(item, str) and item
                )
            )

            if is_blacklisted or in_demo:
                logger.warning(
                    "请求被拦截: {} {} | ip={} | 原因={}",
                    request.method,
                    path,
                    client_ip,
                    "IP黑名单" if is_blacklisted else "演示模式",
                )
                return ErrorResponse(msg="IP已被黑名单" if is_blacklisted else "演示环境，禁止操作")

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

    @staticmethod
    async def _load_config(request: Request) -> dict:
        """加载中间件配置，失败时返回全部默认值。"""
        redis = getattr(request.app.state, "redis", None)
        if not redis:
            return {SysParamKey.DEMO_ENABLE: False, SysParamKey.IP_WHITE_LIST: (), SysParamKey.IP_BLACK_LIST: ()}
        try:
            config_keys = [
                f"{RedisInitKeyConfig.SYSTEM_CONFIG.key}:{SysParamKey.DEMO_ENABLE.value}",
                f"{RedisInitKeyConfig.SYSTEM_CONFIG.key}:{SysParamKey.IP_WHITE_LIST.value}",
                f"{RedisInitKeyConfig.SYSTEM_CONFIG.key}:{SysParamKey.IP_BLACK_LIST.value}",
            ]
            config_values = await RedisCURD(redis).mget(config_keys)
            result: dict[str, Any] = {SysParamKey.DEMO_ENABLE: False, SysParamKey.IP_WHITE_LIST: (), SysParamKey.IP_BLACK_LIST: ()}
            raw_demo, raw_white, raw_black = config_values
            for raw, key in ((raw_demo, SysParamKey.DEMO_ENABLE), (raw_white, SysParamKey.IP_WHITE_LIST), (raw_black, SysParamKey.IP_BLACK_LIST)):
                if not raw:
                    continue
                try:
                    payload = json.loads(raw)
                except json.JSONDecodeError:
                    logger.error("解析系统配置 %s 失败", key)
                    continue
                if not isinstance(payload, dict) or payload.get("status", 0) != 0:
                    continue
                cv = payload.get("config_value")
                if cv is None:
                    continue
                if key == SysParamKey.DEMO_ENABLE:
                    result[key.value] = cv in (True, "true", "1", "yes", "on")
                else:
                    result[key.value] = json.loads(cv) if isinstance(cv, str) else cv
            return result
        except Exception:
            return {SysParamKey.DEMO_ENABLE: False, SysParamKey.IP_WHITE_LIST: (), SysParamKey.IP_BLACK_LIST: ()}


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 CustomHTTPSRedirectMiddleware(BaseHTTPMiddleware):
    """HTTP → HTTPS 重定向中间件（信任前端代理的 X-Forwarded-Proto 头）"""

    async def dispatch(self, request: Request, call_next):
        if request.url.scheme != "https" and request.headers.get("X-Forwarded-Proto") != "https":
            url = request.url.replace(scheme="https")
            return RedirectResponse(url, status_code=301)
        return await call_next(request)


class CustomTrustedHostMiddleware(TrustedHostMiddleware):
    """可信主机 Host 头校验中间件"""

    def __init__(self, app: ASGIApp) -> None:
        super().__init__(app, allowed_hosts=settings.ALLOWED_HOSTS)


class CorrelationIdMiddleware(BaseHTTPMiddleware):
    """请求 ID 中间件"""

    def __init__(self, app: ASGIApp) -> None:
        self._header = "X-Correlation-ID"
        super().__init__(app)

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        cid = request.headers.get(self._header) or str(uuid.uuid4())
        token = set_correlation_id(cid)
        try:
            response = await call_next(request)
            response.headers[self._header] = cid
            return response
        finally:
            reset_correlation_id(token)