gengx

wechat_rpa/secret_box.py

@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+"""模型密钥的落库加密。
+
+模型清单里存的是各家厂商的 API Key。这些密钥现在明文躺在每台客户端的
+ai_settings.json 里——收进后端本来就是为了解决这个问题，如果后端再明文存一遍，
+等于把泄漏面从"很多台客户端"换成"一台服务器 + 一个数据库备份文件"，并没有真的
+变安全。
+
+主密钥单独放一个 0600 的文件，不进数据库、不进代码库、不进备份脚本的默认范围。
+数据库整个被拖走也解不开密钥。
+
+用 AES-GCM：带认证标签，密文被改过会解密失败而不是悄悄返回一串垃圾——那种
+"垃圾密钥"会变成一整天查不出来的 401。
+"""
+
+from __future__ import annotations
+
+import base64
+import os
+import secrets
+from pathlib import Path
+
+_KEY_ENV = "WECOM_BACKEND_SECRET_KEY"
+_KEY_FILENAME = "backend_secret.key"
+
+
+class SecretBoxError(RuntimeError):
+    pass
+
+
+def _key_path(data_dir: Path | str) -> Path:
+    return Path(data_dir) / _KEY_FILENAME
+
+
+def load_or_create_key(data_dir: Path | str) -> bytes:
+    """取主密钥。环境变量优先，方便容器部署不落盘。"""
+    from_env = os.environ.get(_KEY_ENV, "").strip()
+    if from_env:
+        try:
+            key = base64.urlsafe_b64decode(from_env + "=" * (-len(from_env) % 4))
+        except Exception as exc:
+            raise SecretBoxError(f"{_KEY_ENV} 不是合法的 base64") from exc
+        if len(key) != 32:
+            raise SecretBoxError(f"{_KEY_ENV} 解出来不是 32 字节")
+        return key
+
+    path = _key_path(data_dir)
+    if path.exists():
+        key = path.read_bytes().strip()
+        key = base64.urlsafe_b64decode(key + b"=" * (-len(key) % 4))
+        if len(key) != 32:
+            raise SecretBoxError(f"{path} 里的主密钥长度不对，拒绝继续")
+        return key
+
+    key = secrets.token_bytes(32)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    # 先按 0600 建文件再写内容，避免密钥在世界可读的瞬间落盘
+    handle = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    try:
+        os.write(handle, base64.urlsafe_b64encode(key))
+    finally:
+        os.close(handle)
+    try:
+        os.chmod(str(path), 0o600)
+    except OSError:
+        pass
+    return key
+
+
+def encrypt(plaintext: str, key: bytes) -> str:
+    """→ base64(nonce | tag | 密文)。空串原样返回空串。"""
+    text = str(plaintext or "")
+    if not text:
+        return ""
+    from Crypto.Cipher import AES
+
+    nonce = secrets.token_bytes(12)
+    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
+    body, tag = cipher.encrypt_and_digest(text.encode("utf-8"))
+    return base64.urlsafe_b64encode(nonce + tag + body).decode("ascii")
+
+
+def decrypt(token: str, key: bytes) -> str:
+    """解不开就抛异常，绝不返回半截垃圾——那会变成查不出来的 401。"""
+    raw = str(token or "")
+    if not raw:
+        return ""
+    from Crypto.Cipher import AES
+
+    try:
+        blob = base64.urlsafe_b64decode(raw + "=" * (-len(raw) % 4))
+        nonce, tag, body = blob[:12], blob[12:28], blob[28:]
+        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
+        return cipher.decrypt_and_verify(body, tag).decode("utf-8")
+    except Exception as exc:
+        raise SecretBoxError("模型密钥解密失败：主密钥不匹配或密文被改动") from exc
+
+
+def masked(plaintext: str) -> str:
+    """给界面看的遮罩。永远不回显完整密钥。"""
+    text = str(plaintext or "")
+    if not text:
+        return ""
+    if len(text) <= 8:
+        return "*" * len(text)
+    return f"{text[:4]}{'*' * 6}{text[-4:]}"