From b280cc6f99726e41cd8bd9df65b907706d4707f1 Mon Sep 17 00:00:00 2001
From: ZZX9599 <536509593@qq.com>
Date: Mon, 15 Sep 2025 19:45:11 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81=E9=A3=8E?=
 =?UTF-8?q?=E6=A0=BC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 encryption/encrypt_decorator.py | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/encryption/encrypt_decorator.py b/encryption/encrypt_decorator.py
index 6607ad4..1a7df9d 100644
--- a/encryption/encrypt_decorator.py
+++ b/encryption/encrypt_decorator.py
@@ -1,32 +1,42 @@
+import datetime
 import json
 from functools import wraps
+from typing import Any
 
 from encryption.encryption import aes_encrypt
 from schema.response_schema import APIResponse
+# 假设SensitiveResponse等是Pydantic模型，需导入BaseModel
+from pydantic import BaseModel
 
 
 def encrypt_response(field: str = "data"):
-    """接口返回值加密装饰器：默认加密APIResponse的data字段"""
+    """接口返回值加密装饰器：正确序列化自定义对象为JSON"""
 
     def decorator(func):
-        @wraps(func)  # 保留原函数元信息（避免FastAPI路由异常）
+        @wraps(func)
         async def wrapper(*args, **kwargs):
-            # 执行原接口函数，获取原始响应（APIResponse对象）
             original_response: APIResponse = await func(*args, **kwargs)
+            field_value = getattr(original_response, field)
 
-            # 若需加密的字段为空，直接返回原响应（如注册接口data=None）
-            if not getattr(original_response, field):
+            if not field_value:
                 return original_response
 
-            # 复杂数据转JSON字符串（支持datetime、字典、列表等）
-            field_value = getattr(original_response, field)
-            field_value_json = json.dumps(field_value, default=str)  # 处理特殊类型
+            # 自定义JSON序列化函数：处理Pydantic模型和datetime
+            def json_default(obj: Any) -> Any:
+                # 处理Pydantic模型（转换为字典）
+                if isinstance(obj, BaseModel):
+                    return obj.model_dump()  # Pydantic v2用model_dump()，v1用dict()
+                # 处理datetime（转换为ISO格式字符串）
+                if isinstance(obj, datetime):
+                    return obj.isoformat()
+                # 其他无法序列化的类型，可根据需要扩展
+                return str(obj)  # 作为最后兜底
 
-            # AES加密并替换原字段
+            # 使用自定义序列化函数，确保生成标准JSON
+            field_value_json = json.dumps(field_value, default=json_default)
             encrypted_data = aes_encrypt(field_value_json)
             setattr(original_response, field, encrypted_data)
 
-            # 返回加密后的响应
             return original_response
 
         return wrapper