sqxp/backend/app/core/api_request_log_middleware.py

"""统一记录 API 请求/响应出入参的中间件"""
import asyncio
import json
import logging
import time
from typing import Any, Optional, Tuple

from fastapi import HTTPException
from jose import jwt
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

from app.core.config import settings
from app.core.database import SessionLocal

logger = logging.getLogger("sqxp")

_SENSITIVE_KEYS = frozenset(
    {
        "password",
        "oldpassword",
        "newpassword",
        "encryptedlogin",
        "authorization",
        "token",
        "refreshtoken",
        "apikey",
    }
)

_OMIT_REQ_BYTES = 10 * 1024 * 1024


def _truncate(s: str, max_len: int) -> str:
    if not s:
        return s
    if len(s) <= max_len:
        return s
    return s[:max_len] + f"\n... [truncated, total {len(s)} chars]"


def _redact_obj(obj: Any) -> Any:
    if isinstance(obj, dict):
        out = {}
        for k, v in obj.items():
            lk = str(k).lower()
            if (
                lk in _SENSITIVE_KEYS
                or "password" in lk
                or lk in ("encryptedlogin", "refreshtoken")
            ):
                out[k] = "***"
            else:
                out[k] = _redact_obj(v)
        return out
    if isinstance(obj, list):
        return [_redact_obj(x) for x in obj[:200]] + (["..."] if len(obj) > 200 else [])
    return obj


def _format_body_bytes(raw: bytes, content_type: str, max_bytes: int) -> str:
    if not raw:
        return ""
    if len(raw) > _OMIT_REQ_BYTES:
        return f"[body omitted: {len(raw)} bytes]"
    if len(raw) > max_bytes:
        raw = raw[:max_bytes]
    ct = (content_type or "").lower()
    if "multipart/form-data" in ct:
        return "[multipart/form-data body omitted]"
    if "application/json" in ct or ct == "" or "text/" in ct:
        try:
            text = raw.decode("utf-8", errors="replace")
            try:
                parsed = json.loads(text)
                parsed = _redact_obj(parsed)
                return _truncate(
                    json.dumps(parsed, ensure_ascii=False),
                    settings.API_REQUEST_LOG_MAX_BODY,
                )
            except json.JSONDecodeError:
                return _truncate(text, settings.API_REQUEST_LOG_MAX_BODY)
        except Exception:
            return f"[decode error, {len(raw)} bytes]"
    return _truncate(raw.decode("utf-8", errors="replace"), settings.API_REQUEST_LOG_MAX_BODY)


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


def _should_skip(request: Request) -> bool:
    p = request.url.path
    if request.method == "OPTIONS":
        return True
    if p in ("/", "/health", "/docs", "/redoc", "/openapi.json"):
        return True
    if p.startswith("/docs") or p.startswith("/redoc"):
        return True
    return False


def _persist_log(
    *,
    http_method: str,
    path: str,
    query_string: Optional[str],
    request_body: Optional[str],
    response_body: Optional[str],
    status_code: int,
    duration_ms: int,
    user_id: Optional[str],
    username: Optional[str],
    role: Optional[str],
    ip_address: Optional[str],
    user_agent: Optional[str],
) -> None:
    from datetime import datetime

    from app.core.utils import generate_id
    from app.models.models import SysApiRequestLog

    db = SessionLocal()
    try:
        row = SysApiRequestLog(
            id=generate_id(),
            http_method=http_method,
            path=_truncate(path, 500),
            query_string=_truncate(query_string or "", 2000) or None,
            request_body=request_body,
            response_body=response_body,
            status_code=status_code,
            duration_ms=duration_ms,
            user_id=user_id,
            username=_truncate(username or "", 100) if username else None,
            role=role,
            ip_address=ip_address,
            user_agent=_truncate(user_agent or "", 500) if user_agent else None,
            is_del=0,
            create_at=datetime.now(),
            create_by=None,
            create_name=None,
        )
        db.add(row)
        db.commit()
    except Exception as e:
        logger.warning("api_request_log persist failed: %s", e)
        db.rollback()
    finally:
        db.close()


def _schedule_persist_log(**kwargs) -> None:
    """异步提交日志写入任务，不阻塞当前请求。"""
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        # 无事件循环时（极少数场景）回退为同步写入，保证日志不丢。
        try:
            _persist_log(**kwargs)
        except Exception as e:
            logger.warning("api_request_log persist fallback failed: %s", e)
        return

    task = loop.create_task(asyncio.to_thread(_persist_log, **kwargs))

    def _on_done(t: asyncio.Task) -> None:
        try:
            t.result()
        except Exception as e:
            logger.warning("api_request_log async task failed: %s", e)

    task.add_done_callback(_on_done)


class ApiRequestLogMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if not settings.API_REQUEST_LOG_ENABLED or _should_skip(request):
            return await call_next(request)

        max_b = settings.API_REQUEST_LOG_MAX_BODY_BYTES
        content_type = request.headers.get("content-type", "")

        try:
            body = await request.body()
        except Exception as e:
            body = b""
            logger.warning("api_request_log read body failed: %s", e)

        req_text = _format_body_bytes(body, content_type, max_b)

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

        request = Request(request.scope, receive)

        uid, uname, role = _claims_from_auth(request.headers.get("authorization"))
        ip = request.headers.get("x-forwarded-for", "").split(",")[0].strip()
        if not ip and request.client:
            ip = request.client.host
        ua = request.headers.get("user-agent", "")

        start = time.perf_counter()
        status_code = 500
        resp_text = ""
        raw_resp = b""
        out: Optional[Response] = None
        try:
            response = await call_next(request)
            status_code = response.status_code
            if hasattr(response, "body_iterator"):
                resp_chunks: list[bytes] = []
                async for chunk in response.body_iterator:
                    resp_chunks.append(chunk)
                raw_resp = b"".join(resp_chunks)
            elif getattr(response, "body", None) is not None:
                raw_resp = response.body
            ct_out = response.headers.get("content-type", "")
            resp_text = _format_body_bytes(raw_resp, ct_out, max_b)

            headers = {k: v for k, v in response.headers.items() if k.lower() != "content-length"}
            out = Response(
                content=raw_resp,
                status_code=response.status_code,
                headers=headers,
                media_type=response.media_type,
                background=response.background,
            )
        except HTTPException as e:
            status_code = e.status_code
            d = e.detail
            if not isinstance(d, str):
                try:
                    d = json.dumps(d, ensure_ascii=False)
                except Exception:
                    d = str(d)
            resp_text = _truncate(d, settings.API_REQUEST_LOG_MAX_BODY)
            raise
        except Exception as e:
            resp_text = _truncate(str(e), settings.API_REQUEST_LOG_MAX_BODY)
            raise
        finally:
            elapsed_ms = int((time.perf_counter() - start) * 1000)
            path = request.url.path
            qs = request.url.query
            _schedule_persist_log(
                http_method=request.method,
                path=path,
                query_string=qs or None,
                request_body=req_text or None,
                response_body=resp_text or None,
                status_code=status_code,
                duration_ms=elapsed_ms,
                user_id=uid,
                username=uname,
                role=role,
                ip_address=ip,
                user_agent=ua,
            )

        return out