WechatOnCloud/bridge/woc_bridge/db/coordinator.py

"""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/<pid>/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/<pid>/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))