buff3.0/backend-client/app/utils/api_context.py

"""
API接口上下文管理
用于在请求处理过程中传递明文数据（入参明文、出参明文）
使用 contextvars 确保线程安全
"""
import contextvars
from typing import Any, Dict, Optional

# 请求上下文变量
request_context: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
    "request_context", default=None
)


def set_request_context(data: Dict[str, Any]) -> None:
    """设置请求上下文"""
    request_context.set(data)


def get_request_context() -> Optional[Dict[str, Any]]:
    """获取请求上下文"""
    return request_context.get()


def clear_request_context() -> None:
    """清除请求上下文"""
    request_context.set(None)


class RequestContext:
    """请求上下文管理器"""

    def __init__(self, encrypted_request: str = None, encrypted_response: str = None):
        self.encrypted_request = encrypted_request
        self.encrypted_response = encrypted_response
        self.decrypted_request: Optional[Dict[str, Any]] = None
        self.decrypted_response: Optional[Dict[str, Any]] = None
        self._token = None

    def __enter__(self) -> "RequestContext":
        """进入上下文"""
        data = {
            "encrypted_request": self.encrypted_request,
            "encrypted_response": self.encrypted_response,
            "decrypted_request": self.decrypted_request,
            "decrypted_response": self.decrypted_response,
        }
        self._token = request_context.set(data)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文"""
        if self._token:
            request_context.reset(self._token)
        return False

    def set_decrypted_request(self, data: Dict[str, Any]) -> None:
        """设置解密后的请求数据"""
        self.decrypted_request = data
        # 更新上下文
        current = get_request_context()
        if current:
            current["decrypted_request"] = data
            request_context.set(current)

    def set_decrypted_response(self, data: Dict[str, Any]) -> None:
        """设置解密后的响应数据"""
        self.decrypted_response = data
        # 更新上下文
        current = get_request_context()
        if current:
            current["decrypted_response"] = data
            request_context.set(current)


def get_decrypted_request() -> Optional[Dict[str, Any]]:
    """获取解密后的请求数据"""
    ctx = get_request_context()
    return ctx.get("decrypted_request") if ctx else None


def get_decrypted_response() -> Optional[Dict[str, Any]]:
    """获取解密后的响应数据"""
    ctx = get_request_context()
    return ctx.get("decrypted_response") if ctx else None


def get_encrypted_request() -> Optional[str]:
    """获取加密的请求数据"""
    ctx = get_request_context()
    return ctx.get("encrypted_request") if ctx else None


def get_encrypted_response() -> Optional[str]:
    """获取加密的响应数据"""
    ctx = get_request_context()
    return ctx.get("encrypted_response") if ctx else None


# 便捷的日志记录函数
def log_decrypted_request(request_data: Dict[str, Any]) -> None:
    """记录解密后的请求数据（入参明文）"""
    from app.utils.api_logger import api_logger
    from app.utils.api_context import get_decrypted_request as get_req
    
    ctx = get_request_context()
    if ctx:
        endpoint = ctx.get("endpoint", "unknown")
        method = ctx.get("method", "unknown")
        user_id = ctx.get("user_id")
        api_logger.log_request_decrypted(endpoint, method, request_data, user_id)


def log_decrypted_response(response_data: Dict[str, Any]) -> None:
    """记录解密后的响应数据（出参明文）"""
    from app.utils.api_logger import api_logger
    from app.utils.api_context import get_decrypted_response as get_res
    
    ctx = get_request_context()
    if ctx:
        endpoint = ctx.get("endpoint", "unknown")
        method = ctx.get("method", "unknown")
        user_id = ctx.get("user_id")
        api_logger.log_response_ready(endpoint, method, response_data, user_id)


def log_complete(
    decrypted_request: Optional[Dict[str, Any]] = None,
    decrypted_response: Optional[Dict[str, Any]] = None,
    encrypted_response: Optional[str] = None,
    processing_time_ms: int = 0,
    status_code: int = 200,
    error_message: Optional[str] = None,
) -> None:
    """记录完整的API调用信息（包含明文和密文）"""
    from app.utils.api_logger import api_logger
    
    ctx = get_request_context()
    if ctx:
        endpoint = ctx.get("endpoint", "unknown")
        method = ctx.get("method", "unknown")
        user_id = ctx.get("user_id")
        client_ip = ctx.get("client_ip")
        encrypted_request = ctx.get("encrypted_request")
        
        api_logger.log_api_call(
            endpoint=endpoint,
            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,
        )