WechatOnCloud/bridge/woc_bridge/db/key_cache.py

"""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)
            }
            # 恢复元信息：从文件加载的密钥在持久化前已通过验证
            if self._salt_to_key:
                self.source = "file"
                self.verified = True
        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