api_context.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. """
  2. API接口上下文管理
  3. 用于在请求处理过程中传递明文数据(入参明文、出参明文)
  4. 使用 contextvars 确保线程安全
  5. """
  6. import contextvars
  7. from typing import Any, Dict, Optional
  8. # 请求上下文变量
  9. request_context: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
  10. "request_context", default=None
  11. )
  12. def set_request_context(data: Dict[str, Any]) -> None:
  13. """设置请求上下文"""
  14. request_context.set(data)
  15. def get_request_context() -> Optional[Dict[str, Any]]:
  16. """获取请求上下文"""
  17. return request_context.get()
  18. def clear_request_context() -> None:
  19. """清除请求上下文"""
  20. request_context.set(None)
  21. class RequestContext:
  22. """请求上下文管理器"""
  23. def __init__(self, encrypted_request: str = None, encrypted_response: str = None):
  24. self.encrypted_request = encrypted_request
  25. self.encrypted_response = encrypted_response
  26. self.decrypted_request: Optional[Dict[str, Any]] = None
  27. self.decrypted_response: Optional[Dict[str, Any]] = None
  28. self._token = None
  29. def __enter__(self) -> "RequestContext":
  30. """进入上下文"""
  31. data = {
  32. "encrypted_request": self.encrypted_request,
  33. "encrypted_response": self.encrypted_response,
  34. "decrypted_request": self.decrypted_request,
  35. "decrypted_response": self.decrypted_response,
  36. }
  37. self._token = request_context.set(data)
  38. return self
  39. def __exit__(self, exc_type, exc_val, exc_tb):
  40. """退出上下文"""
  41. if self._token:
  42. request_context.reset(self._token)
  43. return False
  44. def set_decrypted_request(self, data: Dict[str, Any]) -> None:
  45. """设置解密后的请求数据"""
  46. self.decrypted_request = data
  47. # 更新上下文
  48. current = get_request_context()
  49. if current:
  50. current["decrypted_request"] = data
  51. request_context.set(current)
  52. def set_decrypted_response(self, data: Dict[str, Any]) -> None:
  53. """设置解密后的响应数据"""
  54. self.decrypted_response = data
  55. # 更新上下文
  56. current = get_request_context()
  57. if current:
  58. current["decrypted_response"] = data
  59. request_context.set(current)
  60. def get_decrypted_request() -> Optional[Dict[str, Any]]:
  61. """获取解密后的请求数据"""
  62. ctx = get_request_context()
  63. return ctx.get("decrypted_request") if ctx else None
  64. def get_decrypted_response() -> Optional[Dict[str, Any]]:
  65. """获取解密后的响应数据"""
  66. ctx = get_request_context()
  67. return ctx.get("decrypted_response") if ctx else None
  68. def get_encrypted_request() -> Optional[str]:
  69. """获取加密的请求数据"""
  70. ctx = get_request_context()
  71. return ctx.get("encrypted_request") if ctx else None
  72. def get_encrypted_response() -> Optional[str]:
  73. """获取加密的响应数据"""
  74. ctx = get_request_context()
  75. return ctx.get("encrypted_response") if ctx else None
  76. # 便捷的日志记录函数
  77. def log_decrypted_request(request_data: Dict[str, Any]) -> None:
  78. """记录解密后的请求数据(入参明文)"""
  79. from app.utils.api_logger import api_logger
  80. from app.utils.api_context import get_decrypted_request as get_req
  81. ctx = get_request_context()
  82. if ctx:
  83. endpoint = ctx.get("endpoint", "unknown")
  84. method = ctx.get("method", "unknown")
  85. user_id = ctx.get("user_id")
  86. api_logger.log_request_decrypted(endpoint, method, request_data, user_id)
  87. def log_decrypted_response(response_data: Dict[str, Any]) -> None:
  88. """记录解密后的响应数据(出参明文)"""
  89. from app.utils.api_logger import api_logger
  90. from app.utils.api_context import get_decrypted_response as get_res
  91. ctx = get_request_context()
  92. if ctx:
  93. endpoint = ctx.get("endpoint", "unknown")
  94. method = ctx.get("method", "unknown")
  95. user_id = ctx.get("user_id")
  96. api_logger.log_response_ready(endpoint, method, response_data, user_id)
  97. def log_complete(
  98. decrypted_request: Optional[Dict[str, Any]] = None,
  99. decrypted_response: Optional[Dict[str, Any]] = None,
  100. encrypted_response: Optional[str] = None,
  101. processing_time_ms: int = 0,
  102. status_code: int = 200,
  103. error_message: Optional[str] = None,
  104. ) -> None:
  105. """记录完整的API调用信息(包含明文和密文)"""
  106. from app.utils.api_logger import api_logger
  107. ctx = get_request_context()
  108. if ctx:
  109. endpoint = ctx.get("endpoint", "unknown")
  110. method = ctx.get("method", "unknown")
  111. user_id = ctx.get("user_id")
  112. client_ip = ctx.get("client_ip")
  113. encrypted_request = ctx.get("encrypted_request")
  114. api_logger.log_api_call(
  115. endpoint=endpoint,
  116. method=method,
  117. encrypted_request=encrypted_request,
  118. decrypted_request=decrypted_request,
  119. encrypted_response=encrypted_response,
  120. decrypted_response=decrypted_response,
  121. processing_time_ms=processing_time_ms,
  122. status_code=status_code,
  123. user_id=user_id,
  124. client_ip=client_ip,
  125. error_message=error_message,
  126. )