内容安全审核

encryption/encrypt_decorator.py

@@ -0,0 +1,40 @@
+import json
+from datetime import datetime
+from functools import wraps
+from typing import Any
+
+from encryption.encryption import aes_encrypt
+from schema.response_schema import APIResponse
+from pydantic import BaseModel
+
+
+def encrypt_response(field: str = "data"):
+    """接口返回值加密装饰器：正确序列化自定义对象为JSON"""
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            original_response: APIResponse = await func(*args, **kwargs)
+            field_value = getattr(original_response, field)
+
+            if not field_value:
+                return original_response
+
+            # 自定义JSON序列化函数：处理Pydantic模型和datetime
+            def json_default(obj: Any) -> Any:
+                if isinstance(obj, BaseModel):
+                    return obj.model_dump() 
+                if isinstance(obj, datetime):
+                    return obj.isoformat()
+                return str(obj) 
+
+            # 使用自定义序列化函数、确保生成标准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
+
+    return decorator