116 lines
4.3 KiB
Python
116 lines
4.3 KiB
Python
import base64
|
|
import json
|
|
import logging
|
|
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
from fastapi import HTTPException, Request, Response
|
|
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
|
|
|
from app.config import settings
|
|
|
|
logger = logging.getLogger("videogen")
|
|
|
|
|
|
def _get_aesgcm() -> AESGCM:
|
|
"""Get AESGCM instance from configured key."""
|
|
key_bytes = base64.b64decode(settings.ENCRYPTION_KEY)
|
|
if len(key_bytes) != 32:
|
|
# Pad or truncate to 32 bytes
|
|
key_bytes = (key_bytes + b"\x00" * 32)[:32]
|
|
return AESGCM(key_bytes)
|
|
|
|
|
|
def encrypt_data(plaintext: bytes) -> bytes:
|
|
"""Encrypt data using AES-256-GCM. Returns nonce + ciphertext."""
|
|
import os
|
|
aesgcm = _get_aesgcm()
|
|
nonce = os.urandom(12)
|
|
ciphertext = aesgcm.encrypt(nonce, plaintext, None)
|
|
return base64.b64encode(nonce + ciphertext)
|
|
|
|
|
|
def decrypt_data(data: bytes) -> bytes:
|
|
"""Decrypt AES-256-GCM encrypted data."""
|
|
raw = base64.b64decode(data)
|
|
nonce = raw[:12]
|
|
ciphertext = raw[12:]
|
|
aesgcm = _get_aesgcm()
|
|
return aesgcm.decrypt(nonce, ciphertext, None)
|
|
|
|
|
|
class RequestEncryptMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(
|
|
self, request: Request, call_next: RequestResponseEndpoint
|
|
) -> Response:
|
|
# 白名单：支付回调接口不需要加密/解密
|
|
path = request.url.path
|
|
if "/payments/alipay/callback" in path or "/payments/wechat/callback" in path:
|
|
return await call_next(request)
|
|
|
|
encrypted = request.headers.get("X-Encrypted", "").lower() == "true"
|
|
if not encrypted:
|
|
return await call_next(request)
|
|
|
|
# GET requests have no body to decrypt — just pass through
|
|
if request.method == "GET":
|
|
response = await call_next(request)
|
|
response_body = b""
|
|
async for chunk in response.body_iterator:
|
|
if isinstance(chunk, str):
|
|
response_body += chunk.encode()
|
|
else:
|
|
response_body += chunk
|
|
encrypted_response = encrypt_data(response_body)
|
|
return Response(
|
|
content=json.dumps({"data": encrypted_response.decode()}),
|
|
status_code=response.status_code,
|
|
headers={"X-Encrypted": "true", "Content-Type": "application/json"},
|
|
)
|
|
|
|
try:
|
|
body = await request.body()
|
|
# Frontend sends {"data": "<encrypted_base64>"}
|
|
body_json = json.loads(body)
|
|
encrypted_b64 = body_json.get("data", "")
|
|
decrypted = decrypt_data(encrypted_b64.encode())
|
|
|
|
# Replace request body with decrypted content
|
|
async def receive():
|
|
return {"type": "http.request", "body": decrypted, "more_body": False}
|
|
|
|
request._receive = receive
|
|
request._body = decrypted
|
|
|
|
response = await call_next(request)
|
|
|
|
# Encrypt response
|
|
response_body = b""
|
|
async for chunk in response.body_iterator:
|
|
if isinstance(chunk, str):
|
|
response_body += chunk.encode()
|
|
else:
|
|
response_body += chunk
|
|
|
|
encrypted_response = encrypt_data(response_body)
|
|
# Wrap in {"data": "..."} to match frontend's expected format
|
|
return Response(
|
|
content=json.dumps({"data": encrypted_response.decode()}),
|
|
status_code=response.status_code,
|
|
headers={"X-Encrypted": "true", "Content-Type": "application/json"},
|
|
)
|
|
except HTTPException as http_exc:
|
|
# BaseHTTPMiddleware 中 call_next 抛出的 HTTPException 会直接传播，
|
|
# 这里捕获后返回对应状态码的响应，避免被外层 except Exception 吞掉
|
|
return Response(
|
|
content=json.dumps({"detail": http_exc.detail}, ensure_ascii=False),
|
|
status_code=http_exc.status_code,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
except Exception:
|
|
logger.exception("Request decryption failed")
|
|
return Response(
|
|
content=json.dumps({"detail": "解密失败"}, ensure_ascii=False),
|
|
status_code=400,
|
|
media_type="application/json",
|
|
)
|