33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
import json
|
||
from functools import wraps
|
||
from schema.response_schema import APIResponse
|
||
from utils.encrypt_utils import aes_encrypt
|
||
|
||
|
||
def encrypt_response(field: str = "data"):
|
||
"""接口返回值加密装饰器：默认加密APIResponse的data字段"""
|
||
|
||
def decorator(func):
|
||
@wraps(func) # 保留原函数元信息（避免FastAPI路由异常）
|
||
async def wrapper(*args, **kwargs):
|
||
# 执行原接口函数，获取原始响应（APIResponse对象）
|
||
original_response: APIResponse = await func(*args, **kwargs)
|
||
|
||
# 若需加密的字段为空，直接返回原响应（如注册接口data=None）
|
||
if not getattr(original_response, field):
|
||
return original_response
|
||
|
||
# 复杂数据转JSON字符串（支持datetime、字典、列表等）
|
||
field_value = getattr(original_response, field)
|
||
field_value_json = json.dumps(field_value, default=str) # 处理特殊类型
|
||
|
||
# AES加密并替换原字段
|
||
encrypted_data = aes_encrypt(field_value_json)
|
||
setattr(original_response, field, encrypted_data)
|
||
|
||
# 返回加密后的响应
|
||
return original_response
|
||
|
||
return wrapper
|
||
|
||
return decorator |