kris/WechatOnCloud
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
5548f8d89f
WechatOnCloud/bridge/db_reader.py
Kris 5548f8d89f feat(server,db_reader): add SSE real-time message push feature 
This commit implements the full SSE real-time message push capability:
1. Add new `/api/messages/stream` SSE endpoint for real-time message delivery
2. Add MessageStreamer backend that polls DB with mtime awareness and broadcasts events
3. Add helper methods in DbReader: get_db_mtime and get_max_create_time
4. Optimize existing DB connections with `isolation_level=None`
5. Update bridge version to 1.5.0 and add `sse_push` to capability list
6. Add proper error handling and fallback logic for DB unavailability
2026-07-07 20:17:06 +08:00

1095 lines
43 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""微信本地 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"
def get_db_mtime(self, rel_path: str) -> Optional[tuple[float, float]]:
"""返回指定 DB 及其 WAL 文件的 mtime，用于增量轮询感知变化。
Args:
rel_path: 相对路径，如 message/message_0.db
Returns:
(db_mtime, wal_mtime)：DB 不存在时返回 None；
WAL 不存在时 wal_mtime 为 0.0；mtime 读取失败记为 0.0。
"""
db_path = self._get_db_abs_path(rel_path)
if db_path is None:
return None
try:
db_mtime = os.path.getmtime(db_path)
except OSError:
db_mtime = 0.0
wal_path = db_path + "-wal"
try:
wal_mtime = os.path.getmtime(wal_path) if os.path.exists(wal_path) else 0.0
except OSError:
wal_mtime = 0.0
return db_mtime, wal_mtime
# ------------------------------------------------------------------
# 辅助：表/列探测
# ------------------------------------------------------------------
@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, isolation_level=None)
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, isolation_level=None)
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 get_max_create_time(self) -> Optional[int]:
"""返回 message 表中最大的 create_time，用于初始化推送游标。
DB 不可读或 message 表不存在时返回 None。供 MessageStreamer 在启动时
把全局游标对齐到当前最新消息，避免向新连接推送全量历史。
"""
try:
db_path = self._ensure_decrypted("message/message_0.db")
except BridgeError:
return None
conn = sqlite3.connect(db_path, isolation_level=None)
try:
cur = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name='message'"
)
if cur.fetchone() is None:
return None
cur = conn.execute("SELECT MAX(create_time) FROM message")
row = cur.fetchone()
return int(row[0]) if row and row[0] is not None else None
except Exception:
return None
finally:
conn.close()
# ------------------------------------------------------------------
# 联系人字段映射（公共）
# ------------------------------------------------------------------
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, isolation_level=None)
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, isolation_level=None)
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, isolation_level=None)
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, isolation_level=None)
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, isolation_level=None)
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")
Reference in New Issue View Git Blame Copy Permalink