优化代码风格

This commit is contained in:
ZZX9599
2025-09-15 18:24:33 +08:00
parent ced84e49bc
commit 3cb83b292e
4 changed files with 67 additions and 88 deletions

View File

@ -1,39 +1,33 @@
import json
from functools import wraps from functools import wraps
from fastapi import HTTPException
from schema.response_schema import APIResponse from schema.response_schema import APIResponse
from encryption.encryption import AESCipher from utils.encrypt_utils import aes_encrypt
def encrypt_response(func): def encrypt_response(field: str = "data"):
""" """接口返回值加密装饰器默认加密APIResponse的data字段"""
返回值加密装饰器:
- 仅对 APIResponse 的 data 字段加密code/message 不加密,便于前端判断基础状态)
- 若 data 为 None如注册接口不加密
"""
@wraps(func) # 保留原函数元信息(如 __name__、__doc__避免 FastAPI 路由异常) def decorator(func):
async def wrapper(*args, **kwargs): @wraps(func) # 保留原函数元信息避免FastAPI路由异常
try: async def wrapper(*args, **kwargs):
# 1. 执行原接口函数,获取返回值APIResponse 对象) # 执行原接口函数,获取原始响应APIResponse对象
response: APIResponse = await func(*args, **kwargs) original_response: APIResponse = await func(*args, **kwargs)
# 2. 仅当 data 不为 None 时加密 # 若需加密的字段为空直接返回原响应如注册接口data=None
if response.data is not None: if not getattr(original_response, field):
# 加密 data 字段(字典类型) return original_response
encrypted_result = AESCipher.encrypt(response.data)
# 替换原 data 为加密后的数据(包含密文和 IV
response.data = {
"is_encrypted": True, # 标记是否加密,便于前端处理
**encrypted_result
}
return response # 复杂数据转JSON字符串支持datetime、字典、列表等
field_value = getattr(original_response, field)
field_value_json = json.dumps(field_value, default=str) # 处理特殊类型
except Exception as e: # AES加密并替换原字段
# 加密过程异常时,返回 500 错误 encrypted_data = aes_encrypt(field_value_json)
raise HTTPException( setattr(original_response, field, encrypted_data)
status_code=500,
detail=f"返回值加密失败:{str(e)}"
) from e
return wrapper # 返回加密后的响应
return original_response
return wrapper
return decorator

View File

@ -1,72 +1,56 @@
import base64
import os import os
from dotenv import load_dotenv import base64
from Crypto.Cipher import AES from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes from fastapi import HTTPException
# 加载环境变量(从 .env 文件读取密钥 # 硬编码AES密钥32字节AES-256
load_dotenv() AES_SECRET_KEY = b"jr1vA6tfWMHOYi6UXw67UuO6fdak2rMa"
AES_BLOCK_SIZE = 16 # AES固定块大小
# 校验密钥长度确保符合AES规范
valid_key_lengths = [16, 24, 32]
if len(AES_SECRET_KEY) not in valid_key_lengths:
raise ValueError(
f"AES密钥长度必须为{valid_key_lengths}字节,当前为{len(AES_SECRET_KEY)}字节"
)
class AESCipher: def aes_encrypt(plaintext: str) -> dict:
"""AES-CBC 对称加密工具类""" """AES-CBC模式加密(返回密文+IV均为Base64编码"""
# 从环境变量获取密钥AES-256 需 32 字节密钥AES-128 需 16 字节) try:
SECRET_KEY = "jr1vA6tfWMHOYi6UXw67UuO6fdak2rMa".encode("utf-8") # 生成随机IV16字节
# AES 块大小固定为 16 字节 iv = os.urandom(AES_BLOCK_SIZE)
BLOCK_SIZE = 16
@classmethod # 创建加密器
def _validate_key(cls): cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv)
"""校验密钥长度AES-256 需 32 字节AES-128 需 16 字节)"""
if len(cls.SECRET_KEY) not in (16, 32):
raise ValueError("AES 密钥长度必须为 16 字节AES-128或 32 字节AES-256")
@classmethod # 明文填充并加密
def encrypt(cls, data: dict) -> dict: padded_plaintext = pad(plaintext.encode("utf-8"), AES_BLOCK_SIZE)
""" ciphertext = base64.b64encode(cipher.encrypt(padded_plaintext)).decode("utf-8")
加密函数:将字典类型的 data 加密 iv_base64 = base64.b64encode(iv).decode("utf-8")
返回:{encrypted_data: 加密后Base64字符串, iv: 16字节IV的Base64字符串}
"""
cls._validate_key()
# 1. 生成 16 字节随机 IV每次加密都生成新 IV无需保密但需和解密一致
iv = get_random_bytes(cls.BLOCK_SIZE)
# 2. 初始化 AES-CBC 加密器
cipher = AES.new(cls.SECRET_KEY, AES.MODE_CBC, iv)
# 3. 数据序列化字典转JSON字符串→ 编码为字节 → 填充PKCS7
data_str = str(data) # 若需更严谨,可使用 json.dumps(data, ensure_ascii=False)
data_bytes = data_str.encode("utf-8")
padded_data = pad(data_bytes, cls.BLOCK_SIZE, style="pkcs7")
# 4. 加密 → 转为 Base64 字符串(便于接口传输)
encrypted_bytes = cipher.encrypt(padded_data)
encrypted_data = base64.b64encode(encrypted_bytes).decode("utf-8")
iv_b64 = base64.b64encode(iv).decode("utf-8")
return { return {
"encrypted_data": encrypted_data, "ciphertext": ciphertext,
"iv": iv_b64 # IV 需随密文一起返回,供前端解密 "iv": iv_base64,
"algorithm": "AES-CBC"
} }
except Exception as e:
raise HTTPException(status_code=500, detail=f"AES加密失败{str(e)}") from e
@classmethod
def decrypt(cls, encrypted_data: str, iv_b64: str) -> dict:
"""
解密函数:将加密后的 Base64 字符串解密为字典
参数encrypted_data加密数据、iv_b64加密时的IV
"""
cls._validate_key()
# 1. 解码 Base64IV 和 密文) def aes_decrypt(ciphertext: str, iv: str) -> str:
iv = base64.b64decode(iv_b64) """AES-CBC模式解密"""
encrypted_bytes = base64.b64decode(encrypted_data) try:
# 解码Base64
ciphertext_bytes = base64.b64decode(ciphertext)
iv_bytes = base64.b64decode(iv)
# 2. 初始化 AES-CBC 解密器 # 创建解密器
cipher = AES.new(cls.SECRET_KEY, AES.MODE_CBC, iv) cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv_bytes)
# 3. 解密 → 去除填充 → 解码为字符串 → 转为字典(此处简化,实际可用 json.loads # 解密并去填充
decrypted_bytes = unpad(cipher.decrypt(encrypted_bytes), cls.BLOCK_SIZE, style="pkcs7") decrypted_bytes = unpad(cipher.decrypt(ciphertext_bytes), AES_BLOCK_SIZE)
decrypted_str = decrypted_bytes.decode("utf-8") return decrypted_bytes.decode("utf-8")
return eval(decrypted_str) # 生产环境建议用 json.loads避免 eval 安全风险 except Exception as e:
raise HTTPException(status_code=500, detail=f"AES解密失败{str(e)}") from e

View File

@ -3,6 +3,7 @@ from mysql.connector import Error as MySQLError
from typing import Optional from typing import Optional
from ds.db import db from ds.db import db
from encryption.encrypt_decorator import encrypt_response
from schema.sensitive_schema import ( from schema.sensitive_schema import (
SensitiveCreateRequest, SensitiveCreateRequest,
SensitiveUpdateRequest, SensitiveUpdateRequest,
@ -120,6 +121,7 @@ async def get_sensitive(
# 3. 获取敏感信息分页列表(重构:支持分页+关键词搜索) # 3. 获取敏感信息分页列表(重构:支持分页+关键词搜索)
# ------------------------------ # ------------------------------
@router.get("", response_model=APIResponse, summary="获取敏感信息分页列表(支持关键词搜索)") @router.get("", response_model=APIResponse, summary="获取敏感信息分页列表(支持关键词搜索)")
@encrypt_response()
async def get_sensitive_list( async def get_sensitive_list(
page: int = Query(1, ge=1, description="页码默认1最小1"), page: int = Query(1, ge=1, description="页码默认1最小1"),
page_size: int = Query(10, ge=1, le=100, description="每页条数默认101-100"), page_size: int = Query(10, ge=1, le=100, description="每页条数默认101-100"),

View File

@ -27,7 +27,6 @@ router = APIRouter(
# 1. 用户注册接口 # 1. 用户注册接口
# ------------------------------ # ------------------------------
@router.post("/register", response_model=APIResponse, summary="用户注册") @router.post("/register", response_model=APIResponse, summary="用户注册")
@encrypt_response
async def user_register(request: UserRegisterRequest): async def user_register(request: UserRegisterRequest):
""" """
用户注册: 用户注册: