Kris
a2201d8be3
新增完整的woc-bridge业务服务，实现微信API代理、密钥自动提取、消息队列限流等功能： 1. 新增bridge核心代码与s6服务配置 2. 新增ptrace权限初始化脚本与docker配置 3. 新增M2M鉴权API代理与环境变量配置 4. 新增SQLCipher解密、密钥缓存、二维码截图等工具模块 5. 完善docker构建与实例部署配置
1046 lines
41 KiB
Python
1046 lines
41 KiB
Python
"""微信本地 DB 只读访问。
|
||
|
||
通过解密缓存避免锁定原 DB，再用 sqlite3 读取消息。所有方法为同步阻塞操作，
|
||
在 server.py 中通过 asyncio.to_thread() 包装调用，避免阻塞 event loop。
|
||
|
||
架构升级（bridge 1.3.0+）：
|
||
- 密钥统一由 KeyCache 管理（single source of truth），DbReader 不再自持密钥存储
|
||
- 密钥持久化、salt 映射、默认 key、已拒绝 salt 全部由 KeyCache 负责
|
||
- 解密后的明文 DB 缓存到 /config/woc-decrypted/<wxid>/，按 mtime 失效
|
||
- WAL 文件自动合并到明文缓存
|
||
- /api/status 不触发内存扫描，直接判断 keys 文件和缓存有效性
|
||
- 新增 clear_default_key 转发，env/api key 验证失败时仅清默认 key
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import atexit
|
||
import os
|
||
import re as _re
|
||
import sqlite3
|
||
import threading
|
||
from typing import Optional
|
||
|
||
from decryptor import Decryptor
|
||
from key_cache import KeyCache
|
||
from models import BridgeError
|
||
|
||
|
||
# SQLCipher page 常量
|
||
PAGE_SIZE = 4096
|
||
SALT_SIZE = 16
|
||
|
||
# 消息类型 → render_type 映射（规格 12.2 节）
|
||
_TYPE_TO_RENDER_TYPE: dict[int, str] = {
|
||
1: "text",
|
||
3: "image",
|
||
34: "voice",
|
||
43: "video",
|
||
49: "file",
|
||
10000: "system",
|
||
10002: "system",
|
||
}
|
||
|
||
|
||
class DbReader:
|
||
"""微信本地 DB 只读访问器。
|
||
|
||
核心变化：
|
||
1. 以 <wxid>/db_storage 为根，通过相对路径访问各类 DB
|
||
2. 密钥统一由 key_cache（KeyCache）管理，DbReader 仅查询不存储
|
||
3. 解密缓存到 decrypted_dir（/config/woc-decrypted/<wxid>/）
|
||
4. 查询时按需解密 + 自动 WAL 合并
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
db_root: str = "/config",
|
||
key_cache: Optional[KeyCache] = None,
|
||
decrypted_dir: str = "/config/woc-decrypted",
|
||
) -> None:
|
||
"""初始化 DbReader。
|
||
|
||
Args:
|
||
db_root: 微信数据根目录，默认 /config
|
||
key_cache: 密钥统一存储（KeyCache 实例）。为 None 时内部创建
|
||
默认 KeyCache（/config/woc-keys.json），便于独立测试
|
||
decrypted_dir: 解密缓存根目录
|
||
"""
|
||
self.db_root = db_root
|
||
self.decrypted_dir = decrypted_dir
|
||
# 密钥统一来源：所有 key 查询/注入通过 key_cache
|
||
self.key_cache: KeyCache = key_cache if key_cache is not None else KeyCache()
|
||
|
||
self._lock = threading.Lock()
|
||
|
||
# 解密缓存的内存索引：rel_path -> (db_mtime, wal_mtime, decrypted_path)
|
||
self._decrypted_cache: dict[str, tuple[float, float, str]] = {}
|
||
|
||
# 媒体路径缓存：msg_id -> path（或 None），避免每次 get_media 都 os.walk
|
||
self._media_path_cache: dict[str, Optional[str]] = {}
|
||
self._media_path_cache_max = 2000
|
||
|
||
# 进程退出时清理未持久化的临时文件
|
||
atexit.register(self._cleanup_temp_files)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 当前账号 / db_storage 路径探测
|
||
# ------------------------------------------------------------------
|
||
def _get_current_wxid(self) -> Optional[str]:
|
||
"""通过 /config/xwechat_files/all_users/login/<wxid>/key_info.db mtime 判断当前登录账号。"""
|
||
login_dir = os.path.join(self.db_root, "xwechat_files", "all_users", "login")
|
||
if not os.path.isdir(login_dir):
|
||
return None
|
||
candidates = []
|
||
for name in os.listdir(login_dir):
|
||
key_info = os.path.join(login_dir, name, "key_info.db")
|
||
if os.path.isfile(key_info):
|
||
candidates.append((name, os.path.getmtime(key_info)))
|
||
if not candidates:
|
||
return None
|
||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||
return candidates[0][0]
|
||
|
||
def _find_db_storage_dir(self) -> Optional[str]:
|
||
"""查找当前登录账号的 db_storage 目录。"""
|
||
base = os.path.join(self.db_root, "xwechat_files")
|
||
if not os.path.isdir(base):
|
||
return None
|
||
|
||
candidates = []
|
||
for name in os.listdir(base):
|
||
if name == "all_users":
|
||
continue
|
||
path = os.path.join(base, name, "db_storage")
|
||
if os.path.isdir(path):
|
||
# 确认目录下确实有 DB
|
||
has_db = False
|
||
for root, _, files in os.walk(path):
|
||
if any(f.endswith(".db") for f in files):
|
||
has_db = True
|
||
break
|
||
if has_db:
|
||
candidates.append((path, os.path.getmtime(path)))
|
||
|
||
if not candidates:
|
||
return None
|
||
|
||
# 优先匹配当前 wxid
|
||
current_wxid = self._get_current_wxid()
|
||
if current_wxid:
|
||
for path, _ in candidates:
|
||
if current_wxid in path:
|
||
return path
|
||
|
||
# 否则选最近修改的
|
||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||
return candidates[0][0]
|
||
|
||
def _find_db_path(self) -> Optional[str]:
|
||
"""返回一个可直接用于验证/探测的 DB 文件路径（contact.db 或 message_0.db）。"""
|
||
storage = self._find_db_storage_dir()
|
||
if storage is None:
|
||
return None
|
||
for rel in ("contact/contact.db", "message/message_0.db"):
|
||
path = os.path.join(storage, rel)
|
||
if os.path.isfile(path):
|
||
return path
|
||
return None
|
||
|
||
def _get_db_abs_path(self, rel_path: str) -> Optional[str]:
|
||
"""把相对路径（如 contact/contact.db）转成绝对路径。
|
||
|
||
优先按拼接路径返回；若文件不存在则按 basename 在 storage 下递归查找，
|
||
兼容 WeChat 4.x 不同版本 DB 实际存放路径可能与硬编码相对路径不一致的情况。
|
||
"""
|
||
storage = self._find_db_storage_dir()
|
||
if storage is None:
|
||
return None
|
||
abs_path = os.path.normpath(os.path.join(storage, rel_path))
|
||
# 防止路径穿越
|
||
if not abs_path.startswith(os.path.normpath(storage) + os.sep):
|
||
return None
|
||
# 快速路径：拼接路径存在则直接返回
|
||
if os.path.isfile(abs_path):
|
||
return abs_path
|
||
# fallback：按 basename 在 storage 下递归查找
|
||
target_name = os.path.basename(rel_path)
|
||
for root, _, files in os.walk(storage):
|
||
if target_name in files:
|
||
return os.path.join(root, target_name)
|
||
return None
|
||
|
||
def _get_wxid_from_storage(self, storage_dir: Optional[str] = None) -> Optional[str]:
|
||
"""从 db_storage 路径推断 wxid。"""
|
||
if storage_dir is None:
|
||
storage_dir = self._find_db_storage_dir()
|
||
if storage_dir is None:
|
||
return None
|
||
# 路径形如 /config/xwechat_files/wxid_xxx_1234/db_storage
|
||
parts = os.path.normpath(storage_dir).split(os.sep)
|
||
for part in reversed(parts):
|
||
if part.startswith("wxid_"):
|
||
return part
|
||
return None
|
||
|
||
# ------------------------------------------------------------------
|
||
# 密钥访问（全部转发到 KeyCache，DbReader 不再自持密钥存储）
|
||
# ------------------------------------------------------------------
|
||
def set_keys(
|
||
self,
|
||
key_map: dict[str, str],
|
||
source: str = "auto_extract",
|
||
pid: Optional[int] = None,
|
||
start_time: Optional[float] = None,
|
||
) -> None:
|
||
"""批量设置 salt -> enc_key 映射。
|
||
|
||
转发到 KeyCache.set_key_map，由 KeyCache 负责持久化。
|
||
keys 变化后清空解密缓存索引，避免旧密钥解密的脏数据被复用。
|
||
|
||
Args:
|
||
key_map: {salt_hex: enc_key_hex}
|
||
source: 密钥来源（auto_extract / api / env）
|
||
pid: 自动提取时关联的微信进程 PID（用于进程重启后失效检测）
|
||
start_time: 自动提取时关联的进程启动时间
|
||
"""
|
||
wxid = self._get_current_wxid()
|
||
self.key_cache.set_wxid(wxid)
|
||
self.key_cache.set_key_map(
|
||
key_map,
|
||
source=source,
|
||
verified=True,
|
||
pid=pid,
|
||
start_time=start_time,
|
||
)
|
||
# keys 变化后清空解密缓存，让查询重新生成
|
||
self._decrypted_cache.clear()
|
||
|
||
def set_key(
|
||
self,
|
||
key_hex: Optional[str],
|
||
salt_hex: Optional[str] = None,
|
||
) -> None:
|
||
"""设置单个密钥（转发到 KeyCache）。
|
||
|
||
任何 key 变更都会清空解密缓存索引，避免旧密钥解密的脏数据被复用。
|
||
"""
|
||
if key_hex is None:
|
||
self.key_cache.clear()
|
||
else:
|
||
self.key_cache.set_key(key_hex, salt_hex=salt_hex)
|
||
self._decrypted_cache.clear()
|
||
|
||
def set_default_key(self, key_hex: Optional[str]) -> None:
|
||
"""设置默认密钥（env 注入专用，转发到 KeyCache）。
|
||
|
||
不影响 salt_to_key 持久化映射，env key 仅作为默认兜底。
|
||
"""
|
||
self.key_cache.set_default_key(key_hex, source="env", verified=False)
|
||
|
||
def clear_default_key(self) -> None:
|
||
"""仅清空默认 key（保留持久化 salt_to_key 映射）。
|
||
|
||
用于 env/api key 验证失败后清除默认 key，不影响已持久化的多 salt 映射。
|
||
keys 变化后清空解密缓存索引。
|
||
"""
|
||
self.key_cache.clear_default_key()
|
||
self._decrypted_cache.clear()
|
||
|
||
def get_keys(self) -> dict[str, str]:
|
||
"""返回当前持久化的 salt -> key 映射副本。"""
|
||
return self.key_cache.get_keys()
|
||
|
||
def has_keys(self) -> bool:
|
||
"""是否有可用的密钥（持久化或默认 key）。"""
|
||
return self.key_cache.has_keys()
|
||
|
||
def _get_effective_key(self, salt_hex: Optional[str] = None) -> Optional[str]:
|
||
"""返回指定 salt 的有效密钥（转发到 KeyCache）。"""
|
||
return self.key_cache.get_key(salt_hex)
|
||
|
||
def _mark_default_key_mismatch(self, salt_hex: str) -> None:
|
||
"""标记默认 key 不匹配该 salt（转发到 KeyCache）。"""
|
||
self.key_cache.mark_default_key_mismatch(salt_hex)
|
||
|
||
def _read_db_salt(self, db_path: str) -> Optional[str]:
|
||
"""读取 DB 第 1 页前 16 字节 salt。"""
|
||
try:
|
||
with open(db_path, "rb") as f:
|
||
page1 = f.read(PAGE_SIZE)
|
||
if len(page1) < SALT_SIZE:
|
||
return None
|
||
return page1[:SALT_SIZE].hex()
|
||
except OSError:
|
||
return None
|
||
|
||
# ------------------------------------------------------------------
|
||
# 解密缓存
|
||
# ------------------------------------------------------------------
|
||
def _cleanup_temp_files(self) -> None:
|
||
"""进程退出时清理临时解密文件（如果有）。"""
|
||
# 持久化缓存保留，不删除
|
||
pass
|
||
|
||
def _decrypted_path_for(self, rel_path: str, wxid: Optional[str] = None) -> str:
|
||
"""计算解密缓存文件路径。"""
|
||
if wxid is None:
|
||
wxid = self._get_wxid_from_storage() or "unknown"
|
||
# rel_path 本身通常以 .db 结尾（如 contact/contact.db），直接保留即可
|
||
safe_rel = rel_path.replace(os.sep, "_").replace("/", "_")
|
||
cache_dir = os.path.join(self.decrypted_dir, wxid)
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
return os.path.join(cache_dir, safe_rel)
|
||
|
||
def _is_plain_db(self, db_path: str) -> bool:
|
||
"""DB 是否为明文 SQLite。"""
|
||
try:
|
||
with open(db_path, "rb") as f:
|
||
header = f.read(16)
|
||
return header == b"SQLite format 3\x00"
|
||
except OSError:
|
||
return False
|
||
|
||
def _ensure_decrypted(self, rel_path: str) -> str:
|
||
"""确保指定相对路径的 DB 已解密到缓存，返回明文 DB 绝对路径。
|
||
|
||
流程：
|
||
1. 定位原 DB
|
||
2. 明文 DB 直接返回原路径
|
||
3. 读取 salt，查找对应 key
|
||
4. 检查解密缓存是否有效（mtime + key 一致）
|
||
5. 无效则重新解密，并合并 WAL
|
||
"""
|
||
db_path = self._get_db_abs_path(rel_path)
|
||
if db_path is None:
|
||
raise BridgeError(
|
||
code="DB_NOT_FOUND",
|
||
message=f"未找到微信 DB: {rel_path}",
|
||
)
|
||
if not os.path.isfile(db_path):
|
||
raise BridgeError(
|
||
code="DB_NOT_FOUND",
|
||
message=f"DB 文件不存在: {rel_path}",
|
||
)
|
||
|
||
# 明文直接返回
|
||
if self._is_plain_db(db_path):
|
||
return db_path
|
||
|
||
# 加密 DB：需要 key
|
||
salt_hex = self._read_db_salt(db_path)
|
||
if not salt_hex:
|
||
raise BridgeError(
|
||
code="DB_ENCRYPTED",
|
||
message=f"无法读取 DB salt: {rel_path}",
|
||
)
|
||
|
||
key_hex = self._get_effective_key(salt_hex)
|
||
if not key_hex:
|
||
raise BridgeError(
|
||
code="DB_NEED_INIT",
|
||
message="DB 已加密且没有可用密钥，请先调用 /api/db/init",
|
||
)
|
||
|
||
wxid = self._get_wxid_from_storage()
|
||
decrypted_path = self._decrypted_path_for(rel_path, wxid)
|
||
wal_path = db_path + "-wal"
|
||
|
||
try:
|
||
db_mtime = os.path.getmtime(db_path)
|
||
wal_mtime = os.path.getmtime(wal_path) if os.path.exists(wal_path) else 0.0
|
||
except OSError as e:
|
||
raise BridgeError(
|
||
code="DB_NOT_FOUND",
|
||
message=f"无法获取 DB mtime: {e}",
|
||
)
|
||
|
||
# 检查内存缓存索引
|
||
cached = self._decrypted_cache.get(rel_path)
|
||
if cached and cached[0] == db_mtime and cached[1] == wal_mtime and os.path.exists(cached[2]):
|
||
return cached[2]
|
||
|
||
with self._lock:
|
||
# 双重检查
|
||
cached = self._decrypted_cache.get(rel_path)
|
||
if cached and cached[0] == db_mtime and cached[1] == wal_mtime and os.path.exists(cached[2]):
|
||
return cached[2]
|
||
|
||
# 解密
|
||
decryptor = Decryptor(key_hex)
|
||
result = decryptor.decrypt_db(db_path, decrypted_path)
|
||
if not result.get("success"):
|
||
error = result.get("error", "unknown")
|
||
if os.path.exists(decrypted_path):
|
||
os.remove(decrypted_path)
|
||
# 默认 key 解密失败时标记 salt，避免后续重复尝试
|
||
if key_hex == self.key_cache.get_default_key() and salt_hex:
|
||
self._mark_default_key_mismatch(salt_hex)
|
||
raise BridgeError(
|
||
code="DB_ENCRYPTED",
|
||
message=f"解密 {rel_path} 失败: {error}",
|
||
)
|
||
|
||
# 合并 WAL
|
||
if os.path.exists(wal_path):
|
||
try:
|
||
decryptor.decrypt_wal(wal_path, decrypted_path)
|
||
except Exception:
|
||
# WAL 合并失败不影响主 DB 读取，只是可能缺最新数据
|
||
pass
|
||
|
||
os.chmod(decrypted_path, 0o644)
|
||
self._decrypted_cache[rel_path] = (db_mtime, wal_mtime, decrypted_path)
|
||
return decrypted_path
|
||
|
||
def invalidate_decrypted_cache(self) -> None:
|
||
"""清空解密缓存索引（不删除文件，下次按 mtime 重新判断）。"""
|
||
with self._lock:
|
||
self._decrypted_cache.clear()
|
||
|
||
# ------------------------------------------------------------------
|
||
# DB 可达性状态
|
||
# ------------------------------------------------------------------
|
||
def check_db_status(self) -> str:
|
||
"""检测 DB 状态。
|
||
|
||
Returns:
|
||
"ok": 有 DB 且可直接/解密读取
|
||
"not_found": 未找到 DB
|
||
"need_init": DB 加密但没有密钥文件/内存 key
|
||
"key_invalid": 有密钥但验证失败（极少见）
|
||
"""
|
||
storage = self._find_db_storage_dir()
|
||
if storage is None:
|
||
return "not_found"
|
||
|
||
# 找一个需要读取的 DB 探针（contact.db 或 message_0.db）
|
||
probe_rel_paths = ["contact/contact.db", "message/message_0.db"]
|
||
probe_path = None
|
||
for rel in probe_rel_paths:
|
||
p = os.path.join(storage, rel)
|
||
if os.path.isfile(p):
|
||
probe_path = p
|
||
break
|
||
|
||
if probe_path is None:
|
||
# 目录存在但找不到探针文件，可能 DB 还没完全生成
|
||
return "not_found"
|
||
|
||
if self._is_plain_db(probe_path):
|
||
return "ok"
|
||
|
||
if not self.has_keys():
|
||
return "need_init"
|
||
|
||
# 有 keys，快速验证是否能解密
|
||
salt_hex = self._read_db_salt(probe_path)
|
||
if salt_hex and self._get_effective_key(salt_hex):
|
||
return "ok"
|
||
|
||
return "key_invalid"
|
||
|
||
def is_db_accessible(self) -> bool:
|
||
"""DB 是否可读。"""
|
||
return self.check_db_status() == "ok"
|
||
|
||
# ------------------------------------------------------------------
|
||
# 辅助：表/列探测
|
||
# ------------------------------------------------------------------
|
||
@staticmethod
|
||
def _table_columns(conn: sqlite3.Connection, table: str) -> list[str]:
|
||
"""用 PRAGMA table_info 取指定表的所有列名。"""
|
||
try:
|
||
cur = conn.execute(f"PRAGMA table_info({table})")
|
||
return [str(row[1]) for row in cur.fetchall() if row[1]]
|
||
except sqlite3.Error:
|
||
return []
|
||
|
||
@staticmethod
|
||
def _pick_column(columns: list[str], candidates: list[str]) -> Optional[str]:
|
||
"""从候选列名中返回第一个匹配（不区分大小写）。"""
|
||
lower = {c.lower(): c for c in columns}
|
||
for cand in candidates:
|
||
if cand.lower() in lower:
|
||
return lower[cand.lower()]
|
||
return None
|
||
|
||
@staticmethod
|
||
def _find_contact_table(conn: sqlite3.Connection) -> Optional[str]:
|
||
"""探测联系人表名。"""
|
||
try:
|
||
cur = conn.execute(
|
||
"SELECT name FROM sqlite_master WHERE type='table'"
|
||
)
|
||
all_tables = {row[0] for row in cur.fetchall() if row[0]}
|
||
except sqlite3.Error:
|
||
return None
|
||
for cand in ("contact", "contacts", "rcontact"):
|
||
if cand in all_tables:
|
||
return cand
|
||
for name in all_tables:
|
||
if "contact" in name.lower():
|
||
return name
|
||
return None
|
||
|
||
@staticmethod
|
||
def _find_table_by_name(conn: sqlite3.Connection, candidates: list[str]) -> Optional[str]:
|
||
"""按候选名查找表。"""
|
||
try:
|
||
cur = conn.execute(
|
||
"SELECT name FROM sqlite_master WHERE type='table'"
|
||
)
|
||
all_tables = {row[0] for row in cur.fetchall() if row[0]}
|
||
except sqlite3.Error:
|
||
return None
|
||
for cand in candidates:
|
||
if cand in all_tables:
|
||
return cand
|
||
return None
|
||
|
||
# ------------------------------------------------------------------
|
||
# 当前账号信息
|
||
# ------------------------------------------------------------------
|
||
def get_self_info(self) -> dict:
|
||
"""读取当前登录账号的 wxid 与 nickname。"""
|
||
result = {"wxid": "", "nickname": ""}
|
||
storage = self._find_db_storage_dir()
|
||
if storage is None:
|
||
return result
|
||
wxid = self._get_wxid_from_storage(storage)
|
||
result["wxid"] = wxid or ""
|
||
try:
|
||
db_path = self._ensure_decrypted("contact/contact.db")
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
table = self._find_contact_table(conn)
|
||
if table is None:
|
||
return result
|
||
cols = self._table_columns(conn, table)
|
||
username_col = self._pick_column(cols, ["username", "wxid"])
|
||
nickname_col = self._pick_column(cols, ["nickname", "nick_name"])
|
||
if username_col is None or not wxid:
|
||
return result
|
||
cur = conn.execute(
|
||
f"SELECT {nickname_col or 'NULL'} AS nickname FROM {table} "
|
||
f"WHERE {username_col} = ? LIMIT 1",
|
||
(wxid,),
|
||
)
|
||
row = cur.fetchone()
|
||
if row:
|
||
result["nickname"] = str(row["nickname"] or "")
|
||
finally:
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
return result
|
||
|
||
# ------------------------------------------------------------------
|
||
# 消息查询
|
||
# ------------------------------------------------------------------
|
||
def _row_to_message(self, row: sqlite3.Row) -> dict:
|
||
"""将 SQL 行映射为 Message dict。"""
|
||
msg_id = str(row["msg_id"])
|
||
talker = row["talker"] or ""
|
||
is_sender = bool(row["is_sender"])
|
||
msg_type = int(row["type"]) if row["type"] is not None else 0
|
||
content = row["content"] or ""
|
||
create_time = int(row["create_time"]) if row["create_time"] is not None else 0
|
||
|
||
render_type = _TYPE_TO_RENDER_TYPE.get(msg_type, "text")
|
||
session_type = "group" if talker.endswith("@chatroom") else "p2p"
|
||
|
||
sender = ""
|
||
if session_type == "group" and is_sender is False and "\n" in content:
|
||
head, _, _ = content.partition("\n")
|
||
if head.endswith(":"):
|
||
sender = head[:-1]
|
||
content = content.split("\n", 1)[1]
|
||
|
||
return {
|
||
"msg_id": msg_id,
|
||
"talker": talker,
|
||
"sender": sender,
|
||
"is_sender": is_sender,
|
||
"type": msg_type,
|
||
"render_type": render_type,
|
||
"content": content,
|
||
"create_time": create_time,
|
||
"session_type": session_type,
|
||
}
|
||
|
||
def get_messages_since(self, cursor: int, limit: int = 50) -> dict:
|
||
"""读取 create_time > cursor 的增量消息。"""
|
||
db_path = self._ensure_decrypted("message/message_0.db")
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
cur = conn.execute(
|
||
"SELECT name FROM sqlite_master WHERE type='table' AND name='message'"
|
||
)
|
||
if cur.fetchone() is None:
|
||
raise BridgeError(
|
||
code="DB_NOT_FOUND",
|
||
message="message 表不存在",
|
||
)
|
||
cur = conn.execute(
|
||
"SELECT msg_id, talker, is_sender, type, content, create_time "
|
||
"FROM message WHERE create_time > ? ORDER BY create_time ASC LIMIT ?",
|
||
(cursor, limit),
|
||
)
|
||
rows = cur.fetchall()
|
||
finally:
|
||
conn.close()
|
||
|
||
messages = [self._row_to_message(row) for row in rows]
|
||
next_cursor = messages[-1]["create_time"] if messages else cursor
|
||
has_more = len(messages) >= limit
|
||
return {
|
||
"messages": messages,
|
||
"next_cursor": next_cursor,
|
||
"has_more": has_more,
|
||
}
|
||
|
||
# ------------------------------------------------------------------
|
||
# 联系人字段映射（公共）
|
||
# ------------------------------------------------------------------
|
||
def _pick_contact_columns(self, cols: list[str]) -> dict[str, Optional[str]]:
|
||
"""根据实际表列名，挑选 contact 表各字段对应的列名。"""
|
||
return {
|
||
"username": self._pick_column(cols, ["username", "wxid"]),
|
||
"nickname": self._pick_column(cols, ["nickname", "nick_name"]),
|
||
"remark": self._pick_column(cols, ["remark", "conRemark"]),
|
||
"alias": self._pick_column(cols, ["alias"]),
|
||
"encrypt_username": self._pick_column(cols, ["encrypt_username", "encryptUsername"]),
|
||
"quan_pin": self._pick_column(cols, ["quan_pin", "quanPin"]),
|
||
"pin_yin_initial": self._pick_column(cols, ["pin_yin_initial", "pinYinInitial"]),
|
||
"big_head_url": self._pick_column(cols, ["big_head_url", "bigHeadImgUrl"]),
|
||
"small_head_url": self._pick_column(cols, ["small_head_url", "smallHeadImgUrl", "head_img_url"]),
|
||
"description": self._pick_column(cols, ["description", "signature"]),
|
||
"local_type": self._pick_column(cols, ["local_type", "localType"]),
|
||
"verify_flag": self._pick_column(cols, ["verify_flag", "verifyFlag"]),
|
||
"delete_flag": self._pick_column(cols, ["delete_flag", "deleteFlag"]),
|
||
"chat_room_type": self._pick_column(cols, ["chat_room_type", "chatRoomType"]),
|
||
}
|
||
|
||
@staticmethod
|
||
def _avatar_url_from_row(row: sqlite3.Row, wxid: str) -> str:
|
||
"""优先使用 DB 里的头像 URL，没有则回退到本地 avatar 接口。"""
|
||
for col in ("small_head_url", "big_head_url", "avatar"):
|
||
val = row[col] if col in row.keys() else None
|
||
if val:
|
||
return str(val)
|
||
return f"/api/media/avatar/{wxid}"
|
||
|
||
def _row_to_contact(self, row: sqlite3.Row, contact_type: str = "friend") -> dict:
|
||
"""把 sqlite3.Row 转成 Contact 兼容字典。"""
|
||
wxid = str(row["username"] or "")
|
||
|
||
def _val(col: str) -> Optional[str]:
|
||
if col not in row.keys():
|
||
return None
|
||
v = row[col]
|
||
if v is None:
|
||
return None
|
||
return str(v) if not isinstance(v, bytes) else None
|
||
|
||
def _int_val(col: str) -> Optional[int]:
|
||
if col not in row.keys():
|
||
return None
|
||
v = row[col]
|
||
if v is None:
|
||
return None
|
||
try:
|
||
return int(v)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
return {
|
||
"wxid": wxid,
|
||
"nickname": _val("nickname") or "",
|
||
"remark": _val("remark") or "",
|
||
"avatar_url": self._avatar_url_from_row(row, wxid),
|
||
"type": contact_type,
|
||
"alias": _val("alias"),
|
||
"encrypt_username": _val("encrypt_username"),
|
||
"quan_pin": _val("quan_pin"),
|
||
"pin_yin_initial": _val("pin_yin_initial"),
|
||
"big_head_url": _val("big_head_url"),
|
||
"small_head_url": _val("small_head_url"),
|
||
"description": _val("description"),
|
||
"local_type": _int_val("local_type"),
|
||
"verify_flag": _int_val("verify_flag"),
|
||
"delete_flag": _int_val("delete_flag"),
|
||
"chat_room_type": _int_val("chat_room_type"),
|
||
}
|
||
|
||
def _select_contact_columns(self, col_map: dict[str, Optional[str]]) -> str:
|
||
"""生成 SELECT 子句，只选实际存在的列。"""
|
||
selects = [f"{col_map['username']} AS username"]
|
||
for alias, col in col_map.items():
|
||
if alias == "username" or col is None:
|
||
continue
|
||
selects.append(f"{col} AS {alias}")
|
||
return ", ".join(selects)
|
||
|
||
def get_contacts(self, keyword: str = "", limit: int = 50) -> dict:
|
||
"""联系人查询。"""
|
||
empty = {"contacts": [], "total": 0}
|
||
try:
|
||
db_path = self._ensure_decrypted("contact/contact.db")
|
||
except BridgeError:
|
||
return empty
|
||
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
table = self._find_contact_table(conn)
|
||
if table is None:
|
||
return empty
|
||
cols = self._table_columns(conn, table)
|
||
col_map = self._pick_contact_columns(cols)
|
||
username_col = col_map["username"]
|
||
nickname_col = col_map["nickname"]
|
||
remark_col = col_map["remark"]
|
||
if username_col is None:
|
||
return empty
|
||
|
||
select_sql = self._select_contact_columns(col_map)
|
||
if keyword:
|
||
escaped = keyword.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||
like = f"%{escaped}%"
|
||
where_clauses = [f"{username_col} LIKE ? ESCAPE '\\'"]
|
||
params: list = [like]
|
||
if nickname_col:
|
||
where_clauses.append(f"{nickname_col} LIKE ? ESCAPE '\\'")
|
||
params.append(like)
|
||
if remark_col:
|
||
where_clauses.append(f"{remark_col} LIKE ? ESCAPE '\\'")
|
||
params.append(like)
|
||
where_sql = " OR ".join(where_clauses)
|
||
sql = (
|
||
f"SELECT {select_sql} "
|
||
f"FROM {table} WHERE ({where_sql}) "
|
||
f"AND {username_col} NOT LIKE '%@chatroom' "
|
||
f"LIMIT ?"
|
||
)
|
||
params.append(limit)
|
||
cur = conn.execute(sql, params)
|
||
else:
|
||
sql = (
|
||
f"SELECT {select_sql} "
|
||
f"FROM {table} "
|
||
f"WHERE {username_col} NOT LIKE '%@chatroom' "
|
||
f"LIMIT ?"
|
||
)
|
||
cur = conn.execute(sql, (limit,))
|
||
contacts = [self._row_to_contact(row, contact_type="friend") for row in cur.fetchall()]
|
||
return {"contacts": contacts, "total": len(contacts)}
|
||
except sqlite3.Error:
|
||
return empty
|
||
finally:
|
||
conn.close()
|
||
|
||
def get_contact_detail(self, wxid: str) -> Optional[dict]:
|
||
"""单条联系人查询。"""
|
||
try:
|
||
db_path = self._ensure_decrypted("contact/contact.db")
|
||
except BridgeError:
|
||
return None
|
||
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
table = self._find_contact_table(conn)
|
||
if table is None:
|
||
return None
|
||
cols = self._table_columns(conn, table)
|
||
col_map = self._pick_contact_columns(cols)
|
||
username_col = col_map["username"]
|
||
if username_col is None:
|
||
return None
|
||
select_sql = self._select_contact_columns(col_map)
|
||
sql = f"SELECT {select_sql} FROM {table} WHERE {username_col} = ? LIMIT 1"
|
||
cur = conn.execute(sql, (wxid,))
|
||
row = cur.fetchone()
|
||
if row is None:
|
||
return None
|
||
return self._row_to_contact(row, contact_type="friend")
|
||
except sqlite3.Error:
|
||
return None
|
||
finally:
|
||
conn.close()
|
||
|
||
# ------------------------------------------------------------------
|
||
# 群聊与群成员查询
|
||
# ------------------------------------------------------------------
|
||
def get_groups(self, limit: int = 50) -> dict:
|
||
"""查询群聊列表。"""
|
||
empty = {"groups": [], "total": 0}
|
||
try:
|
||
db_path = self._ensure_decrypted("contact/contact.db")
|
||
except BridgeError:
|
||
return empty
|
||
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
table = self._find_contact_table(conn)
|
||
if table is None:
|
||
return empty
|
||
cols = self._table_columns(conn, table)
|
||
col_map = self._pick_contact_columns(cols)
|
||
username_col = col_map["username"]
|
||
if username_col is None:
|
||
return empty
|
||
select_sql = self._select_contact_columns(col_map)
|
||
sql = (
|
||
f"SELECT {select_sql} "
|
||
f"FROM {table} WHERE {username_col} LIKE '%@chatroom' LIMIT ?"
|
||
)
|
||
cur = conn.execute(sql, (limit,))
|
||
groups = [self._row_to_contact(row, contact_type="group") for row in cur.fetchall()]
|
||
return {"groups": groups, "total": len(groups)}
|
||
except sqlite3.Error:
|
||
return empty
|
||
finally:
|
||
conn.close()
|
||
|
||
def get_group_members(self, group_wxid: str) -> dict:
|
||
"""查询群成员。"""
|
||
empty = {"group_wxid": group_wxid, "members": [], "total": 0}
|
||
try:
|
||
db_path = self._ensure_decrypted("contact/contact.db")
|
||
except BridgeError:
|
||
return empty
|
||
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
table = self._find_table_by_name(
|
||
conn,
|
||
["chatroom_member", "chatroom_members", "group_member", "group_members"],
|
||
)
|
||
if table is None:
|
||
return empty
|
||
cols = self._table_columns(conn, table)
|
||
member_wxid_col = self._pick_column(cols, ["username", "wxid", "member_wxid"])
|
||
nickname_col = self._pick_column(cols, ["nickname", "nick_name"])
|
||
display_name_col = self._pick_column(cols, ["display_name", "displayName", "group_nickname"])
|
||
group_col = self._pick_column(cols, ["chatroom", "group_id", "chatroom_username", "room_id"])
|
||
admin_col = self._pick_column(
|
||
cols,
|
||
["is_admin", "admin", "is_manager", "manager", "room_admin", "is_room_admin"],
|
||
)
|
||
if member_wxid_col is None or group_col is None:
|
||
return empty
|
||
sql = (
|
||
f"SELECT {member_wxid_col} AS wxid, "
|
||
f"{nickname_col or 'NULL'} AS nickname, "
|
||
f"{display_name_col or 'NULL'} AS display_name, "
|
||
f"{admin_col or 'NULL'} AS is_admin "
|
||
f"FROM {table} WHERE {group_col} = ?"
|
||
)
|
||
cur = conn.execute(sql, (group_wxid,))
|
||
|
||
owner_wxid = self._get_chatroom_owner(conn, group_wxid)
|
||
|
||
members: list[dict] = []
|
||
for row in cur.fetchall():
|
||
wxid = str(row["wxid"] or "")
|
||
raw_admin = row["is_admin"]
|
||
if isinstance(raw_admin, str) and raw_admin.lower() in ("false", "0", "no", ""):
|
||
raw_admin = False
|
||
is_admin = bool(raw_admin) or (bool(wxid) and wxid == owner_wxid)
|
||
members.append({
|
||
"wxid": wxid,
|
||
"nickname": str(row["nickname"] or ""),
|
||
"display_name": str(row["display_name"] or ""),
|
||
"is_admin": is_admin,
|
||
})
|
||
return {"group_wxid": group_wxid, "members": members, "total": len(members)}
|
||
except sqlite3.Error:
|
||
return empty
|
||
finally:
|
||
conn.close()
|
||
|
||
def _get_chatroom_owner(self, conn: sqlite3.Connection, group_wxid: str) -> Optional[str]:
|
||
"""从 chatroom 表查询群主 wxid。"""
|
||
table = self._find_table_by_name(conn, ["chatroom", "chatrooms", "chat_room"])
|
||
if table is None:
|
||
return None
|
||
cols = self._table_columns(conn, table)
|
||
owner_col = self._pick_column(cols, ["roomowner", "room_owner", "owner"])
|
||
if owner_col is None:
|
||
return None
|
||
group_col = self._pick_column(
|
||
cols,
|
||
["chatroom", "chatroomname", "chatroom_username", "group_id", "username", "room_id"],
|
||
)
|
||
if group_col is None:
|
||
return None
|
||
try:
|
||
cur = conn.execute(
|
||
f"SELECT {owner_col} AS owner FROM {table} "
|
||
f"WHERE {group_col} = ? LIMIT 1",
|
||
(group_wxid,),
|
||
)
|
||
row = cur.fetchone()
|
||
if row is None:
|
||
return None
|
||
return str(row["owner"] or "") or None
|
||
except sqlite3.Error:
|
||
return None
|
||
|
||
# ------------------------------------------------------------------
|
||
# 媒体路径解析
|
||
# ------------------------------------------------------------------
|
||
def get_media_path(self, msg_id: str) -> Optional[str]:
|
||
"""根据 msg_id 解析媒体文件路径。"""
|
||
info = self._get_media_info(msg_id)
|
||
if info is None:
|
||
return None
|
||
return info.get("path")
|
||
|
||
def get_media_bytes(self, msg_id: str) -> Optional[tuple[bytes, str]]:
|
||
"""读取媒体文件内容。"""
|
||
path = self.get_media_path(msg_id)
|
||
if path is None or not os.path.isfile(path) or not os.access(path, os.R_OK):
|
||
return None
|
||
try:
|
||
with open(path, "rb") as f:
|
||
data = f.read()
|
||
except OSError:
|
||
return None
|
||
mime = self._guess_mime(path)
|
||
return (data, mime)
|
||
|
||
def _get_media_info(self, msg_id: str) -> Optional[dict]:
|
||
"""从 DB 查询单条消息的 type / content。"""
|
||
try:
|
||
db_path = self._ensure_decrypted("message/message_0.db")
|
||
except BridgeError:
|
||
return None
|
||
|
||
conn = sqlite3.connect(db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
try:
|
||
cols = self._table_columns(conn, "message")
|
||
if not cols:
|
||
return None
|
||
msg_id_col = self._pick_column(cols, ["msg_id"])
|
||
type_col = self._pick_column(cols, ["type"])
|
||
content_col = self._pick_column(cols, ["content"])
|
||
if not (msg_id_col and type_col and content_col):
|
||
return None
|
||
cur = conn.execute(
|
||
f"SELECT {type_col} AS type, {content_col} AS content "
|
||
f"FROM message WHERE {msg_id_col} = ? LIMIT 1",
|
||
(msg_id,),
|
||
)
|
||
row = cur.fetchone()
|
||
if row is None:
|
||
return None
|
||
msg_type = int(row["type"]) if row["type"] is not None else 0
|
||
content = str(row["content"] or "")
|
||
path = self._resolve_media_path(msg_id, msg_type, content)
|
||
return {"type": msg_type, "content": content, "path": path}
|
||
except sqlite3.Error:
|
||
return None
|
||
finally:
|
||
conn.close()
|
||
|
||
def _resolve_media_path(self, msg_id: str, msg_type: int, content: str) -> Optional[str]:
|
||
"""根据消息类型与内容启发式解析媒体文件路径。
|
||
|
||
优化点：
|
||
- 媒体路径按 msg_id 缓存，避免每次 get_media 都 os.walk 全盘扫描
|
||
- 优先扫描已知媒体目录（attachment/message/file/image/video/voice），
|
||
避免遍历 db_storage 下的 .db 文件
|
||
- 移除 5000 文件硬上限（改为按目录优先级扫描，命中即返回）
|
||
"""
|
||
if msg_type not in (3, 34, 43, 49):
|
||
return None
|
||
# 命中缓存
|
||
if msg_id in self._media_path_cache:
|
||
return self._media_path_cache[msg_id]
|
||
|
||
result: Optional[str] = None
|
||
db_root_abs = os.path.realpath(self.db_root)
|
||
|
||
# 策略 1：从 content 提取路径
|
||
if content:
|
||
matches = _re.findall(r"/[^\s\"'<>]+", content)
|
||
for cand in matches:
|
||
if not os.path.isfile(cand) or not os.access(cand, os.R_OK):
|
||
continue
|
||
cand_abs = os.path.realpath(cand)
|
||
if cand_abs.startswith(db_root_abs + os.sep):
|
||
result = cand
|
||
break
|
||
|
||
# 策略 2：在已知媒体目录下按 msg_id 前缀查找
|
||
if result is None and msg_id:
|
||
storage = self._find_db_storage_dir()
|
||
search_dirs: list[str] = []
|
||
if storage:
|
||
# 媒体文件通常在 db_storage 同级或上层目录
|
||
parent = os.path.dirname(storage)
|
||
for sub in ("attachment", "message", "file", "image", "video", "voice", "audio"):
|
||
p = os.path.join(parent, sub)
|
||
if os.path.isdir(p):
|
||
search_dirs.append(p)
|
||
# db_storage 本身也作为兜底
|
||
search_dirs.append(storage)
|
||
if not search_dirs and os.path.isdir(self.db_root):
|
||
search_dirs = [self.db_root]
|
||
|
||
for search_dir in search_dirs:
|
||
if result:
|
||
break
|
||
for dirpath, _, filenames in os.walk(search_dir):
|
||
if result:
|
||
break
|
||
for name in filenames:
|
||
if name.endswith(".db"):
|
||
continue
|
||
stem, _ = os.path.splitext(name)
|
||
if stem == msg_id or stem.startswith(msg_id + "_") or stem.startswith(msg_id + "-"):
|
||
full = os.path.join(dirpath, name)
|
||
if os.access(full, os.R_OK):
|
||
result = full
|
||
break
|
||
|
||
# 入缓存（含 None 结果，避免重复扫描未命中的 msg_id）
|
||
if len(self._media_path_cache) >= self._media_path_cache_max:
|
||
# FIFO 淘汰最旧
|
||
oldest = next(iter(self._media_path_cache))
|
||
del self._media_path_cache[oldest]
|
||
self._media_path_cache[msg_id] = result
|
||
return result
|
||
|
||
@staticmethod
|
||
def _guess_mime(path: str) -> str:
|
||
"""根据扩展名粗略猜测 MIME 类型。"""
|
||
ext = os.path.splitext(path)[1].lower()
|
||
return {
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".png": "image/png",
|
||
".gif": "image/gif",
|
||
".bmp": "image/bmp",
|
||
".webp": "image/webp",
|
||
".mp3": "audio/mpeg",
|
||
".wav": "audio/wav",
|
||
".amr": "audio/amr",
|
||
".mp4": "video/mp4",
|
||
".pdf": "application/pdf",
|
||
".doc": "application/msword",
|
||
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
".xls": "application/vnd.ms-excel",
|
||
".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||
".zip": "application/zip",
|
||
".txt": "text/plain",
|
||
}.get(ext, "application/octet-stream")
|