70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
"""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
|