shanxia2.0/backend/app/core/api_crypto_middleware.py

"""Web API 加解密中间件。"""
from __future__ import annotations

import asyncio
import json
import logging
import time
from typing import Any, Optional, Tuple
from urllib.parse import parse_qsl, urlencode

import redis.asyncio as redis
from jose import jwt
from redis import exceptions as redis_exceptions
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

from app.core import api_crypto
from app.core.config import settings
from app.core.redis_client import get_redis_client
from app.core.response import error_response

logger = logging.getLogger("sqxp")

_local_nonce_cache: dict[str, float] = {}


def _nonce_key(user_id: str, nonce: str) -> str:
    return f"api-enc:nonce:{user_id}:{nonce}"


def _should_skip_path(path: str) -> bool:
    excludes = [x.strip() for x in (settings.API_ENCRYPT_EXCLUDE_PATHS or "").split(",") if x.strip()]
    return path in excludes


def _is_crypto_api(path: str) -> bool:
    """仅对运营端与小程序业务 API 应用统一密文协议。"""
    prefix = settings.API_V1_PREFIX
    return path.startswith(f"{prefix}/web/") or path.startswith(f"{prefix}/mini/")


def _claims_from_auth(auth_header: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
    if not auth_header or not auth_header.startswith("Bearer "):
        return None, None
    token = auth_header[7:].strip()
    if not token:
        return None, None
    try:
        payload = jwt.decode(
            token,
            settings.SECRET_KEY,
            algorithms=[settings.ALGORITHM],
            options={"verify_signature": False, "verify_exp": False},
        )
        return payload.get("sub"), token
    except Exception:
        return None, token


def _get_redis_client() -> Optional[redis.Redis]:
    return get_redis_client()


async def _mark_nonce_once(user_id: str, nonce: str) -> bool:
    ttl_sec = max(30, int(settings.API_ENCRYPT_NONCE_TTL_SEC))
    key = _nonce_key(user_id, nonce)
    client = _get_redis_client()
    if client is not None:
        try:
            ok = await asyncio.wait_for(
                client.set(key, "1", ex=ttl_sec, nx=True),
                timeout=settings.REDIS_OP_TIMEOUT_SEC,
            )
            return bool(ok)
        except (asyncio.TimeoutError, redis_exceptions.RedisError):
            logger.warning("api crypto nonce redis unavailable, fallback to local memory")

    now = time.time()
    expire_before = now - ttl_sec
    stale_keys = [k for k, v in _local_nonce_cache.items() if v < expire_before]
    for stale in stale_keys:
        _local_nonce_cache.pop(stale, None)
    if key in _local_nonce_cache:
        return False
    _local_nonce_cache[key] = now
    return True


def _json_error_response(code: int, message: str) -> Response:
    body = error_response(code, message)
    return Response(content=json.dumps(body, ensure_ascii=False), media_type="application/json", status_code=200)


async def _coerce_post_method_override_to_get(request: Request, call_next) -> Response:
    """前端 Axios 在开启加密时会把 GET 改为 POST 并带 X-HTTP-Method-Override: GET。

    当后端未启用 API_ENCRYPT_ENABLED 时，原逻辑直接放行，路由仍按 POST 匹配，导致仅注册了 GET 的接口返回 405。
    此处在不解密的情况下，将此类请求还原为 GET（空或明文 JSON 体）；若仍为 enc=1 密文则提示配置不一致。
    """
    if request.method != "POST":
        return await call_next(request)
    override = (request.headers.get("x-http-method-override") or "").upper()
    if override != "GET":
        return await call_next(request)

    body_bytes = await request.body()
    req_json: dict = {}
    if body_bytes:
        try:
            raw = json.loads(body_bytes.decode("utf-8"))
            req_json = raw if isinstance(raw, dict) else {}
        except Exception:
            req_json = {}

    if int(req_json.get("enc", 0) or 0) == 1:
        return _json_error_response(
            400,
            "客户端已开启接口加密(enc=1)，但服务端未启用 API_ENCRYPT_ENABLED；"
            "请在后端开启加密并配置 API_ENCRYPT_FIXED_SECRET，或将前端 VITE_API_ENCRYPT_ENABLED 设为 false。",
        )

    scope = dict(request.scope)
    scope["method"] = "GET"
    payload = req_json if isinstance(req_json, dict) else {}
    q0 = (request.scope.get("query_string") or b"").decode("utf-8")
    pairs = dict(parse_qsl(q0, keep_blank_values=True))
    for k, v in payload.items():
        if isinstance(v, (dict, list)):
            continue
        pairs[k] = "" if v is None else str(v)
    scope["query_string"] = urlencode(pairs, doseq=True).encode("utf-8") if pairs else b""

    async def receive():
        return {"type": "http.request", "body": b"", "more_body": False}

    new_request = Request(scope, receive)
    return await call_next(new_request)


class ApiCryptoMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # 专用 multipart 路由在解析前限流；不得被 JSON/GET override 预读整份上传。
        if request.url.path == f"{settings.API_V1_PREFIX}/web/doclib/upload":
            return await call_next(request)
        if not settings.API_ENCRYPT_ENABLED:
            return await _coerce_post_method_override_to_get(request, call_next)
        if not settings.API_ENCRYPT_FIXED_SECRET:
            if settings.API_ENCRYPT_FAIL_OPEN:
                logger.warning("api crypto bypassed: API_ENCRYPT_FIXED_SECRET is empty")
                return await call_next(request)
            return _json_error_response(500, "服务端加密配置缺失")
        path = request.url.path
        if not _is_crypto_api(path) or _should_skip_path(path):
            return await call_next(request)

        content_type = (request.headers.get("content-type") or "").lower()
        if "multipart/form-data" in content_type:
            return await call_next(request)

        body_bytes = await request.body()
        req_json = {}
        if body_bytes:
            try:
                req_json = json.loads(body_bytes.decode("utf-8"))
                if not isinstance(req_json, dict):
                    req_json = {}
            except Exception:
                req_json = {}

        is_encrypted_req = int(req_json.get("enc", 0) or 0) == 1
        if not is_encrypted_req:
            logger.info("api crypto bypassed plaintext request: %s", path)
            return await call_next(request)

        user_id, _ = _claims_from_auth(request.headers.get("authorization"))
        if not user_id:
            return _json_error_response(401, "未授权或 token 无效")

        ts = req_json.get("ts")
        nonce = req_json.get("nonce")
        edata = req_json.get("edata")
        if ts is None or not nonce or not isinstance(edata, str):
            logger.warning("api crypto invalid envelope: path=%s", path)
            return _json_error_response(400, "加密报文格式错误")
        if not api_crypto.is_ts_in_window(int(ts), settings.API_ENCRYPT_TS_SKEW_SEC):
            logger.warning("api crypto ts expired: path=%s user_id=%s", path, user_id)
            return _json_error_response(400, "请求已过期，请重试")
        if not await _mark_nonce_once(str(user_id), str(nonce)):
            logger.warning("api crypto replay detected: path=%s user_id=%s nonce=%s", path, user_id, nonce)
            return _json_error_response(409, "重复请求")

        key = api_crypto.derive_key(
            fixed_secret=settings.API_ENCRYPT_FIXED_SECRET,
            user_id=str(user_id),
            ts=int(ts),
            nonce=str(nonce),
        )
        try:
            plain = api_crypto.decrypt_json(edata, key)
        except Exception as exc:
            logger.warning("api crypto decrypt failed: path=%s user_id=%s err=%s", path, user_id, exc)
            return _json_error_response(400, "密文解密失败")

        override_method = request.headers.get("x-http-method-override", "").upper()
        payload = plain.get("data", plain)

        if override_method == "GET":
            request.scope["method"] = "GET"
            query_dict = payload if isinstance(payload, dict) else {}
            request.scope["query_string"] = urlencode(query_dict, doseq=True).encode("utf-8")
            body_for_route = b""
        else:
            # JSON body 可能是对象（多数 CRUD）或数组（如批量轮播图 POST）
            if isinstance(payload, (dict, list)):
                body_for_route = json.dumps(payload, ensure_ascii=False).encode("utf-8")
            else:
                body_for_route = json.dumps({}, ensure_ascii=False).encode("utf-8")

        async def receive():
            return {"type": "http.request", "body": body_for_route, "more_body": False}

        request = Request(request.scope, receive)
        response = await call_next(request)

        content_type = (response.headers.get("content-type") or "").lower()
        is_json_response = "application/json" in content_type or response.media_type == "application/json"
        if not is_json_response:
            return response

        raw_resp = b""
        if hasattr(response, "body_iterator"):
            chunks: list[bytes] = []
            async for chunk in response.body_iterator:
                chunks.append(chunk)
            raw_resp = b"".join(chunks)
        elif getattr(response, "body", None) is not None:
            raw_resp = response.body
        if not raw_resp:
            return response

        try:
            resp_obj = json.loads(raw_resp.decode("utf-8"))
            if not isinstance(resp_obj, dict):
                return response
        except Exception:
            return response

        try:
            resp_ts = int(time.time() * 1000)
            resp_nonce = api_crypto.gen_nonce_hex()
            resp_key = api_crypto.derive_key(
                fixed_secret=settings.API_ENCRYPT_FIXED_SECRET,
                user_id=str(user_id),
                ts=resp_ts,
                nonce=resp_nonce,
            )
            resp_edata = api_crypto.encrypt_json(resp_obj, resp_key)
            wrapped = {
                "code": resp_obj.get("code", 200),
                "encrypted": True,
                "ts": resp_ts,
                "nonce": resp_nonce,
                "edata": resp_edata,
            }
        except Exception as exc:
            logger.warning("api crypto response encrypt failed: path=%s user_id=%s err=%s", path, user_id, exc)
            if settings.API_ENCRYPT_FAIL_OPEN:
                return response
            return _json_error_response(500, "响应加密失败")
        headers = {k: v for k, v in response.headers.items() if k.lower() != "content-length"}
        return Response(
            content=json.dumps(wrapped, ensure_ascii=False),
            status_code=response.status_code,
            headers=headers,
            media_type="application/json",
            background=response.background,
        )