mtsd/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 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.response import error_response

logger = logging.getLogger("sqxp")

_redis_client: Optional[redis.Redis] = None
_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_web_api(path: str) -> bool:
    return path.startswith(f"{settings.API_V1_PREFIX}/web/")


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]:
    global _redis_client
    if not settings.REDIS_HOST:
        return None
    if _redis_client is None:
        _redis_client = redis.Redis(
            host=settings.REDIS_HOST,
            port=settings.REDIS_PORT,
            db=settings.REDIS_DB,
            password=settings.REDIS_PASSWORD or None,
            socket_timeout=settings.REDIS_TIMEOUT_SEC,
            ssl=settings.REDIS_SSL_ENABLED,
            decode_responses=True,
        )
    return _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=2.0)
            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)


class ApiCryptoMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if not settings.API_ENCRYPT_ENABLED:
            return await call_next(request)
        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_web_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 not isinstance(payload, dict):
            payload = {}

        if override_method == "GET":
            request.scope["method"] = "GET"
            request.scope["query_string"] = urlencode(payload, doseq=True).encode("utf-8")
            body_for_route = b""
        else:
            body_for_route = json.dumps(payload, 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,
        )