buff3.0/backend-client/app/middleware/enhanced_logging.py

"""
增强的API日志中间件
记录所有接口的入参密文、入参明文、出参明文、出参密文
日志文件按天滚动，每天一个日志文件
"""
import time
import json
from typing import Any, Callable, Dict, Optional
from fastapi import Request
from starlette.types import ASGIApp
from app.utils.api_context import (
    RequestContext,
    set_request_context,
    get_decrypted_request,
    get_decrypted_response,
    get_encrypted_request,
    get_encrypted_response,
)
from app.utils.api_logger import api_logger


class EnhancedAPILoggingMiddleware:
    """增强的API日志中间件"""

    def __init__(self, app: ASGIApp):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return

        start_time = time.time()
        method = scope.get("method")
        server = scope.get("server") or ("localhost", 80)
        path = scope.get("raw_path", b"").decode("utf-8", errors="ignore")
        query_string = scope.get("query_string", b"").decode("utf-8", errors="ignore")
        scheme = scope.get("scheme", "http")
        url = f"{scheme}://{server[0]}:{server[1]}{path}" + (f"?{query_string}" if query_string else "")
        client = scope.get("client")
        client_ip = client[0] if client else None
        headers = {k.decode("utf-8"): v.decode("utf-8") for k, v in scope.get("headers", [])}
        user_agent = headers.get("user-agent", "")

        # 获取用户ID
        user_id = None
        try:
            header_user_id = headers.get("x-user-id") or headers.get("X-User-Id")
            if header_user_id:
                user_id = int(header_user_id)
        except (ValueError, TypeError):
            user_id = None

        # 收集请求体（仅在有 body 时）
        body_bytes = b""
        async def receive_wrapper():
            nonlocal body_bytes
            message = await receive()
            if message.get("type") == "http.request":
                body_bytes += message.get("body", b"")
            return message

        # 收集响应状态与 body 分片
        status_code = 500
        resp_chunks = []
        async def send_wrapper(message):
            nonlocal status_code, resp_chunks
            if message.get("type") == "http.response.start":
                status_code = message.get("status", 0)
            if message.get("type") == "http.response.body":
                chunk = message.get("body", b"")
                if chunk:
                    resp_chunks.append(chunk)
            await send(message)

        # 加密的请求数据
        encrypted_request = None
        if body_bytes:
            try:
                encrypted_request = body_bytes.decode("utf-8", errors="ignore")
            except Exception:
                encrypted_request = str(body_bytes)

        # 创建请求上下文
        ctx = RequestContext(
            encrypted_request=encrypted_request,
            encrypted_response=None,
        )

        error_message = None
        try:
            with ctx:
                # 设置请求上下文
                set_request_context({
                    "encrypted_request": encrypted_request,
                    "encrypted_response": None,
                    "decrypted_request": None,
                    "decrypted_response": None,
                })

                # 调用下游应用
                await self.app(scope, receive_wrapper, send_wrapper)

        except Exception as e:
            error_message = str(e)
            raise

        finally:
            try:
                # 收集响应数据
                response_body_bytes = b"".join(resp_chunks) if resp_chunks else b""
                encrypted_response = None
                if response_body_bytes:
                    try:
                        encrypted_response = response_body_bytes.decode("utf-8", errors="ignore")
                    except Exception:
                        encrypted_response = str(response_body_bytes)

                # 计算执行时间
                processing_time_ms = int((time.time() - start_time) * 1000)

                # 获取解密后的请求和响应数据
                decrypted_request = get_decrypted_request()
                decrypted_response = get_decrypted_response()

                # 记录完整的API调用日志
                api_logger.log_api_call(
                    endpoint=url,
                    method=method,
                    encrypted_request=encrypted_request,
                    decrypted_request=decrypted_request,
                    encrypted_response=encrypted_response,
                    decrypted_response=decrypted_response,
                    processing_time_ms=processing_time_ms,
                    status_code=status_code,
                    user_id=user_id,
                    client_ip=client_ip,
                    error_message=error_message,
                )

            except Exception as e:
                # 日志记录失败不应影响响应
                print(f"记录API日志失败: {e}")


def setup_enhanced_api_logging(app: Callable) -> None:
    """
    为FastAPI应用设置增强的API日志中间件
    在应用创建后调用此函数来注册中间件
    """
    from fastapi import FastAPI
    if isinstance(app, FastAPI):
        app.add_middleware(EnhancedAPILoggingMiddleware)