Kris
a2201d8be3
新增完整的woc-bridge业务服务，实现微信API代理、密钥自动提取、消息队列限流等功能： 1. 新增bridge核心代码与s6服务配置 2. 新增ptrace权限初始化脚本与docker配置 3. 新增M2M鉴权API代理与环境变量配置 4. 新增SQLCipher解密、密钥缓存、二维码截图等工具模块 5. 完善docker构建与实例部署配置
342 lines
13 KiB
Python
342 lines
13 KiB
Python
"""SQLCipher 密钥统一存储与持久化（single source of truth）。
|
||
|
||
设计目标：
|
||
- 统一密钥来源：salt -> enc_key 映射、默认 key、已拒绝 salt、来源/验证/pid 元信息
|
||
全部在此类管理，DbReader 与 server.py 都通过此类读写，避免双份同步问题
|
||
- 持久化：自动加载/保存 woc-keys.json，容器重启后复用
|
||
- 线程安全：内置 threading.Lock
|
||
- 进程失效检测：pid / start_time 变化时 is_valid_for_pid 返回 False
|
||
|
||
使用方式：
|
||
key_cache = KeyCache(keys_file="/config/woc-keys.json")
|
||
|
||
# 注入（server.py / db_reader.py 都可调用）
|
||
key_cache.set_default_key(env_key, source="env", verified=False)
|
||
key_cache.set_key(key_hex, salt_hex=salt, source="api", verified=True)
|
||
key_cache.set_key_map(key_map, source="auto_extract", verified=True, pid=pid)
|
||
|
||
# 查询（DbReader 用）
|
||
key = key_cache.get_key(salt_hex) # 严格按 salt，回退到未拒绝的默认 key
|
||
has = key_cache.has_keys()
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import datetime
|
||
import json
|
||
import os
|
||
import threading
|
||
from typing import Any, Optional
|
||
|
||
|
||
class KeyCache:
|
||
"""DB 解密密钥统一存储。
|
||
|
||
所有密钥状态（salt->key 映射、默认 key、已拒绝 salt、来源/验证/pid）
|
||
全部在此类管理，DbReader 与 server.py 都通过此类读写，避免双份同步。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
keys_file: str = "/config/woc-keys.json",
|
||
) -> None:
|
||
"""
|
||
Args:
|
||
keys_file: 持久化密钥文件路径（/config/woc-keys.json）
|
||
"""
|
||
self.keys_file = keys_file
|
||
|
||
# 持久化的 salt -> enc_key 映射
|
||
self._salt_to_key: dict[str, str] = {}
|
||
# 默认 key（env 注入或单 key 兼容），不持久化（未按 salt 验证）
|
||
self._default_key: Optional[str] = None
|
||
# 已验证默认 key 不匹配的 salt 集合，避免反复触发整库解密失败
|
||
self._rejected_default_salts: set[str] = set()
|
||
# 持久化文件中记录的 wxid
|
||
self._wxid: Optional[str] = None
|
||
|
||
# 元信息（不持久化，进程级状态）
|
||
self.source: Optional[str] = None # env / api / auto_extract
|
||
self.verified: bool = False
|
||
self.pid: Optional[int] = None
|
||
self.start_time: Optional[float] = None
|
||
|
||
self._lock = threading.Lock()
|
||
self._load_from_file()
|
||
|
||
# ------------------------------------------------------------------
|
||
# 文件持久化
|
||
# ------------------------------------------------------------------
|
||
def _load_from_file(self) -> None:
|
||
"""从 keys_file 加载密钥到内存。"""
|
||
if not os.path.isfile(self.keys_file):
|
||
return
|
||
try:
|
||
with open(self.keys_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
if not isinstance(data, dict):
|
||
return
|
||
self._wxid = data.get("_wxid")
|
||
self._salt_to_key = {
|
||
k: v for k, v in data.items()
|
||
if not k.startswith("_") and isinstance(v, str)
|
||
}
|
||
except (json.JSONDecodeError, OSError):
|
||
pass
|
||
|
||
def _save_to_file(self) -> None:
|
||
"""持久化 salt -> key 映射到 keys_file（调用方需持锁）。"""
|
||
try:
|
||
os.makedirs(os.path.dirname(self.keys_file), exist_ok=True)
|
||
data: dict[str, Any] = {
|
||
"_wxid": self._wxid,
|
||
"_updated_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||
}
|
||
data.update(self._salt_to_key)
|
||
with open(self.keys_file, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||
except OSError:
|
||
pass
|
||
|
||
def set_wxid(self, wxid: Optional[str]) -> None:
|
||
"""更新持久化文件中的 wxid 字段。"""
|
||
with self._lock:
|
||
self._wxid = wxid
|
||
self._save_to_file()
|
||
|
||
# ------------------------------------------------------------------
|
||
# 查询
|
||
# ------------------------------------------------------------------
|
||
def get_key(self, salt_hex: Optional[str] = None) -> Optional[str]:
|
||
"""返回指定 salt 的有效密钥。
|
||
|
||
严格按 salt 查找；未命中时仅对未拒绝过的 salt 回退到默认 key，
|
||
避免反复用错误默认 key 触发整库解密失败。
|
||
"""
|
||
with self._lock:
|
||
if salt_hex:
|
||
if salt_hex in self._salt_to_key:
|
||
return self._salt_to_key[salt_hex]
|
||
if self._default_key and salt_hex not in self._rejected_default_salts:
|
||
return self._default_key
|
||
return None
|
||
return self._default_key
|
||
|
||
def get_default_key(self) -> Optional[str]:
|
||
"""返回默认密钥（向后兼容旧接口）。"""
|
||
with self._lock:
|
||
return self._default_key
|
||
|
||
def get_keys(self) -> dict[str, str]:
|
||
"""返回当前持久化的 salt -> key 映射副本。"""
|
||
with self._lock:
|
||
return dict(self._salt_to_key)
|
||
|
||
def has_keys(self) -> bool:
|
||
"""是否有可用的密钥（持久化或默认 key）。"""
|
||
with self._lock:
|
||
return bool(self._salt_to_key or self._default_key)
|
||
|
||
def has_persisted_keys(self) -> bool:
|
||
"""是否有持久化到文件的 salt -> key 映射。"""
|
||
with self._lock:
|
||
return bool(self._salt_to_key)
|
||
|
||
def key_prefix(self) -> Optional[str]:
|
||
"""返回缓存密钥的摘要（不泄露完整 key）。
|
||
|
||
多 salt 时返回数量 + 第一个 key 的前 4/后 4 位。
|
||
"""
|
||
with self._lock:
|
||
if self._default_key:
|
||
k = self._default_key
|
||
return f"{k[:4]}...{k[-4:]}"
|
||
if not self._salt_to_key:
|
||
return None
|
||
first_key = next(iter(self._salt_to_key.values()))
|
||
prefix = f"{first_key[:4]}...{first_key[-4:]}"
|
||
if len(self._salt_to_key) > 1:
|
||
return f"{prefix} (+{len(self._salt_to_key) - 1} salts)"
|
||
return prefix
|
||
|
||
@property
|
||
def cached(self) -> bool:
|
||
"""是否有缓存的 key（持久化或默认）。"""
|
||
with self._lock:
|
||
return self._default_key is not None or bool(self._salt_to_key)
|
||
|
||
def is_valid_for_pid(
|
||
self,
|
||
current_pid: Optional[int],
|
||
current_start_time: Optional[float] = None,
|
||
) -> bool:
|
||
"""检查缓存的 key 是否对当前进程仍有效。
|
||
|
||
pid 变化或 start_time 变化时返回 False（需重新提取）。
|
||
"""
|
||
if not self.cached:
|
||
return False
|
||
with self._lock:
|
||
if current_pid is not None and self.pid is not None and current_pid != self.pid:
|
||
return False
|
||
if current_start_time is not None and self.start_time is not None and current_start_time != self.start_time:
|
||
return False
|
||
return True
|
||
|
||
# ------------------------------------------------------------------
|
||
# 注入
|
||
# ------------------------------------------------------------------
|
||
def set_key(
|
||
self,
|
||
key_hex: Optional[str],
|
||
salt_hex: Optional[str] = None,
|
||
source: Optional[str] = None,
|
||
verified: bool = True,
|
||
pid: Optional[int] = None,
|
||
start_time: Optional[float] = None,
|
||
) -> None:
|
||
"""设置单个密钥。
|
||
|
||
Args:
|
||
key_hex: 64 位十六进制密钥；为 None 时清空所有密钥与元信息
|
||
salt_hex: 为 None 时设默认 key，否则设 salt -> key 映射
|
||
source: env / api / auto_extract
|
||
verified: 是否已通过 HMAC 验证
|
||
pid / start_time: 自动提取时关联的进程信息
|
||
"""
|
||
with self._lock:
|
||
if key_hex is None:
|
||
# 清空所有
|
||
self._salt_to_key.clear()
|
||
self._default_key = None
|
||
self._rejected_default_salts.clear()
|
||
self.source = None
|
||
self.verified = False
|
||
self.pid = None
|
||
self.start_time = None
|
||
self._save_to_file()
|
||
return
|
||
|
||
if salt_hex is None:
|
||
self._default_key = key_hex
|
||
else:
|
||
self._salt_to_key[salt_hex] = key_hex
|
||
# 新 key 注入时清除该 salt 的拒绝标记
|
||
self._rejected_default_salts.discard(salt_hex)
|
||
self._save_to_file()
|
||
|
||
if source is not None:
|
||
self.source = source
|
||
self.verified = verified
|
||
if pid is not None:
|
||
self.pid = pid
|
||
if start_time is not None:
|
||
self.start_time = start_time
|
||
|
||
def set_default_key(
|
||
self,
|
||
key_hex: Optional[str],
|
||
source: Optional[str] = None,
|
||
verified: bool = False,
|
||
) -> None:
|
||
"""设置默认密钥（env 注入专用），不影响 salt_to_key 映射。
|
||
|
||
与 set_key(salt_hex=None) 的区别：
|
||
- set_key(salt_hex=None) 会清空 salt_to_key 并覆盖 woc-keys.json
|
||
- set_default_key 仅设 _default_key，保留已持久化的多 salt 映射
|
||
|
||
用于 WOC_DB_KEY 环境变量加载：容器重启时先从 woc-keys.json 恢复多 salt
|
||
映射，再用 env key 作为默认兜底，二者共存而非互斥。
|
||
"""
|
||
with self._lock:
|
||
self._default_key = key_hex
|
||
if source is not None:
|
||
self.source = source
|
||
self.verified = verified
|
||
# 默认 key 不持久化到文件（未按 salt 验证，避免污染 woc-keys.json）
|
||
|
||
def set_key_map(
|
||
self,
|
||
key_map: dict[str, str],
|
||
source: Optional[str] = None,
|
||
verified: bool = True,
|
||
pid: Optional[int] = None,
|
||
start_time: Optional[float] = None,
|
||
persist: bool = True,
|
||
) -> None:
|
||
"""批量设置 salt -> enc_key 映射。
|
||
|
||
Args:
|
||
key_map: {salt_hex: enc_key_hex}
|
||
source: env / api / auto_extract
|
||
verified: 是否已通过 HMAC 验证
|
||
pid / start_time: 自动提取时关联的进程信息
|
||
persist: 是否持久化到文件（默认 True）
|
||
"""
|
||
with self._lock:
|
||
self._salt_to_key = dict(key_map)
|
||
# 新 keymap 注入时清空拒绝标记（旧 salt 可能已重新匹配）
|
||
self._rejected_default_salts.clear()
|
||
if source is not None:
|
||
self.source = source
|
||
self.verified = verified
|
||
if pid is not None:
|
||
self.pid = pid
|
||
if start_time is not None:
|
||
self.start_time = start_time
|
||
if persist:
|
||
self._save_to_file()
|
||
|
||
def clear(self) -> None:
|
||
"""清空所有密钥与元信息（含持久化文件）。"""
|
||
with self._lock:
|
||
self._salt_to_key.clear()
|
||
self._default_key = None
|
||
self._rejected_default_salts.clear()
|
||
self.source = None
|
||
self.verified = False
|
||
self.pid = None
|
||
self.start_time = None
|
||
self._save_to_file()
|
||
|
||
def clear_default_key(self) -> None:
|
||
"""仅清空默认 key（保留 salt_to_key 持久化映射）。
|
||
|
||
用于 env key 验证失败后清除默认 key，不影响已持久化的多 salt 映射。
|
||
"""
|
||
with self._lock:
|
||
self._default_key = None
|
||
self._rejected_default_salts.clear()
|
||
|
||
def mark_default_key_mismatch(self, salt_hex: str) -> None:
|
||
"""标记默认 key 不匹配该 salt。
|
||
|
||
由 DbReader._ensure_decrypted 在解密失败且使用的是默认 key 时调用，
|
||
避免后续相同 salt 的查询反复触发整库解密。
|
||
"""
|
||
with self._lock:
|
||
if self._default_key:
|
||
self._rejected_default_salts.add(salt_hex)
|
||
|
||
def set_meta(
|
||
self,
|
||
source: Optional[str] = None,
|
||
verified: Optional[bool] = None,
|
||
pid: Optional[int] = None,
|
||
start_time: Optional[float] = None,
|
||
) -> None:
|
||
"""更新元信息（不影响密钥存储）。
|
||
|
||
用于外部已通过其他途径设置 key 后补充元信息（如 source="api"）。
|
||
所有参数为 None 时不更新对应字段。
|
||
"""
|
||
with self._lock:
|
||
if source is not None:
|
||
self.source = source
|
||
if verified is not None:
|
||
self.verified = verified
|
||
if pid is not None:
|
||
self.pid = pid
|
||
if start_time is not None:
|
||
self.start_time = start_time
|