feat: Rust 完整重写 wx-cli(单一二进制,支持 macOS/Linux/Windows)

实现所有核心模块:
- src/crypto/: SQLCipher 4 页解密 + WAL 应用(AES-256-CBC)
- src/scanner/: 三平台内存扫描(macOS Mach VM / Linux /proc/mem / Windows ReadProcessMemory)
- src/daemon/: tokio 异步 daemon,Unix socket IPC,mtime-aware DB 缓存,WAL 监听推送
- src/cli/: clap CLI,自动启动 daemon,完整命令实现
- src/config.rs: 跨平台配置加载,兼容 Python 版 config.json 格式
- src/ipc.rs: 换行符分隔 JSON 协议,与 Python 版兼容
- .github/workflows/release.yml: 四平台自动构建发布

cargo build --release 验证通过,生成 4.8MB macOS arm64 单一二进制
This commit is contained in:
jackwener
2026-04-16 14:37:10 +08:00
parent 0d5ac82349
commit d475f6219b
27 changed files with 5294 additions and 0 deletions
+215
View File
@@ -0,0 +1,215 @@
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::config;
use crate::crypto;
use crate::crypto::wal;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct MtimeEntry {
db_mt: f64,
wal_mt: f64,
path: String,
}
#[derive(Debug, Clone)]
struct CacheEntry {
db_mtime: f64,
wal_mtime: f64,
decrypted_path: PathBuf,
}
/// 解密后数据库的 mtime-aware 缓存
///
/// 当数据库文件(.db)或 WAL 文件(.db-wal)的 mtime 发生变化时,
/// 自动重新解密并更新缓存。跨进程重启可通过持久化 mtime 文件复用已解密的 DB。
pub struct DbCache {
db_dir: PathBuf,
cache_dir: PathBuf,
all_keys: HashMap<String, String>, // rel_key -> enc_key(hex)
inner: Arc<Mutex<HashMap<String, CacheEntry>>>,
}
impl DbCache {
pub async fn new(
db_dir: PathBuf,
all_keys: HashMap<String, String>,
) -> Result<Self> {
let cache_dir = config::cache_dir();
tokio::fs::create_dir_all(&cache_dir).await?;
let inner: HashMap<String, CacheEntry> = HashMap::new();
let cache = DbCache {
db_dir,
cache_dir,
all_keys,
inner: Arc::new(Mutex::new(inner)),
};
cache.load_persistent().await;
Ok(cache)
}
fn cache_file_path(&self, rel_key: &str) -> PathBuf {
let hash = format!("{:x}", md5::compute(rel_key.as_bytes()));
let short = &hash[..12];
self.cache_dir.join(format!("{}.db", short))
}
/// 从持久化文件加载 mtime 记录,复用未过期的解密文件
async fn load_persistent(&self) {
let mtime_file = config::mtime_file();
let content = match tokio::fs::read_to_string(&mtime_file).await {
Ok(c) => c,
Err(_) => return,
};
let saved: HashMap<String, MtimeEntry> = match serde_json::from_str(&content) {
Ok(v) => v,
Err(_) => return,
};
let mut inner = self.inner.lock().await;
let mut reused = 0usize;
for (rel_key, entry) in &saved {
let dec_path = PathBuf::from(&entry.path);
if !dec_path.exists() {
continue;
}
let db_path = self.db_dir.join(rel_key.replace('\\', std::path::MAIN_SEPARATOR_STR).replace('/', std::path::MAIN_SEPARATOR_STR));
let wal_path_str = format!("{}-wal", db_path.display());
let wal_path = Path::new(&wal_path_str);
let db_mt = mtime_f64(&db_path);
let wal_mt = if wal_path.exists() { mtime_f64(wal_path) } else { 0.0 };
if (db_mt - entry.db_mt).abs() < 0.001 && (wal_mt - entry.wal_mt).abs() < 0.001 {
inner.insert(rel_key.clone(), CacheEntry {
db_mtime: db_mt,
wal_mtime: wal_mt,
decrypted_path: dec_path,
});
reused += 1;
}
}
if reused > 0 {
eprintln!("[cache] 复用 {} 个已解密 DB", reused);
}
}
/// 持久化 mtime 记录
async fn save_persistent(&self) {
let mtime_file = config::mtime_file();
let inner = self.inner.lock().await;
let data: HashMap<String, MtimeEntry> = inner.iter().map(|(k, v)| {
(k.clone(), MtimeEntry {
db_mt: v.db_mtime,
wal_mt: v.wal_mtime,
path: v.decrypted_path.to_string_lossy().into_owned(),
})
}).collect();
drop(inner);
if let Ok(json) = serde_json::to_string_pretty(&data) {
let _ = tokio::fs::write(&mtime_file, json).await;
}
}
/// 获取解密后的数据库路径
///
/// 如果 mtime 未变,直接返回缓存路径;否则重新解密
pub async fn get(&self, rel_key: &str) -> Result<Option<PathBuf>> {
let enc_key_hex = match self.all_keys.get(rel_key) {
Some(k) => k.clone(),
None => return Ok(None),
};
let db_path = self.db_dir.join(
rel_key.replace('\\', std::path::MAIN_SEPARATOR_STR)
.replace('/', std::path::MAIN_SEPARATOR_STR)
);
if !db_path.exists() {
return Ok(None);
}
let wal_path_str = format!("{}-wal", db_path.display());
let wal_path = Path::new(&wal_path_str).to_path_buf();
let db_mt = mtime_f64(&db_path);
let wal_mt = if wal_path.exists() { mtime_f64(&wal_path) } else { 0.0 };
// 检查缓存
{
let inner = self.inner.lock().await;
if let Some(entry) = inner.get(rel_key) {
if (entry.db_mtime - db_mt).abs() < 0.001
&& (entry.wal_mtime - wal_mt).abs() < 0.001
&& entry.decrypted_path.exists()
{
return Ok(Some(entry.decrypted_path.clone()));
}
}
}
// 需要重新解密
let out_path = self.cache_file_path(rel_key);
let enc_key_bytes = hex_to_32bytes(&enc_key_hex)
.with_context(|| format!("密钥格式错误: {}", rel_key))?;
let t0 = std::time::Instant::now();
let db_path2 = db_path.clone();
let out_path2 = out_path.clone();
let key_copy = enc_key_bytes;
tokio::task::spawn_blocking(move || {
crypto::full_decrypt(&db_path2, &out_path2, &key_copy)
}).await??;
// 应用 WAL
if wal_path.exists() {
let out_path3 = out_path.clone();
let wal_path3 = wal_path.clone();
let key_copy2 = enc_key_bytes;
tokio::task::spawn_blocking(move || {
wal::apply_wal(&wal_path3, &out_path3, &key_copy2)
}).await??;
}
let elapsed_ms = t0.elapsed().as_millis();
eprintln!("[cache] 解密 {} ({}ms)", rel_key, elapsed_ms);
// 更新内存缓存
{
let mut inner = self.inner.lock().await;
inner.insert(rel_key.to_string(), CacheEntry {
db_mtime: db_mt,
wal_mtime: wal_mt,
decrypted_path: out_path.clone(),
});
}
self.save_persistent().await;
Ok(Some(out_path))
}
}
fn mtime_f64(path: &Path) -> f64 {
std::fs::metadata(path)
.and_then(|m| m.modified())
.map(|t| t.duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs_f64())
.unwrap_or(0.0)
}
fn hex_to_32bytes(s: &str) -> Result<[u8; 32]> {
if s.len() != 64 {
anyhow::bail!("密钥 hex 长度应为 64,实际为 {}", s.len());
}
let mut out = [0u8; 32];
for i in 0..32 {
out[i] = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16)
.with_context(|| format!("非法 hex 字符 at {}", i * 2))?;
}
Ok(out)
}
+280
View File
@@ -0,0 +1,280 @@
pub mod cache;
pub mod query;
pub mod watcher;
pub mod server;
use anyhow::Result;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::broadcast;
use crate::config;
/// daemon 入口
///
/// 当 WX_DAEMON_MODE 环境变量设置时,main() 调用此函数
pub fn run() {
let rt = tokio::runtime::Runtime::new().expect("无法创建 tokio runtime");
if let Err(e) = rt.block_on(async_run()) {
eprintln!("[daemon] 启动失败: {}", e);
std::process::exit(1);
}
}
async fn async_run() -> Result<()> {
// 确保工作目录存在
let cli_dir = config::cli_dir();
tokio::fs::create_dir_all(&cli_dir).await?;
tokio::fs::create_dir_all(config::cache_dir()).await?;
// 写 PID 文件
let pid = std::process::id();
tokio::fs::write(config::pid_path(), pid.to_string()).await?;
// 注册 SIGTERM / SIGINT 处理
setup_signal_handler().await;
eprintln!("[daemon] wx-daemon 启动 (PID {})", pid);
// 加载配置
let cfg = config::load_config()?;
eprintln!("[daemon] DB_DIR: {}", cfg.db_dir.display());
// 加载密钥
let keys_content = tokio::fs::read_to_string(&cfg.keys_file).await
.map_err(|e| anyhow::anyhow!("读取密钥文件 {:?} 失败: {}", cfg.keys_file, e))?;
let keys_raw: serde_json::Value = serde_json::from_str(&keys_content)?;
let all_keys = extract_keys(&keys_raw);
eprintln!("[daemon] 密钥数量: {}", all_keys.len());
// 初始化 DbCache
let db = Arc::new(cache::DbCache::new(cfg.db_dir.clone(), all_keys.clone()).await?);
// 收集消息 DB 列表
let msg_db_keys: Vec<String> = all_keys.keys()
.filter(|k| {
let k = k.replace('\\', "/");
k.contains("message/message_") && k.ends_with(".db")
})
.cloned()
.collect();
// 预热:加载联系人 + 解密 session.db
eprintln!("[daemon] 预热...");
let names_raw = query::load_names(&*db).await.unwrap_or_else(|e| {
eprintln!("[daemon] 加载联系人失败: {}", e);
query::Names {
map: HashMap::new(),
md5_to_uname: HashMap::new(),
msg_db_keys: Vec::new(),
}
});
let mut names = names_raw;
names.msg_db_keys = msg_db_keys;
let _ = db.get("session/session.db").await;
eprintln!("[daemon] 预热完成,联系人 {} 个", names.map.len());
let names_arc = Arc::new(std::sync::RwLock::new(names));
// 启动 WAL watcher
let (watch_tx, _) = broadcast::channel::<crate::ipc::WatchEvent>(500);
let session_wal = cfg.db_dir.join("session").join("session.db-wal");
// SAFETY: 我们确保 db 和 names_arc 在 daemon 生命周期内有效
// 使用 Arc 传递引用避免 'static 问题
let db_arc = Arc::clone(&db);
let names_arc2 = Arc::clone(&names_arc);
let tx_clone = watch_tx.clone();
let session_wal2 = session_wal.clone();
tokio::spawn(async move {
run_watcher(db_arc, names_arc2, tx_clone, session_wal2).await;
});
// 启动 IPC server(阻塞)
server::serve(Arc::clone(&db), Arc::clone(&names_arc), watch_tx).await?;
Ok(())
}
async fn run_watcher(
db: Arc<cache::DbCache>,
names: Arc<std::sync::RwLock<query::Names>>,
tx: broadcast::Sender<crate::ipc::WatchEvent>,
session_wal: PathBuf,
) {
use std::collections::HashMap;
use std::time::Duration;
use crate::ipc::WatchEvent;
let mut last_mtime = 0.0f64;
let mut last_ts: HashMap<String, i64> = HashMap::new();
let mut initialized = false;
loop {
tokio::time::sleep(Duration::from_millis(500)).await;
if tx.receiver_count() == 0 {
continue;
}
let wal_mtime = match mtime_f64(&session_wal) {
Some(m) => m,
None => continue,
};
if (wal_mtime - last_mtime).abs() < 0.001 {
continue;
}
last_mtime = wal_mtime;
let path = match db.get("session/session.db").await {
Ok(Some(p)) => p,
_ => continue,
};
let path2 = path.clone();
let rows: Vec<(String, Vec<u8>, i64, i64, String)> = match tokio::task::spawn_blocking(move || {
let conn = rusqlite::Connection::open(&path2)?;
let mut stmt = conn.prepare(
"SELECT username, summary, last_timestamp, last_msg_type, last_msg_sender
FROM SessionTable WHERE last_timestamp > 0
ORDER BY last_timestamp DESC LIMIT 50"
)?;
let rows = stmt.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, Vec<u8>>(1).unwrap_or_default(),
row.get::<_, i64>(2)?,
row.get::<_, i64>(3).unwrap_or(0),
row.get::<_, String>(4).unwrap_or_default(),
))
})?.collect::<rusqlite::Result<Vec<_>>>()?;
Ok::<_, anyhow::Error>(rows)
}).await {
Ok(Ok(r)) => r,
_ => continue,
};
let names_guard = match names.read() {
Ok(g) => g,
Err(_) => continue,
};
for (username, summary_bytes, ts, msg_type, sender) in &rows {
if !initialized {
last_ts.insert(username.clone(), *ts);
continue;
}
let prev_ts = last_ts.get(username).copied().unwrap_or(0);
if *ts <= prev_ts {
continue;
}
last_ts.insert(username.clone(), *ts);
let display = names_guard.display(username);
let is_group = username.contains("@chatroom");
let summary = decompress_or_str(summary_bytes);
let summary = if summary.contains(":\n") {
summary.splitn(2, ":\n").nth(1).unwrap_or(&summary).to_string()
} else {
summary
};
let sender_display = if !sender.is_empty() {
names_guard.map.get(sender).cloned().unwrap_or_else(|| sender.clone())
} else {
String::new()
};
let event = WatchEvent {
event: "message".into(),
time: Some(fmt_hhmm(*ts)),
chat: Some(display),
username: Some(username.clone()),
is_group: Some(is_group),
sender: Some(sender_display),
content: Some(summary),
msg_type: Some(query::fmt_type(*msg_type)),
timestamp: Some(*ts),
};
let _ = tx.send(event);
}
if !initialized {
initialized = true;
}
}
}
fn mtime_f64(path: &std::path::Path) -> Option<f64> {
std::fs::metadata(path).ok()?
.modified().ok()
.map(|t| t.duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs_f64())
}
fn decompress_or_str(data: &[u8]) -> String {
if data.is_empty() { return String::new(); }
if let Ok(dec) = zstd::decode_all(data) {
if let Ok(s) = String::from_utf8(dec) { return s; }
}
String::from_utf8_lossy(data).into_owned()
}
fn fmt_hhmm(ts: i64) -> String {
use chrono::{Local, TimeZone};
Local.timestamp_opt(ts, 0)
.single()
.map(|dt| dt.format("%H:%M").to_string())
.unwrap_or_else(|| ts.to_string())
}
/// 从 all_keys.json 提取 rel_key -> enc_key 映射
///
/// 兼容两种格式:
/// - `{ "rel/path.db": { "enc_key": "hex" } }`(Python 版原生格式)
/// - `{ "rel/path.db": "hex" }`(简化格式)
fn extract_keys(json: &serde_json::Value) -> HashMap<String, String> {
let mut result = HashMap::new();
if let Some(obj) = json.as_object() {
for (k, v) in obj {
if k.starts_with('_') { continue; }
let enc_key = if let Some(s) = v.as_str() {
s.to_string()
} else if let Some(obj2) = v.as_object() {
obj2.get("enc_key")
.and_then(|e| e.as_str())
.unwrap_or_default()
.to_string()
} else {
continue;
};
if !enc_key.is_empty() {
// 统一路径分隔符
let rel = k.replace('\\', "/");
result.insert(rel, enc_key);
}
}
}
result
}
/// 设置信号处理(Unix: SIGTERM/SIGINT)
async fn setup_signal_handler() {
#[cfg(unix)]
tokio::spawn(async move {
use tokio::signal::unix::{signal, SignalKind};
let mut term = signal(SignalKind::terminate()).expect("无法监听 SIGTERM");
let mut int = signal(SignalKind::interrupt()).expect("无法监听 SIGINT");
tokio::select! {
_ = term.recv() => {},
_ = int.recv() => {},
}
cleanup_and_exit();
});
}
fn cleanup_and_exit() {
let _ = std::fs::remove_file(config::sock_path());
let _ = std::fs::remove_file(config::pid_path());
std::process::exit(0);
}
+676
View File
@@ -0,0 +1,676 @@
use anyhow::{Context, Result};
use chrono::{Local, TimeZone};
use regex::Regex;
use rusqlite::Connection;
use serde_json::{json, Value};
use std::collections::HashMap;
use super::cache::DbCache;
/// 联系人名称缓存
#[derive(Clone)]
pub struct Names {
/// username -> display_name
pub map: HashMap<String, String>,
/// md5(username) -> username(用于从 Msg_<md5> 表名反推联系人)
pub md5_to_uname: HashMap<String, String>,
/// 消息 DB 的相对路径列表(message/message_N.db)
pub msg_db_keys: Vec<String>,
}
impl Names {
pub fn display(&self, username: &str) -> String {
self.map.get(username).cloned().unwrap_or_else(|| username.to_string())
}
}
/// 加载联系人缓存(从 contact/contact.db)
pub async fn load_names(db: &DbCache) -> Result<Names> {
let path = db.get("contact/contact.db").await?;
let mut map = HashMap::new();
if let Some(p) = path {
let p2 = p.clone();
let rows: Vec<(String, String, String)> = tokio::task::spawn_blocking(move || {
let conn = Connection::open(&p2).context("打开 contact.db 失败")?;
let mut stmt = conn.prepare(
"SELECT username, nick_name, remark FROM contact"
)?;
let rows = stmt.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1).unwrap_or_default(),
row.get::<_, String>(2).unwrap_or_default(),
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok::<_, anyhow::Error>(rows)
}).await??;
for (uname, nick, remark) in rows {
let display = if !remark.is_empty() { remark }
else if !nick.is_empty() { nick }
else { uname.clone() };
map.insert(uname, display);
}
}
let md5_to_uname: HashMap<String, String> = map.keys()
.map(|u| (format!("{:x}", md5::compute(u.as_bytes())), u.clone()))
.collect();
Ok(Names { map, md5_to_uname, msg_db_keys: Vec::new() })
}
/// 查询最近会话列表
pub async fn q_sessions(db: &DbCache, names: &Names, limit: usize) -> Result<Value> {
let path = db.get("session/session.db").await?
.context("无法解密 session.db")?;
let path2 = path.clone();
let limit_val = limit;
let rows: Vec<(String, i64, Vec<u8>, i64, i64, String, String)> = tokio::task::spawn_blocking(move || {
let conn = Connection::open(&path2)?;
let mut stmt = conn.prepare(
"SELECT username, unread_count, summary, last_timestamp,
last_msg_type, last_msg_sender, last_sender_display_name
FROM SessionTable
WHERE last_timestamp > 0
ORDER BY last_timestamp DESC LIMIT ?"
)?;
let rows = stmt.query_map([limit_val as i64], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, i64>(1).unwrap_or(0),
row.get::<_, Vec<u8>>(2).unwrap_or_default(),
row.get::<_, i64>(3).unwrap_or(0),
row.get::<_, i64>(4).unwrap_or(0),
row.get::<_, String>(5).unwrap_or_default(),
row.get::<_, String>(6).unwrap_or_default(),
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok::<_, anyhow::Error>(rows)
}).await??;
let mut results = Vec::new();
for (username, unread, summary_bytes, ts, msg_type, sender, sender_name) in rows {
let display = names.display(&username);
let is_group = username.contains("@chatroom");
// 尝试 zstd 解压 summary
let summary = decompress_or_str(&summary_bytes);
let summary = strip_group_prefix(&summary);
let sender_display = if is_group && !sender.is_empty() {
names.map.get(&sender).cloned().unwrap_or_else(|| {
if !sender_name.is_empty() { sender_name.clone() } else { sender.clone() }
})
} else {
String::new()
};
results.push(json!({
"chat": display,
"username": username,
"is_group": is_group,
"unread": unread,
"last_msg_type": fmt_type(msg_type),
"last_sender": sender_display,
"summary": summary,
"timestamp": ts,
"time": fmt_time(ts, "%m-%d %H:%M"),
}));
}
Ok(json!({ "sessions": results }))
}
/// 查询聊天记录
pub async fn q_history(
db: &DbCache,
names: &Names,
chat: &str,
limit: usize,
offset: usize,
since: Option<i64>,
until: Option<i64>,
) -> Result<Value> {
let username = resolve_username(chat, names)
.with_context(|| format!("找不到联系人: {}", chat))?;
let display = names.display(&username);
let is_group = username.contains("@chatroom");
let tables = find_msg_tables(db, names, &username).await?;
if tables.is_empty() {
return Ok(json!({ "error": format!("找不到 {} 的消息记录", display) }));
}
let mut all_msgs: Vec<Value> = Vec::new();
for (db_path, table_name) in &tables {
let path = db_path.clone();
let tname = table_name.clone();
let uname = username.clone();
let is_group2 = is_group;
let names_map = names.map.clone();
let since2 = since;
let until2 = until;
let limit2 = limit;
let offset2 = offset;
let msgs: Vec<Value> = tokio::task::spawn_blocking(move || {
query_messages(&path, &tname, &uname, is_group2, &names_map, since2, until2, limit2 + offset2, 0)
}).await??;
all_msgs.extend(msgs);
}
all_msgs.sort_by_key(|m| std::cmp::Reverse(m["timestamp"].as_i64().unwrap_or(0)));
let paged: Vec<Value> = all_msgs.into_iter().skip(offset).take(limit).collect();
let mut paged = paged;
paged.sort_by_key(|m| m["timestamp"].as_i64().unwrap_or(0));
Ok(json!({
"chat": display,
"username": username,
"is_group": is_group,
"count": paged.len(),
"messages": paged,
}))
}
/// 搜索消息
pub async fn q_search(
db: &DbCache,
names: &Names,
keyword: &str,
chats: Option<Vec<String>>,
limit: usize,
since: Option<i64>,
until: Option<i64>,
) -> Result<Value> {
let mut targets: Vec<(String, String, String, String)> = Vec::new(); // (path, table, display, uname)
if let Some(chat_names) = chats {
for chat_name in &chat_names {
if let Some(uname) = resolve_username(chat_name, names) {
let tables = find_msg_tables(db, names, &uname).await?;
for (p, t) in tables {
targets.push((p.to_string_lossy().into_owned(), t, names.display(&uname), uname.clone()));
}
}
}
} else {
// 全局搜索:遍历所有消息 DB
for rel_key in &names.msg_db_keys {
let path = match db.get(rel_key).await? {
Some(p) => p,
None => continue,
};
let path2 = path.clone();
let md5_lookup = names.md5_to_uname.clone();
let names_map = names.map.clone();
let table_targets: Vec<(String, String, String, String)> = tokio::task::spawn_blocking(move || {
let conn = Connection::open(&path2)?;
let mut stmt = conn.prepare(
"SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'Msg_%'"
)?;
let table_names: Vec<String> = stmt.query_map([], |row| row.get(0))?
.filter_map(|r| r.ok())
.collect();
let re = Regex::new(r"^Msg_[0-9a-f]{32}$").unwrap();
let mut result = Vec::new();
for tname in table_names {
if !re.is_match(&tname) {
continue;
}
let hash = &tname[4..];
let uname = md5_lookup.get(hash).cloned().unwrap_or_default();
let display = if uname.is_empty() {
String::new()
} else {
names_map.get(&uname).cloned().unwrap_or_else(|| uname.clone())
};
result.push((
path2.to_string_lossy().into_owned(),
tname,
display,
uname,
));
}
Ok::<_, anyhow::Error>(result)
}).await??;
targets.extend(table_targets);
}
}
// 按 db_path 分组
let mut by_path: HashMap<String, Vec<(String, String, String)>> = HashMap::new();
for (p, t, d, u) in targets {
by_path.entry(p).or_default().push((t, d, u));
}
let mut results: Vec<Value> = Vec::new();
let kw = keyword.to_string();
for (db_path, table_list) in by_path {
let kw2 = kw.clone();
let since2 = since;
let until2 = until;
let limit2 = limit * 3;
let found: Vec<Value> = tokio::task::spawn_blocking(move || {
let conn = Connection::open(&db_path)?;
let mut all = Vec::new();
for (tname, display, uname) in &table_list {
let is_group = uname.contains("@chatroom");
let rows = search_in_table(&conn, tname, &uname, is_group,
&HashMap::new(), &kw2, since2, until2, limit2)?;
for mut row in rows {
if row.get("chat").map(|v| v.as_str().unwrap_or("")).unwrap_or("").is_empty() {
if let Some(obj) = row.as_object_mut() {
obj.insert("chat".into(), serde_json::Value::String(
if display.is_empty() { tname.clone() } else { display.clone() }
));
}
}
all.push(row);
}
}
Ok::<_, anyhow::Error>(all)
}).await??;
results.extend(found);
}
results.sort_by_key(|r| std::cmp::Reverse(r["timestamp"].as_i64().unwrap_or(0)));
let paged: Vec<Value> = results.into_iter().take(limit).collect();
Ok(json!({ "keyword": keyword, "count": paged.len(), "results": paged }))
}
/// 查询联系人
pub async fn q_contacts(names: &Names, query: Option<&str>, limit: usize) -> Result<Value> {
let mut contacts: Vec<Value> = names.map.iter()
.filter(|(u, _)| !u.starts_with("gh_") && !u.starts_with("biz_"))
.map(|(u, d)| json!({ "username": u, "display": d }))
.collect();
if let Some(q) = query {
let low = q.to_lowercase();
contacts.retain(|c| {
c["display"].as_str().map(|s| s.to_lowercase().contains(&low)).unwrap_or(false)
|| c["username"].as_str().map(|s| s.to_lowercase().contains(&low)).unwrap_or(false)
});
}
contacts.sort_by(|a, b| {
a["display"].as_str().unwrap_or("").cmp(b["display"].as_str().unwrap_or(""))
});
let total = contacts.len();
contacts.truncate(limit);
Ok(json!({ "contacts": contacts, "total": total }))
}
// ─── 内部辅助函数 ────────────────────────────────────────────────────────────
fn resolve_username(chat_name: &str, names: &Names) -> Option<String> {
if names.map.contains_key(chat_name)
|| chat_name.contains("@chatroom")
|| chat_name.starts_with("wxid_")
{
return Some(chat_name.to_string());
}
let low = chat_name.to_lowercase();
// 精确匹配显示名
for (uname, display) in &names.map {
if low == display.to_lowercase() {
return Some(uname.clone());
}
}
// 模糊匹配
for (uname, display) in &names.map {
if display.to_lowercase().contains(&low) {
return Some(uname.clone());
}
}
None
}
async fn find_msg_tables(
db: &DbCache,
names: &Names,
username: &str,
) -> Result<Vec<(std::path::PathBuf, String)>> {
let table_name = format!("Msg_{:x}", md5::compute(username.as_bytes()));
let re = Regex::new(r"^Msg_[0-9a-f]{32}$").unwrap();
if !re.is_match(&table_name) {
return Ok(Vec::new());
}
let mut results = Vec::new();
for rel_key in &names.msg_db_keys {
let path = match db.get(rel_key).await? {
Some(p) => p,
None => continue,
};
let tname = table_name.clone();
let path2 = path.clone();
let exists: Option<i64> = tokio::task::spawn_blocking(move || {
let conn = Connection::open(&path2)?;
// 检查表是否存在
let table_exists: Option<i64> = conn.query_row(
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
[&tname],
|row| row.get(0),
).ok().flatten();
if table_exists.is_none() {
return Ok::<_, anyhow::Error>(None);
}
let max_ts: Option<i64> = conn.query_row(
&format!("SELECT MAX(create_time) FROM [{}]", tname),
[],
|row| row.get(0),
).ok().flatten();
Ok(max_ts)
}).await??;
if exists.is_some() {
results.push((path.clone(), table_name.clone()));
}
}
// 按最大时间戳排序(最新的优先)
Ok(results)
}
fn query_messages(
db_path: &std::path::Path,
table: &str,
chat_username: &str,
is_group: bool,
names_map: &HashMap<String, String>,
since: Option<i64>,
until: Option<i64>,
limit: usize,
offset: usize,
) -> Result<Vec<Value>> {
let conn = Connection::open(db_path)?;
let id2u = load_id2u(&conn);
let mut clauses = Vec::new();
let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
if let Some(s) = since {
clauses.push("create_time >= ?");
params.push(Box::new(s));
}
if let Some(u) = until {
clauses.push("create_time <= ?");
params.push(Box::new(u));
}
let where_clause = if clauses.is_empty() {
String::new()
} else {
format!("WHERE {}", clauses.join(" AND "))
};
let sql = format!(
"SELECT local_id, local_type, create_time, real_sender_id,
message_content, WCDB_CT_message_content
FROM [{}] {} ORDER BY create_time DESC LIMIT ? OFFSET ?",
table, where_clause
);
params.push(Box::new(limit as i64));
params.push(Box::new(offset as i64));
let params_ref: Vec<&dyn rusqlite::types::ToSql> = params.iter().map(|p| p.as_ref()).collect();
let mut stmt = conn.prepare(&sql)?;
let rows = stmt.query_map(params_ref.as_slice(), |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, i64>(2)?,
row.get::<_, i64>(3)?,
row.get::<_, Vec<u8>>(4).unwrap_or_default(),
row.get::<_, i64>(5).unwrap_or(0),
))
})?
.filter_map(|r| r.ok())
.collect::<Vec<_>>();
let mut result = Vec::new();
for (local_id, local_type, ts, real_sender_id, content_bytes, ct) in rows {
let content = decompress_message(&content_bytes, ct);
let sender = sender_label(real_sender_id, &content, is_group, chat_username, &id2u, names_map);
let text = fmt_content(local_id, local_type, &content, is_group);
result.push(json!({
"timestamp": ts,
"time": fmt_time(ts, "%Y-%m-%d %H:%M"),
"sender": sender,
"content": text,
"type": fmt_type(local_type),
"local_id": local_id,
}));
}
Ok(result)
}
fn search_in_table(
conn: &Connection,
table: &str,
chat_username: &str,
is_group: bool,
names_map: &HashMap<String, String>,
keyword: &str,
since: Option<i64>,
until: Option<i64>,
limit: usize,
) -> Result<Vec<Value>> {
let id2u = load_id2u(conn);
let mut clauses = vec!["message_content LIKE ?".to_string()];
let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = vec![Box::new(format!("%{}%", keyword))];
if let Some(s) = since {
clauses.push("create_time >= ?".into());
params.push(Box::new(s));
}
if let Some(u) = until {
clauses.push("create_time <= ?".into());
params.push(Box::new(u));
}
let where_clause = format!("WHERE {}", clauses.join(" AND "));
let sql = format!(
"SELECT local_id, local_type, create_time, real_sender_id,
message_content, WCDB_CT_message_content
FROM [{}] {} ORDER BY create_time DESC LIMIT ?",
table, where_clause
);
params.push(Box::new(limit as i64));
let params_ref: Vec<&dyn rusqlite::types::ToSql> = params.iter().map(|p| p.as_ref()).collect();
let mut stmt = conn.prepare(&sql)?;
let rows = stmt.query_map(params_ref.as_slice(), |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, i64>(2)?,
row.get::<_, i64>(3)?,
row.get::<_, Vec<u8>>(4).unwrap_or_default(),
row.get::<_, i64>(5).unwrap_or(0),
))
})?
.filter_map(|r| r.ok())
.collect::<Vec<_>>();
let mut result = Vec::new();
for (local_id, local_type, ts, real_sender_id, content_bytes, ct) in rows {
let content = decompress_message(&content_bytes, ct);
let sender = sender_label(real_sender_id, &content, is_group, chat_username, &id2u, names_map);
let text = fmt_content(local_id, local_type, &content, is_group);
result.push(json!({
"timestamp": ts,
"time": fmt_time(ts, "%Y-%m-%d %H:%M"),
"chat": "",
"sender": sender,
"content": text,
"type": fmt_type(local_type),
}));
}
Ok(result)
}
fn load_id2u(conn: &Connection) -> HashMap<i64, String> {
let mut map = HashMap::new();
if let Ok(mut stmt) = conn.prepare("SELECT rowid, user_name FROM Name2Id") {
let _ = stmt.query_map([], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?))
}).map(|rows| {
for r in rows.flatten() {
map.insert(r.0, r.1);
}
});
}
map
}
fn sender_label(
real_sender_id: i64,
content: &str,
is_group: bool,
chat_username: &str,
id2u: &HashMap<i64, String>,
names: &HashMap<String, String>,
) -> String {
let sender_uname = id2u.get(&real_sender_id).cloned().unwrap_or_default();
if is_group {
if !sender_uname.is_empty() && sender_uname != chat_username {
return names.get(&sender_uname).cloned().unwrap_or(sender_uname);
}
if content.contains(":\n") {
let raw = content.splitn(2, ":\n").next().unwrap_or("");
return names.get(raw).cloned().unwrap_or_else(|| raw.to_string());
}
return String::new();
}
if !sender_uname.is_empty() && sender_uname != chat_username {
return names.get(&sender_uname).cloned().unwrap_or(sender_uname);
}
String::new()
}
fn decompress_message(data: &[u8], ct: i64) -> String {
if ct == 4 && !data.is_empty() {
// zstd 压缩
if let Ok(dec) = zstd::decode_all(data) {
return String::from_utf8_lossy(&dec).into_owned();
}
}
String::from_utf8_lossy(data).into_owned()
}
fn decompress_or_str(data: &[u8]) -> String {
if data.is_empty() {
return String::new();
}
// 尝试 zstd 解压
if let Ok(dec) = zstd::decode_all(data) {
if let Ok(s) = String::from_utf8(dec) {
return s;
}
}
String::from_utf8_lossy(data).into_owned()
}
fn strip_group_prefix(s: &str) -> String {
if s.contains(":\n") {
s.splitn(2, ":\n").nth(1).unwrap_or(s).to_string()
} else {
s.to_string()
}
}
pub fn fmt_type(t: i64) -> String {
let base = (t as u64 & 0xFFFFFFFF) as i64;
match base {
1 => "文本".into(),
3 => "图片".into(),
34 => "语音".into(),
42 => "名片".into(),
43 => "视频".into(),
47 => "表情".into(),
48 => "位置".into(),
49 => "链接/文件".into(),
50 => "通话".into(),
10000 => "系统".into(),
10002 => "撤回".into(),
_ => format!("type={}", base),
}
}
fn fmt_content(local_id: i64, local_type: i64, content: &str, is_group: bool) -> String {
let base = (local_type as u64 & 0xFFFFFFFF) as i64;
match base {
3 => return format!("[图片] local_id={}", local_id),
47 => return "[表情]".into(),
50 => return "[通话]".into(),
_ => {}
}
let text = if is_group && content.contains(":\n") {
content.splitn(2, ":\n").nth(1).unwrap_or(content)
} else {
content
};
if base == 49 && text.contains("<appmsg") {
if let Some(parsed) = parse_appmsg(text) {
return parsed;
}
}
text.to_string()
}
fn parse_appmsg(text: &str) -> Option<String> {
// 简单 XML 解析,避免引入重量级 XML 库(或直接用 minidom)
// 这里用基本字符串搜索实现
let title = extract_xml_text(text, "title")?;
let atype = extract_xml_text(text, "type").unwrap_or_default();
match atype.as_str() {
"6" => Some(if !title.is_empty() { format!("[文件] {}", title) } else { "[文件]".into() }),
"57" => {
let ref_content = extract_xml_text(text, "content")
.map(|s| {
let s: String = s.split_whitespace().collect::<Vec<_>>().join(" ");
if s.len() > 80 { format!("{}...", &s[..80]) } else { s }
})
.unwrap_or_default();
let quote = if !title.is_empty() { format!("[引用] {}", title) } else { "[引用]".into() };
if !ref_content.is_empty() {
Some(format!("{}\n \u{21b3} {}", quote, ref_content))
} else {
Some(quote)
}
}
"33" | "36" | "44" => Some(if !title.is_empty() { format!("[小程序] {}", title) } else { "[小程序]".into() }),
_ => Some(if !title.is_empty() { format!("[链接] {}", title) } else { "[链接/文件]".into() }),
}
}
fn extract_xml_text(xml: &str, tag: &str) -> Option<String> {
let open = format!("<{}>", tag);
let close = format!("</{}>", tag);
let start = xml.find(&open)?;
let content_start = start + open.len();
let end = xml[content_start..].find(&close)?;
Some(xml[content_start..content_start + end].trim().to_string())
}
fn fmt_time(ts: i64, fmt: &str) -> String {
Local.timestamp_opt(ts, 0)
.single()
.map(|dt| dt.format(fmt).to_string())
.unwrap_or_else(|| ts.to_string())
}
+219
View File
@@ -0,0 +1,219 @@
use anyhow::Result;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::broadcast;
use crate::ipc::{Request, Response, WatchEvent};
use super::cache::DbCache;
use super::query::Names;
/// 启动 IPC server(Unix socket / Windows named pipe)
pub async fn serve(
db: Arc<DbCache>,
names: Arc<std::sync::RwLock<Names>>,
watch_tx: broadcast::Sender<WatchEvent>,
) -> Result<()> {
#[cfg(unix)]
serve_unix(db, names, watch_tx).await?;
#[cfg(windows)]
serve_windows(db, names, watch_tx).await?;
Ok(())
}
#[cfg(unix)]
async fn serve_unix(
db: Arc<DbCache>,
names: Arc<std::sync::RwLock<Names>>,
watch_tx: broadcast::Sender<WatchEvent>,
) -> Result<()> {
use tokio::net::UnixListener;
let sock_path = crate::config::sock_path();
// 删除旧 socket 文件
if sock_path.exists() {
let _ = tokio::fs::remove_file(&sock_path).await;
}
let listener = UnixListener::bind(&sock_path)?;
// 设置权限 0600
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&sock_path, std::fs::Permissions::from_mode(0o600))?;
}
eprintln!("[server] 监听 {}", sock_path.display());
loop {
let (stream, _) = listener.accept().await?;
let db2 = Arc::clone(&db);
let names2 = Arc::clone(&names);
let tx2 = watch_tx.clone();
tokio::spawn(async move {
if let Err(e) = handle_connection_unix(stream, db2, names2, tx2).await {
eprintln!("[server] 连接处理错误: {}", e);
}
});
}
}
#[cfg(unix)]
async fn handle_connection_unix(
stream: tokio::net::UnixStream,
db: Arc<DbCache>,
names: Arc<std::sync::RwLock<Names>>,
watch_tx: broadcast::Sender<WatchEvent>,
) -> Result<()> {
let (reader, mut writer) = stream.into_split();
let mut lines = BufReader::new(reader).lines();
let line = match lines.next_line().await? {
Some(l) => l,
None => return Ok(()),
};
// 解析请求
let req: Request = match serde_json::from_str(&line) {
Ok(r) => r,
Err(e) => {
let resp = Response::err(format!("JSON 解析错误: {}", e));
writer.write_all(resp.to_json_line()?.as_bytes()).await?;
return Ok(());
}
};
match req {
Request::Watch => {
// 流式模式:持续推送事件
let mut rx = watch_tx.subscribe();
let connected = WatchEvent::connected();
writer.write_all(connected.to_json_line()?.as_bytes()).await?;
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(e) => {
if writer.write_all(e.to_json_line()?.as_bytes()).await.is_err() {
break;
}
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(_) => break,
}
}
_ = tokio::time::sleep(Duration::from_secs(30)) => {
// 心跳
let hb = WatchEvent::heartbeat();
if writer.write_all(hb.to_json_line()?.as_bytes()).await.is_err() {
break;
}
}
}
}
}
other => {
let resp = dispatch(other, &db, &names).await;
writer.write_all(resp.to_json_line()?.as_bytes()).await?;
}
}
Ok(())
}
#[cfg(windows)]
async fn serve_windows(
db: Arc<DbCache>,
names: Arc<std::sync::RwLock<Names>>,
watch_tx: broadcast::Sender<WatchEvent>,
) -> Result<()> {
use interprocess::local_socket::{
tokio::prelude::*, GenericNamespaced, ListenerOptions,
};
let pipe_name = r"\\.\pipe\wechat-cli-daemon";
let name = pipe_name.to_ns_name::<GenericNamespaced>()?;
let opts = ListenerOptions::new().name(name);
let listener = opts.create_tokio()?;
eprintln!("[server] 监听 {}", pipe_name);
loop {
let conn = listener.accept().await?;
let db2 = Arc::clone(&db);
let names2 = Arc::clone(&names);
let tx2 = watch_tx.clone();
tokio::spawn(async move {
if let Err(e) = handle_connection_generic(conn, db2, names2, tx2).await {
eprintln!("[server] 连接处理错误: {}", e);
}
});
}
}
async fn dispatch(
req: Request,
db: &DbCache,
names: &std::sync::RwLock<Names>,
) -> Response {
use crate::ipc::Request::*;
use super::query;
match req {
Ping => Response::ok(serde_json::json!({ "pong": true })),
Sessions { limit } => {
// 在 await 前获取并复制所需数据,避免 RwLockGuard 跨 await
let names_snapshot = match clone_names(names) {
Ok(n) => n,
Err(e) => return Response::err(e),
};
match query::q_sessions(db, &names_snapshot, limit).await {
Ok(v) => Response::ok(v),
Err(e) => Response::err(e.to_string()),
}
}
History { chat, limit, offset, since, until } => {
let names_snapshot = match clone_names(names) {
Ok(n) => n,
Err(e) => return Response::err(e),
};
match query::q_history(db, &names_snapshot, &chat, limit, offset, since, until).await {
Ok(v) => Response::ok(v),
Err(e) => Response::err(e.to_string()),
}
}
Search { keyword, chats, limit, since, until } => {
let names_snapshot = match clone_names(names) {
Ok(n) => n,
Err(e) => return Response::err(e),
};
match query::q_search(db, &names_snapshot, &keyword, chats, limit, since, until).await {
Ok(v) => Response::ok(v),
Err(e) => Response::err(e.to_string()),
}
}
Contacts { query, limit } => {
let names_snapshot = match clone_names(names) {
Ok(n) => n,
Err(e) => return Response::err(e),
};
match query::q_contacts(&names_snapshot, query.as_deref(), limit).await {
Ok(v) => Response::ok(v),
Err(e) => Response::err(e.to_string()),
}
}
Watch => Response::err("Watch 命令不应通过 dispatch 处理"),
}
}
/// 克隆 Names 以避免 RwLockGuard 跨 await
fn clone_names(names: &std::sync::RwLock<Names>) -> Result<Names, String> {
let guard = names.read().map_err(|_| "内部错误: names lock poisoned".to_string())?;
Ok(Names {
map: guard.map.clone(),
md5_to_uname: guard.md5_to_uname.clone(),
msg_db_keys: guard.msg_db_keys.clone(),
})
}
+151
View File
@@ -0,0 +1,151 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;
use tokio::sync::broadcast;
use super::cache::DbCache;
use super::query::{fmt_type, Names};
use crate::ipc::WatchEvent;
/// 启动 WAL 变化监听 task
///
/// 每 500ms 检测 session.db-wal 的 mtime,有变化时重新读 session.db,
/// 找到 timestamp 更新的行,broadcast 到所有 watch 客户端
#[allow(dead_code)]
pub async fn start_watcher(
db: &'static DbCache,
names_ref: &'static std::sync::RwLock<Names>,
tx: broadcast::Sender<WatchEvent>,
session_wal_path: PathBuf,
) {
tokio::spawn(async move {
let mut last_mtime = 0.0f64;
let mut last_ts: HashMap<String, i64> = HashMap::new();
let mut initialized = false;
loop {
tokio::time::sleep(Duration::from_millis(500)).await;
// 如果没有订阅者,跳过
if tx.receiver_count() == 0 {
continue;
}
let wal_mtime = match mtime_f64(&session_wal_path) {
Some(m) => m,
None => continue,
};
if (wal_mtime - last_mtime).abs() < 0.001 {
continue;
}
last_mtime = wal_mtime;
// 重新解密 session.db
let path = match db.get("session/session.db").await {
Ok(Some(p)) => p,
_ => continue,
};
let path2 = path.clone();
let rows: Vec<(String, Vec<u8>, i64, i64, String)> = match tokio::task::spawn_blocking(move || {
let conn = rusqlite::Connection::open(&path2)?;
let mut stmt = conn.prepare(
"SELECT username, summary, last_timestamp, last_msg_type, last_msg_sender
FROM SessionTable WHERE last_timestamp > 0
ORDER BY last_timestamp DESC LIMIT 50"
)?;
let rows = stmt.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, Vec<u8>>(1).unwrap_or_default(),
row.get::<_, i64>(2)?,
row.get::<_, i64>(3).unwrap_or(0),
row.get::<_, String>(4).unwrap_or_default(),
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok::<_, anyhow::Error>(rows)
}).await {
Ok(Ok(r)) => r,
_ => continue,
};
let names_guard = names_ref.read().expect("names lock poisoned");
for (username, summary_bytes, ts, msg_type, sender) in &rows {
if !initialized {
last_ts.insert(username.clone(), *ts);
continue;
}
let prev_ts = last_ts.get(username).copied().unwrap_or(0);
if *ts <= prev_ts {
continue;
}
last_ts.insert(username.clone(), *ts);
let display = names_guard.display(username);
let is_group = username.contains("@chatroom");
let summary = decompress_or_str(summary_bytes);
let summary = if summary.contains(":\n") {
summary.splitn(2, ":\n").nth(1).unwrap_or(&summary).to_string()
} else {
summary
};
let sender_display = if !sender.is_empty() {
names_guard.map.get(sender).cloned().unwrap_or_else(|| sender.clone())
} else {
String::new()
};
let event = WatchEvent {
event: "message".into(),
time: Some(fmt_time_hhmm(*ts)),
chat: Some(display),
username: Some(username.clone()),
is_group: Some(is_group),
sender: Some(sender_display),
content: Some(summary),
msg_type: Some(fmt_type(*msg_type)),
timestamp: Some(*ts),
};
let _ = tx.send(event);
}
if !initialized {
initialized = true;
}
}
});
}
fn mtime_f64(path: &std::path::Path) -> Option<f64> {
std::fs::metadata(path)
.and_then(|m| m.modified())
.ok()
.map(|t| t.duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs_f64())
}
fn decompress_or_str(data: &[u8]) -> String {
if data.is_empty() {
return String::new();
}
if let Ok(dec) = zstd::decode_all(data) {
if let Ok(s) = String::from_utf8(dec) {
return s;
}
}
String::from_utf8_lossy(data).into_owned()
}
fn fmt_time_hhmm(ts: i64) -> String {
use chrono::{Local, TimeZone};
Local.timestamp_opt(ts, 0)
.single()
.map(|dt| dt.format("%H:%M").to_string())
.unwrap_or_else(|| ts.to_string())
}