1

video-gen-api/app/middleware/request_encrypt.py

@@ -0,0 +1,102 @@
+import base64
+import json
+import logging
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from fastapi import 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:
+        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 Exception:
+            logger.exception("Request decryption failed")
+            return Response(
+                content=json.dumps({"detail": "解密失败"}, ensure_ascii=False),
+                status_code=400,
+                media_type="application/json",
+            )