11

app/core/aes_crypto.py

@@ -0,0 +1,69 @@
+"""AES-256-GCM 加解密，与 aiclient Tauri 端一致。"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+from typing import Any
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+ENVELOPE_VERSION = 1
+ENCRYPTED_HEADER = "x-aiclient-encrypted"
+
+
+def parse_aes_key(raw: str) -> bytes | None:
+    """解析 32 字节密钥：标准 Base64 或 64 位十六进制。"""
+    s = raw.strip()
+    if not s:
+        return None
+    try:
+        if len(s) == 64 and all(c in "0123456789abcdefABCDEF" for c in s):
+            key = bytes.fromhex(s)
+        else:
+            key = base64.b64decode(s)
+        if len(key) != 32:
+            return None
+        return key
+    except Exception:
+        return None
+
+
+def encrypt_bytes(plaintext: bytes, key: bytes) -> str:
+    nonce = os.urandom(12)
+    ciphertext = AESGCM(key).encrypt(nonce, plaintext, None)
+    return base64.b64encode(nonce + ciphertext).decode("ascii")
+
+
+def decrypt_payload(b64_payload: str, key: bytes) -> bytes:
+    raw = base64.b64decode(b64_payload)
+    if len(raw) < 13:
+        raise ValueError("密文过短")
+    nonce, ciphertext = raw[:12], raw[12:]
+    return AESGCM(key).decrypt(nonce, ciphertext, None)
+
+
+def wrap_encrypted(plaintext: bytes, key: bytes) -> bytes:
+    payload = encrypt_bytes(plaintext, key)
+    return json.dumps({"v": ENVELOPE_VERSION, "payload": payload}, ensure_ascii=False).encode(
+        "utf-8"
+    )
+
+
+def unwrap_encrypted(body: bytes, key: bytes) -> bytes:
+    obj: dict[str, Any] = json.loads(body)
+    if obj.get("v") != ENVELOPE_VERSION or "payload" not in obj:
+        raise ValueError("无效的加密信封")
+    payload = obj["payload"]
+    if not isinstance(payload, str):
+        raise ValueError("payload 应为字符串")
+    return decrypt_payload(payload, key)
+
+
+def is_encrypted_envelope(body: bytes) -> bool:
+    try:
+        obj = json.loads(body)
+        return isinstance(obj, dict) and "payload" in obj and obj.get("v") == ENVELOPE_VERSION
+    except Exception:
+        return False