107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
# -*- 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:]}"
|