39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
|
from functools import wraps
|
|||
|
|
from fastapi import HTTPException
|
|||
|
|
from schema.response_schema import APIResponse
|
|||
|
|
from encryption.encryption import AESCipher
|
|||
|
|
|
|||
|
|
|
|||
|
|
def encrypt_response(func):
|
|||
|
|
"""
|
|||
|
|
返回值加密装饰器：
|
|||
|
|
- 仅对 APIResponse 的 data 字段加密（code/message 不加密，便于前端判断基础状态）
|
|||
|
|
- 若 data 为 None（如注册接口），不加密
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
@wraps(func) # 保留原函数元信息（如 __name__、__doc__，避免 FastAPI 路由异常）
|
|||
|
|
async def wrapper(*args, **kwargs):
|
|||
|
|
try:
|
|||
|
|
# 1. 执行原接口函数，获取返回值（APIResponse 对象）
|
|||
|
|
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
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return response
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
# 加密过程异常时，返回 500 错误
|
|||
|
|
raise HTTPException(
|
|||
|
|
status_code=500,
|
|||
|
|
detail=f"返回值加密失败：{str(e)}"
|
|||
|
|
) from e
|
|||
|
|
|
|||
|
|
return wrapper
|