优化代码风格
This commit is contained in:
@ -1,32 +1,42 @@
|
|||||||
|
import datetime
|
||||||
import json
|
import json
|
||||||
from functools import wraps
|
from functools import wraps
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
from encryption.encryption import aes_encrypt
|
from encryption.encryption import aes_encrypt
|
||||||
from schema.response_schema import APIResponse
|
from schema.response_schema import APIResponse
|
||||||
|
# 假设SensitiveResponse等是Pydantic模型,需导入BaseModel
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
||||||
def encrypt_response(field: str = "data"):
|
def encrypt_response(field: str = "data"):
|
||||||
"""接口返回值加密装饰器:默认加密APIResponse的data字段"""
|
"""接口返回值加密装饰器:正确序列化自定义对象为JSON"""
|
||||||
|
|
||||||
def decorator(func):
|
def decorator(func):
|
||||||
@wraps(func) # 保留原函数元信息(避免FastAPI路由异常)
|
@wraps(func)
|
||||||
async def wrapper(*args, **kwargs):
|
async def wrapper(*args, **kwargs):
|
||||||
# 执行原接口函数,获取原始响应(APIResponse对象)
|
|
||||||
original_response: APIResponse = await func(*args, **kwargs)
|
original_response: APIResponse = await func(*args, **kwargs)
|
||||||
|
field_value = getattr(original_response, field)
|
||||||
|
|
||||||
# 若需加密的字段为空,直接返回原响应(如注册接口data=None)
|
if not field_value:
|
||||||
if not getattr(original_response, field):
|
|
||||||
return original_response
|
return original_response
|
||||||
|
|
||||||
# 复杂数据转JSON字符串(支持datetime、字典、列表等)
|
# 自定义JSON序列化函数:处理Pydantic模型和datetime
|
||||||
field_value = getattr(original_response, field)
|
def json_default(obj: Any) -> Any:
|
||||||
field_value_json = json.dumps(field_value, default=str) # 处理特殊类型
|
# 处理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)
|
encrypted_data = aes_encrypt(field_value_json)
|
||||||
setattr(original_response, field, encrypted_data)
|
setattr(original_response, field, encrypted_data)
|
||||||
|
|
||||||
# 返回加密后的响应
|
|
||||||
return original_response
|
return original_response
|
||||||
|
|
||||||
return wrapper
|
return wrapper
|
||||||
|
|||||||
Reference in New Issue
Block a user