"""DB 状态协调:DB 状态解析 / 密钥自动提取 / DB 可读性检查 / 重试装饰器。 从 server.py 抽出的 DB 协调逻辑,供 routes/ 与 app.py 调用。 所有函数通过 _state(from config.py)访问全局状态。 """ from __future__ import annotations import asyncio import functools import inspect import logging import os import time from dataclasses import dataclass from typing import Optional, get_type_hints from woc_bridge.config import _state, _require_db_reader, _require_xdotool from woc_bridge.db.decryptor import Decryptor, _resolve_page1_key_material from woc_bridge.db.key_extractor import KeyExtractor, _get_wechat_pids from woc_bridge.models import BridgeError logger = logging.getLogger("woc-bridge") # --------------------------------------------------------------------------- # with_db_retry 装饰器 # --------------------------------------------------------------------------- def with_db_retry(func): """路由装饰器:遇到 DB_ENCRYPTED 时强制重新提取 key 并重试一次。 DB 加密异常可能发生在: - _check_db_readable():contact.db 探针 key 失效 - DB 读取操作:message_0.db 等 DB 的 key 失效(探针未覆盖) 捕获 DB_ENCRYPTED 后,阻塞调用 _auto_extract_with_lock(force=True) 重新提取 key, 成功后重试一次路由;提取失败则原样抛出。只重试一次,避免死循环。 发送类接口(send_text 等)同样安全:DB_ENCRYPTED 只会在发送前的 DB 读取阶段(_resolve_display_name)抛出,尚未执行实际 UI 发送。 """ @functools.wraps(func) async def wrapper(*args, **kwargs): try: return await func(*args, **kwargs) except BridgeError as e: if e.code != "DB_ENCRYPTED": raise logger.warning( "%s 遇到 DB_ENCRYPTED,强制重新提取 key 后重试一次: %s", func.__name__, e.message, ) extracted = await _auto_extract_with_lock(force=True) if not extracted: raise logger.info("%s key 重新提取成功,重试请求", func.__name__) return await func(*args, **kwargs) # functools.wraps 复制的 __annotations__ 是字符串化注解(因 from __future__ # import annotations),FastAPI 解析 wrapper 签名时拿到字符串注解,无法识别 # Pydantic body 参数,会把它当 query 参数处理导致 422。用 get_type_hints # 解析为真实类型,并据此重建 __signature__,让 FastAPI 能正确识别参数类型。 try: hints = get_type_hints(func, include_extras=True) wrapper.__annotations__ = hints # 基于原始函数签名重建 wrapper 的 __signature__,替换 (*args, **kwargs) orig_sig = inspect.signature(func) new_params = [] for name, param in orig_sig.parameters.items(): annotation = hints.get(name, param.annotation) new_params.append(param.replace(annotation=annotation)) wrapper.__signature__ = orig_sig.replace( parameters=new_params, return_annotation=hints.get("return", orig_sig.return_annotation), ) except Exception: # 解析失败时保持 functools.wraps 的默认行为 pass return wrapper # --------------------------------------------------------------------------- # _resolve_db_state_tuple(供 MessageStreamer 使用的适配器) # --------------------------------------------------------------------------- async def _resolve_db_state_tuple() -> tuple[bool, Optional[str]]: """供 MessageStreamer 使用的 DB 状态解析适配器。 把 _resolve_db_state() 的 DbState 降维为 (db_accessible, db_error_code), 供 MessageStreamer 解耦 server.py 内部数据结构。 """ state = await _resolve_db_state() return (state.db_accessible, state.db_error_code) # --------------------------------------------------------------------------- # DbState 与 DB 状态缓存 # --------------------------------------------------------------------------- @dataclass class DbState: """DB 状态解析结果(由 _resolve_db_state 返回)。 统一三处调用方(get_status / _check_db_readable / diagnostic_run)的 DB 状态判断逻辑,避免重复代码与行为不一致。 """ status: str # ok / not_found / need_init / key_invalid / unreadable db_accessible: bool # 是否可读(ok 或 有有效 key) db_error_code: Optional[str] # 不可读时的错误码 salt_hex: Optional[str] # 当前 DB 的 salt(加密 DB 才有) key_hex: Optional[str] # 当前 DB 对应的有效 key(无则 None) need_extract: bool # 无有效 key 且需要自动提取 # DB 状态短时缓存:避免高频 /api/status 重复读 salt/验证 key _db_state_cache: Optional[tuple[float, DbState]] = None _DB_STATE_CACHE_TTL = 1.0 # 秒 def _invalidate_db_state_cache() -> None: """失效 DB 状态缓存(key 变化、提取完成时调用)。""" global _db_state_cache _db_state_cache = None # --------------------------------------------------------------------------- # _resolve_db_state # --------------------------------------------------------------------------- async def _resolve_db_state() -> DbState: """统一解析 DB 状态、salt、有效 key、pid 失效检测、env/api key 验证。 本函数只做"状态解析",不触发自动提取;提取由调用方按 need_extract 决定。 单次请求内 salt/pid 只读一次,结果携带在 DbState 中供调用方复用。 带 1 秒短时缓存,减少高频轮询时的重复 I/O。 Returns: DbState:含 status / db_accessible / db_error_code / salt_hex / key_hex / need_extract """ global _db_state_cache # 检查缓存 if _db_state_cache is not None: cached_at, cached_state = _db_state_cache if time.monotonic() - cached_at < _DB_STATE_CACHE_TTL: return cached_state db_reader = _require_db_reader() status = await asyncio.to_thread(db_reader.check_db_status) # 明文 DB 或不可达:直接返回 if status == "ok": result = DbState( status=status, db_accessible=True, db_error_code=None, salt_hex=None, key_hex=None, need_extract=False, ) _db_state_cache = (time.monotonic(), result) return result if status in ("not_found", "unreadable"): result = DbState( status=status, db_accessible=False, db_error_code=status, salt_hex=None, key_hex=None, need_extract=False, ) _db_state_cache = (time.monotonic(), result) return result # status in ("need_init", "key_invalid"):加密 DB salt_hex = await asyncio.to_thread(_read_db_salt, db_reader) key = db_reader._get_effective_key(salt_hex) if salt_hex else db_reader._get_effective_key() # auto_extract 的 key 在进程重启后失效 if key and _state.key_cache.cached and _state.key_cache.source == "auto_extract": wechat_pid_info = await _get_wechat_pid() current_pid = wechat_pid_info[0] if wechat_pid_info else None current_start_time = wechat_pid_info[1] if wechat_pid_info else None if not _state.key_cache.is_valid_for_pid(current_pid, current_start_time): _state.key_cache.clear() db_reader.invalidate_decrypted_cache() key = None # env/api 注入的 key 若未验证,快速验证 page1(4KB),避免无效 key 反复失败 if key and _state.key_cache.cached and _state.key_cache.source in ("env", "api") and not _state.key_cache.verified: verified, _key_mode = await asyncio.to_thread(_verify_db_key, db_reader, key) if verified: _state.key_cache.set_meta(verified=True) else: # env/api key 验证失败:仅清默认 key,保留持久化 salt_to_key db_reader.clear_default_key() key = None # 有有效 key:DB 可读 if key: result = DbState( status=status, db_accessible=True, db_error_code="encrypted_key_ok", salt_hex=salt_hex, key_hex=key, need_extract=False, ) _db_state_cache = (time.monotonic(), result) return result # 无有效 key init_in_progress = _state.init_state.state == "running" if init_in_progress: db_error_code = "init_in_progress" elif _state.config.auto_extract_enabled: db_error_code = "need_init" else: db_error_code = "encrypted_no_key" result = DbState( status=status, db_accessible=False, db_error_code=db_error_code, salt_hex=salt_hex, key_hex=None, need_extract=_state.config.auto_extract_enabled and not init_in_progress, ) _db_state_cache = (time.monotonic(), result) return result # --------------------------------------------------------------------------- # _auto_extract_with_lock # --------------------------------------------------------------------------- async def _auto_extract_with_lock(force: bool = False) -> Optional[dict[str, str]]: """带锁的自动密钥提取,结果同步到 init_state。 供 diagnostic_autofix 直接调用(阻塞等待结果); _check_db_readable 通过 asyncio.create_task 在后台调用(非阻塞); with_db_retry 装饰器在 DB_ENCRYPTED 时强制调用(force=True)。 串行化:通过 extract_lock 确保同一时间只有一个提取任务。 Args: force: True 时跳过 _resolve_db_state 探针检查,强制重新扫描内存。 用于 contact.db 探针通过但 message_0.db key 失效的场景 (探针误判 DB 可读,常规路径不会触发重新提取)。 Returns: salt_hex -> enc_key_hex 映射;失败返回 None """ async with _state.extract_lock(): # 双重检查:持锁后可能已被其他请求提取成功(force 模式跳过此检查) if not force: state = await _resolve_db_state() if state.db_accessible: return _state.db_reader.get_keys() if _state.db_reader else None if _state.init_state.state == "running": # 已在运行(不应该发生,lock 串行化),返回 None return None # 标记运行中,阻止后续并发请求重复扫描 _state.init_state.state = "running" _state.init_state.progress_pct = 0.0 _state.init_state.message = "正在扫描内存提取密钥" try: extracted = await _try_auto_extract_key() if extracted: _state.init_state.set_success( len(extracted), f"成功提取 {len(extracted)} 个密钥", ) else: _state.init_state.set_failed("自动提取密钥失败") return extracted except Exception as e: _state.init_state.set_failed(str(e)) return None finally: # 兜底:异常路径未设置 success/failed 时回退 idle if _state.init_state.state == "running": _state.init_state.state = "idle" _state.init_state.progress_pct = None _state.init_state.message = None # 失效 DB 状态缓存,让下次查询看到新 key _invalidate_db_state_cache() # --------------------------------------------------------------------------- # _check_db_readable # --------------------------------------------------------------------------- async def _check_db_readable() -> None: """检查 DB 是否可读,不可读时抛出对应错误码。 非阻塞设计(1.4.0+): - not_found → DB_NOT_FOUND - unreadable → DB_NOT_FOUND - ok 或有有效 key → 不抛异常 - 无有效 key 且提取进行中 → 抛 DB_ENCRYPTED(提示稍后重试) - 无有效 key 且需提取 → 后台触发提取(不等待),立即抛 DB_ENCRYPTED - 无有效 key 且提取已失败 → 抛 DB_ENCRYPTED(提示注入 key) - 自动提取未启用 → 抛 DB_ENCRYPTED(提示注入 key) Raises: BridgeError(DB_NOT_FOUND / DB_ENCRYPTED) """ db_reader = _require_db_reader() state = await _resolve_db_state() if state.status == "not_found": raise BridgeError(code="DB_NOT_FOUND", message="未在 /config 下找到微信消息 DB") if state.status == "unreadable": raise BridgeError(code="DB_NOT_FOUND", message="DB 文件存在但不可读(权限问题)") if state.db_accessible: # ok 或有有效 key:不抛异常 return # 无有效 key init_state = _state.init_state.state if init_state == "running": pct = _state.init_state.progress_pct or 0 raise BridgeError( code="DB_ENCRYPTED", message=f"DB 已加密,密钥提取进行中({pct:.0f}%),请稍后重试", ) if not state.need_extract: # 自动提取未启用,或已提取过但当前 salt 无匹配 if init_state == "failed": raise BridgeError( code="DB_ENCRYPTED", message="DB 已加密(SQLCipher),自动提取密钥失败,需通过 POST /api/db/decrypt 注入密钥或设置 WOC_DB_KEY", ) raise BridgeError( code="DB_ENCRYPTED", message="DB 已加密(SQLCipher),需通过 POST /api/db/decrypt 注入密钥或设置 WOC_DB_KEY", ) # need_extract=True:后台触发提取(非阻塞),立即返回 503 asyncio.create_task(_auto_extract_with_lock()) raise BridgeError( code="DB_ENCRYPTED", message="DB 已加密(SQLCipher),已触发后台密钥提取,请稍后重试", ) # --------------------------------------------------------------------------- # _verify_db_key # --------------------------------------------------------------------------- def _verify_db_key( db_reader: DbReader, key_hex: str, salt_hex: str | None = None, ) -> tuple[bool, Optional[str]]: """验证 key 是否匹配当前 DB 文件,并返回 key_mode。 用 db_reader._find_db_path 找到 DB 文件,调 Decryptor.verify_key 验证。 验证通过时一并返回 key_mode(enc_key / key_material), 避免调用方再读一次 page1。 Args: db_reader: DB 读取器 key_hex: 64 位十六进制密钥 salt_hex: 未使用(保留参数兼容旧调用),自动读取当前 DB salt Returns: (verified, key_mode): verified 为是否验证通过; key_mode 为密钥形态(验证失败时为 None) """ db_path = db_reader._find_db_path() if db_path is None: return (False, None) try: decryptor = Decryptor(key_hex) verified = decryptor.verify_key(db_path) if not verified: return (False, None) # 验证通过,获取 key_mode(复用已读的 page1 逻辑) key_mode = _get_key_mode(key_hex, db_path) return (True, key_mode) except Exception: return (False, None) # --------------------------------------------------------------------------- # _read_db_salt # --------------------------------------------------------------------------- def _read_db_salt(db_reader: DbReader) -> str | None: """读取当前 DB 文件的 salt(同步阻塞)。""" db_path = db_reader._find_db_path() if db_path is None: return None return db_reader._read_db_salt(db_path) # --------------------------------------------------------------------------- # _get_wechat_pid # --------------------------------------------------------------------------- async def _get_wechat_pid() -> tuple[int, float] | None: """获取微信主进程 PID 与启动时间。 使用与 key_extractor 一致的进程识别逻辑(检查 /proc//comm 与 exe), 避免 `pgrep -f xwechat` 漏掉主进程(主进程命令行不含 xwechat)。 Returns: (pid, start_time) 或 None(进程未运行) """ try: pids = await asyncio.to_thread(_get_wechat_pids) if not pids: return None # 取内存占用最大的进程作为主进程 pid, _rss_kb = pids[0] # 读 /proc//stat 的第 22 字段(start_time,单位 jiffies) try: with open(f"/proc/{pid}/stat", "r") as f: stat_fields = f.read().split() start_time = float(stat_fields[21]) if len(stat_fields) > 21 else 0.0 except (OSError, IndexError, ValueError): start_time = 0.0 return (pid, start_time) except Exception: return None # --------------------------------------------------------------------------- # _try_auto_extract_key # --------------------------------------------------------------------------- async def _try_auto_extract_key() -> dict[str, str] | None: """尝试从微信进程内存自动提取 SQLCipher 密钥。 流程: 1. 检查 auto_extract_enabled 开关 2. 检查微信进程是否运行(pgrep) 3. 获取全部微信 PID 候选(不只取第一个) 4. 找到加密 DB 目录作为探针 5. 对每个 PID 调 KeyExtractor.extract_all_keys() 提取 salt->enc_key 映射 6. 成功则缓存到 KeyCache(含 pid / start_time)并注入 DbReader Returns: salt_hex -> enc_key_hex 映射字典;失败返回 None """ if not _state.config.auto_extract_enabled: return None _require_db_reader() _require_xdotool() # 检查微信进程是否运行 wechat_running = await _state.xdotool.is_wechat_running() if not wechat_running: return None # 找到加密 DB 目录作为探针 db_path = await asyncio.to_thread(_state.db_reader._find_db_path) if db_path is None: logger.warning("自动提取密钥失败:未找到 DB 路径") return None # 优先扫描整个 /config/xwechat_files,让 KeyExtractor 自己探测所有微信 PID # 原 pgrep -f xwechat 会漏掉主进程(主进程命令行里没有 xwechat 字样) db_dir = "/config/xwechat_files" if not os.path.isdir(db_dir): db_dir = os.path.dirname(db_path) logger.info("自动提取密钥开始,探针目录: %s", db_dir) try: # pid=0 让 KeyExtractor 用 /proc 自探测,比 pgrep 更可靠 extractor = KeyExtractor(0, db_path) key_map = await asyncio.to_thread(extractor.extract_all_keys, db_dir) if key_map: # 记录成功时关联的 PID(取 key_map 中第一个地址反推太麻烦, # 这里用 _get_wechat_pid 拿到主进程 PID 做缓存校验) wechat_pid_info = await _get_wechat_pid() wechat_pid = wechat_pid_info[0] if wechat_pid_info else 0 start_time = wechat_pid_info[1] if wechat_pid_info else 0.0 # KeyExtractor 内部已通过 page 1 HMAC 验证,无需再次全库验证 # 统一通过 db_reader.set_keys 转发到 KeyCache.set_key_map, # 避免 key_cache 与 db_reader 两份重复同步 _state.db_reader.set_keys( key_map, source="auto_extract", pid=wechat_pid, start_time=start_time, ) _invalidate_db_state_cache() logger.info( "自动提取 DB 密钥成功(主 PID=%s,%d salts),来源:微信进程内存扫描", wechat_pid, len(key_map), ) return key_map except Exception as e: logger.exception("自动提取密钥异常: %s", e) logger.warning("自动提取 DB 密钥失败:扫描全部微信进程后未找到有效密钥") return None # --------------------------------------------------------------------------- # _get_key_mode # --------------------------------------------------------------------------- def _get_key_mode(key_hex: str, db_path: str) -> str | None: """获取 key 形态(enc_key / key_material)。 通过读取 DB 第 1 页 + _resolve_page1_key_material 验证后返回 mode。 """ try: with open(db_path, "rb") as f: page1 = f.read(4096) key_bytes = bytes.fromhex(key_hex) result = _resolve_page1_key_material(key_bytes, page1) if result: return result[2] except Exception: pass return None # --------------------------------------------------------------------------- # run_auto_extract_sync(供后台线程同步调用的密钥提取入口) # --------------------------------------------------------------------------- def run_auto_extract_sync(force: bool = True) -> Optional[dict[str, str]]: """供后台线程同步调用的密钥提取入口。 通过 asyncio.run_coroutine_threadsafe 在主事件循环中执行 _auto_extract_with_lock,确保与 asyncio 路径共享同一把 extract_lock, 避免并发 set_keys() 竞态。 供 routes/db.py 的 _init_task 调用。 """ try: loop = asyncio.get_event_loop() except RuntimeError: # 无事件循环(纯线程环境),创建临时 loop loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result = loop.run_until_complete(_auto_extract_with_lock(force=force)) loop.close() return result if loop.is_running(): # 主事件循环运行中,用 run_coroutine_threadsafe 投递 future = asyncio.run_coroutine_threadsafe( _auto_extract_with_lock(force=force), loop ) return future.result(timeout=120) # 提取最长 120s else: return loop.run_until_complete(_auto_extract_with_lock(force=force))