优化代码风格

encryption/encrypt_decorator.py

@@ -1,39 +1,33 @@
+import json
 from functools import wraps
-from fastapi import HTTPException
 from schema.response_schema import APIResponse
-from encryption.encryption import AESCipher
+from utils.encrypt_utils import aes_encrypt
 
 
-def encrypt_response(func):
-    """
-    返回值加密装饰器：
-    - 仅对 APIResponse 的 data 字段加密（code/message 不加密，便于前端判断基础状态）
-    - 若 data 为 None（如注册接口），不加密
-    """
+def encrypt_response(field: str = "data"):
+    """接口返回值加密装饰器：默认加密APIResponse的data字段"""
 
-    @wraps(func)  # 保留原函数元信息（如 __name__、__doc__，避免 FastAPI 路由异常）
-    async def wrapper(*args, **kwargs):
-        try:
-            # 1. 执行原接口函数，获取返回值（APIResponse 对象）
-            response: APIResponse = await func(*args, **kwargs)
+    def decorator(func):
+        @wraps(func)  # 保留原函数元信息（避免FastAPI路由异常）
+        async def wrapper(*args, **kwargs):
+            # 执行原接口函数，获取原始响应（APIResponse对象）
+            original_response: APIResponse = await func(*args, **kwargs)
 
-            # 2. 仅当 data 不为 None 时加密
-            if response.data is not None:
-                # 加密 data 字段（字典类型）
-                encrypted_result = AESCipher.encrypt(response.data)
-                # 替换原 data 为加密后的数据（包含密文和 IV）
-                response.data = {
-                    "is_encrypted": True,  # 标记是否加密，便于前端处理
-                    **encrypted_result
-                }
+            # 若需加密的字段为空，直接返回原响应（如注册接口data=None）
+            if not getattr(original_response, field):
+                return original_response
 
-            return response
+            # 复杂数据转JSON字符串（支持datetime、字典、列表等）
+            field_value = getattr(original_response, field)
+            field_value_json = json.dumps(field_value, default=str)  # 处理特殊类型
 
-        except Exception as e:
-            # 加密过程异常时，返回 500 错误
-            raise HTTPException(
-                status_code=500,
-                detail=f"返回值加密失败：{str(e)}"
-            ) from e
+            # AES加密并替换原字段
+            encrypted_data = aes_encrypt(field_value_json)
+            setattr(original_response, field, encrypted_data)
 
-    return wrapper
+            # 返回加密后的响应
+            return original_response
+
+        return wrapper
+
+    return decorator