chuan/WeChatDataAnalysis
1
0
Fork 0
mirror of https://github.com/LifeArchiveProject/WeChatDataAnalysis.git synced 2026-06-18 15:54:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add logging and sync guards around chat search decrypt flow

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
2977094657
2026-04-04 20:02:43 +08:00
Unverified
parent 4ba2c75332
commit 624853f483
4 changed files with 510 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/wechat_decrypt_tool/chat_realtime_autosync.py
+67
View File
@@ -123,9 +123,62 @@ class ChatRealtimeAutoSyncService:
self._mu = threading.Lock()
self._states: dict[str, _AccountState] = {}
self._paused_accounts: dict[str, int] = {}
self._stop = threading.Event()
self._thread: Optional[threading.Thread] = None
def _is_account_paused_locked(self, account: str) -> bool:
key = str(account or "").strip()
if not key:
return False
return int(self._paused_accounts.get(key) or 0) > 0
def is_account_paused(self, account: str) -> bool:
with self._mu:
return self._is_account_paused_locked(account)
def pause_account(self, account: str, reason: str = "") -> int:
key = str(account or "").strip()
if not key:
return 0
with self._mu:
depth = int(self._paused_accounts.get(key) or 0) + 1
self._paused_accounts[key] = depth
st = self._states.get(key)
if st is not None:
st.due_at = 0.0
logger.info(
"[realtime-autosync] pause account=%s reason=%s depth=%s",
key,
str(reason or "").strip() or "-",
int(depth),
)
return depth
def resume_account(self, account: str, reason: str = "") -> int:
key = str(account or "").strip()
if not key:
return 0
with self._mu:
current = int(self._paused_accounts.get(key) or 0)
if current <= 1:
self._paused_accounts.pop(key, None)
depth = 0
else:
depth = current - 1
self._paused_accounts[key] = depth
logger.info(
"[realtime-autosync] resume account=%s reason=%s depth=%s",
key,
str(reason or "").strip() or "-",
int(depth),
)
return depth
def start(self) -> None:
if not self._enabled:
logger.info("[realtime-autosync] disabled by env WECHAT_TOOL_REALTIME_AUTOSYNC=0")
@@ -188,6 +241,12 @@ class ChatRealtimeAutoSyncService:
if self._stop.is_set():
break
if self.is_account_paused(acc):
with self._mu:
st = self._states.setdefault(acc, _AccountState())
st.due_at = 0.0
continue
try:
account_dir = _resolve_account_dir(acc)
except HTTPException:
@@ -238,6 +297,9 @@ class ChatRealtimeAutoSyncService:
for acc, st in self._states.items():
if running >= int(self._workers):
break
if self._is_account_paused_locked(acc):
st.due_at = 0.0
continue
if st.due_at <= 0 or st.due_at > now:
continue
if st.thread is not None and st.thread.is_alive():
@@ -278,6 +340,9 @@ class ChatRealtimeAutoSyncService:
try:
if self._stop.is_set() or (not account):
return
if self.is_account_paused(account):
logger.info("[realtime-autosync] sync skipped account=%s reason=paused", account)
return
res = self._sync_account(account)
inserted = int((res or {}).get("inserted_total") or (res or {}).get("insertedTotal") or 0)
synced = int((res or {}).get("synced") or (res or {}).get("sessionsSynced") or 0)
@@ -297,6 +362,8 @@ class ChatRealtimeAutoSyncService:
account = str(account or "").strip()
if not account:
return {"status": "skipped", "reason": "missing account"}
if self.is_account_paused(account):
return {"status": "skipped", "reason": "paused"}
try:
account_dir = _resolve_account_dir(account)
src/wechat_decrypt_tool/routers/chat.py
+32
View File
@@ -7071,13 +7071,26 @@ async def get_chat_messages_around(
if after > 200:
after = 200
trace_id = f"msg-around-{int(time.time() * 1000)}-{threading.get_ident()}"
logger.info(
"[%s] chat messages around start account=%s username=%s anchor_id=%s before=%s after=%s",
trace_id,
str(account or "").strip(),
str(username or "").strip(),
str(anchor_id or "").strip(),
int(before),
int(after),
)
parts = str(anchor_id).split(":", 2)
if len(parts) != 3:
logger.warning("[%s] chat messages around invalid anchor format anchor_id=%s", trace_id, str(anchor_id or "").strip())
raise HTTPException(status_code=400, detail="Invalid anchor_id.")
anchor_db_stem, anchor_table_name_in, anchor_local_id_str = parts
try:
anchor_local_id = int(anchor_local_id_str)
except Exception:
logger.warning("[%s] chat messages around invalid anchor local_id anchor_id=%s", trace_id, str(anchor_id or "").strip())
raise HTTPException(status_code=400, detail="Invalid anchor_id.")
account_dir = _resolve_account_dir(account)
@@ -7093,6 +7106,13 @@ async def get_chat_messages_around(
anchor_db_path = p
break
if anchor_db_path is None:
logger.warning(
"[%s] chat messages around anchor db missing account=%s username=%s anchor_db=%s",
trace_id,
account_dir.name,
username,
anchor_db_stem,
)
raise HTTPException(status_code=404, detail="Anchor database not found.")
# Open resource DB once (optional), and reuse for all message DBs.
@@ -7491,6 +7511,18 @@ async def get_chat_messages_around(
head_image_db_path=head_image_db_path,
)
logger.info(
"[%s] chat messages around done account=%s username=%s anchor_id=%s canonical_anchor=%s anchor_index=%s returned=%s merged_total=%s",
trace_id,
account_dir.name,
username,
str(anchor_id or "").strip(),
anchor_id_canon,
int(anchor_index),
len(return_messages),
len(merged),
)
return {
"status": "success",
"account": account_dir.name,
src/wechat_decrypt_tool/routers/decrypt.py
+316 -194
View File
@@ -5,12 +5,14 @@ import json
import os
import time
from pathlib import Path
from typing import Any
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel, Field
from starlette.responses import StreamingResponse
from ..app_paths import get_output_databases_dir
from ..chat_realtime_autosync import CHAT_REALTIME_AUTOSYNC
from ..logging_config import get_logger
from ..path_fix import PathFixRoute
from ..key_store import upsert_account_keys_in_store
@@ -21,6 +23,96 @@ logger = get_logger(__name__)
router = APIRouter(route_class=PathFixRoute)
def _normalize_decrypt_guard_accounts(accounts: Any) -> list[str]:
if not accounts:
return []
out: list[str] = []
seen: set[str] = set()
for account in accounts:
key = str(account or "").strip()
if (not key) or (key in seen):
continue
seen.add(key)
out.append(key)
out.sort()
return out
def _resolve_decrypt_guard_accounts(db_storage_path: str) -> list[str]:
try:
scan_result = scan_account_databases_from_path(db_storage_path)
except Exception:
logger.exception("[decrypt] pre-scan accounts failed db_storage_path=%s", db_storage_path)
return []
if scan_result.get("status") == "error":
return []
return _normalize_decrypt_guard_accounts((scan_result.get("account_databases") or {}).keys())
def _get_realtime_sync_all_lock(account: str):
from .chat import _realtime_sync_all_lock
return _realtime_sync_all_lock(account)
def _release_decrypt_account_guards(guards: list[tuple[str, Any]], *, reason: str) -> None:
for account, lock in reversed(list(guards or [])):
try:
lock.release()
logger.info("[decrypt] released realtime sync_all lock account=%s reason=%s", account, reason)
except Exception:
logger.exception("[decrypt] release realtime sync_all lock failed account=%s reason=%s", account, reason)
try:
CHAT_REALTIME_AUTOSYNC.resume_account(account, reason=reason)
logger.info("[decrypt] resumed realtime autosync for account after decrypt account=%s reason=%s", account, reason)
except Exception:
logger.exception(
"[decrypt] resume realtime autosync failed account=%s reason=%s",
account,
reason,
)
def _acquire_decrypt_account_guards(accounts: Any, *, reason: str) -> list[tuple[str, Any]]:
guards: list[tuple[str, Any]] = []
for account in _normalize_decrypt_guard_accounts(accounts):
paused = False
try:
CHAT_REALTIME_AUTOSYNC.pause_account(account, reason=reason)
paused = True
logger.info("[decrypt] paused realtime autosync for account during decrypt account=%s reason=%s", account, reason)
lock = _get_realtime_sync_all_lock(account)
logger.info("[decrypt] waiting realtime sync_all lock account=%s reason=%s", account, reason)
lock.acquire()
logger.info("[decrypt] acquired realtime sync_all lock account=%s reason=%s", account, reason)
guards.append((account, lock))
except Exception:
if paused:
try:
CHAT_REALTIME_AUTOSYNC.resume_account(account, reason=f"{reason}:acquire_failed")
logger.info(
"[decrypt] resumed realtime autosync after guard acquire failure account=%s reason=%s",
account,
f"{reason}:acquire_failed",
)
except Exception:
logger.exception(
"[decrypt] resume realtime autosync after guard acquire failure failed account=%s reason=%s",
account,
reason,
)
_release_decrypt_account_guards(guards, reason=reason)
raise
return guards
class DecryptRequest(BaseModel):
"""解密请求模型"""
@@ -48,11 +140,16 @@ async def decrypt_databases(request: DecryptRequest):
logger.warning(f"密钥格式无效: 长度={len(request.key) if request.key else 0}")
raise HTTPException(status_code=400, detail="密钥格式无效，必须是64位十六进制字符串")
# 使用新的解密API
results = decrypt_wechat_databases(
db_storage_path=request.db_storage_path,
key=request.key,
)
guard_accounts = _resolve_decrypt_guard_accounts(request.db_storage_path)
guards = _acquire_decrypt_account_guards(guard_accounts, reason="decrypt:post")
try:
# 使用新的解密API
results = decrypt_wechat_databases(
db_storage_path=request.db_storage_path,
key=request.key,
)
finally:
_release_decrypt_account_guards(guards, reason="decrypt:post")
if results["status"] == "error":
logger.error(f"解密失败: {results['message']}")
@@ -148,182 +245,99 @@ async def decrypt_databases_stream(
account_sources = scan_result.get("account_sources", {})
total_databases = sum(len(dbs) for dbs in account_databases.values())
yield _sse({"type": "start", "total": total_databases, "message": f"开始解密 {total_databases} 个数据库"})
await asyncio.sleep(0)
# 3) Init output dir & decryptor.
base_output_dir = get_output_databases_dir()
base_output_dir.mkdir(parents=True, exist_ok=True)
decrypt_guards: list[tuple[str, Any]] = []
try:
decryptor = WeChatDatabaseDecryptor(k)
except ValueError as e:
yield _sse({"type": "error", "message": f"密钥错误: {e}"})
return
# 4) Decrypt per account, stream progress.
success_count = 0
fail_count = 0
processed_files: list[str] = []
failed_files: list[str] = []
account_results: dict = {}
diagnostic_warning_count = 0
overall_current = 0
for account, dbs in account_databases.items():
account_output_dir = base_output_dir / account
account_output_dir.mkdir(parents=True, exist_ok=True)
# Save a hint for later UI (same as non-stream endpoint).
try:
source_info = account_sources.get(account, {})
source_db_storage_path = str(source_info.get("db_storage_path") or p)
wxid_dir = str(source_info.get("wxid_dir") or "")
(account_output_dir / "_source.json").write_text(
json.dumps({"db_storage_path": source_db_storage_path, "wxid_dir": wxid_dir}, ensure_ascii=False, indent=2),
encoding="utf-8",
)
except Exception:
pass
account_success = 0
account_processed: list[str] = []
account_failed: list[str] = []
account_db_diagnostics: dict[str, dict] = {}
account_diagnostic_warning_count = 0
for db_info in dbs:
if await request.is_disconnected():
return
overall_current += 1
db_path = str(db_info.get("path") or "")
db_name = str(db_info.get("name") or "")
current_file = f"{account}/{db_name}" if account else db_name
# Emit a "processing" event so UI updates immediately for large db files.
yield _sse(
{
"type": "progress",
"current": overall_current,
"total": total_databases,
"success_count": success_count,
"fail_count": fail_count,
"current_file": current_file,
"status": "processing",
"message": "解密中...",
}
)
output_path = account_output_dir / db_name
task = asyncio.create_task(asyncio.to_thread(decryptor.decrypt_database, db_path, str(output_path)))
# Wait with heartbeat (can't yield while awaiting the thread directly).
last_heartbeat = time.time()
while not task.done():
if await request.is_disconnected():
return
now = time.time()
if now - last_heartbeat > 15:
last_heartbeat = now
# SSE comment heartbeat; browsers ignore but keeps proxies alive.
yield ": ping\n\n"
await asyncio.sleep(0.6)
try:
ok = bool(task.result())
except Exception:
ok = False
db_diagnostic = dict(getattr(decryptor, "last_result", {}) or {})
if not db_diagnostic:
db_diagnostic = {
"db_path": str(db_path),
"db_name": str(db_name),
"output_path": str(output_path),
"success": bool(ok),
}
db_diagnostic["account"] = str(account)
account_db_diagnostics[db_name] = db_diagnostic
if (
(not bool(db_diagnostic.get("success", ok)))
or int(db_diagnostic.get("failed_pages") or 0) > 0
or str(db_diagnostic.get("diagnostic_status") or "") != "ok"
):
account_diagnostic_warning_count += 1
if ok:
account_success += 1
success_count += 1
account_processed.append(str(output_path))
processed_files.append(str(output_path))
status = "success"
msg = "解密成功"
else:
account_failed.append(db_path)
failed_files.append(db_path)
fail_count += 1
status = "fail"
msg = "解密失败"
payload = {
"type": "progress",
"current": overall_current,
"total": total_databases,
"success_count": success_count,
"fail_count": fail_count,
"current_file": current_file,
"status": status,
"message": msg,
}
if db_diagnostic:
payload["diagnostic_status"] = str(db_diagnostic.get("diagnostic_status") or "")
payload["page_failures"] = int(db_diagnostic.get("failed_pages") or 0)
if db_diagnostic.get("failed_page_samples"):
payload["failed_page_samples"] = db_diagnostic.get("failed_page_samples")
if db_diagnostic.get("diagnostics"):
payload["diagnostics"] = db_diagnostic.get("diagnostics")
yield _sse(payload)
if overall_current % 5 == 0:
await asyncio.sleep(0)
account_results[account] = {
"total": len(dbs),
"success": account_success,
"failed": len(dbs) - account_success,
"output_dir": str(account_output_dir),
"processed_files": account_processed,
"failed_files": account_failed,
"db_diagnostics": account_db_diagnostics,
"diagnostic_warning_count": int(account_diagnostic_warning_count),
}
diagnostic_warning_count += int(account_diagnostic_warning_count)
# Build cache table (keep behavior consistent with the POST endpoint).
if os.environ.get("WECHAT_TOOL_BUILD_SESSION_LAST_MESSAGE", "1") != "0":
guard_accounts = _normalize_decrypt_guard_accounts(account_databases.keys())
if guard_accounts:
yield _sse(
{
"type": "phase",
"phase": "session_last_message",
"account": account,
"message": "正在构建会话缓存（最后一条消息）...",
"phase": "decrypt_guard",
"message": "正在暂停实时同步并等待解密写锁...",
}
)
await asyncio.sleep(0)
decrypt_guards = await asyncio.to_thread(
_acquire_decrypt_account_guards,
guard_accounts,
reason="decrypt:sse",
)
yield _sse({"type": "start", "total": total_databases, "message": f"开始解密 {total_databases} 个数据库"})
await asyncio.sleep(0)
# 3) Init output dir & decryptor.
base_output_dir = get_output_databases_dir()
base_output_dir.mkdir(parents=True, exist_ok=True)
try:
decryptor = WeChatDatabaseDecryptor(k)
except ValueError as e:
yield _sse({"type": "error", "message": f"密钥错误: {e}"})
return
# 4) Decrypt per account, stream progress.
success_count = 0
fail_count = 0
processed_files: list[str] = []
failed_files: list[str] = []
account_results: dict = {}
diagnostic_warning_count = 0
overall_current = 0
for account, dbs in account_databases.items():
account_output_dir = base_output_dir / account
account_output_dir.mkdir(parents=True, exist_ok=True)
# Save a hint for later UI (same as non-stream endpoint).
try:
from ..session_last_message import build_session_last_message_table
task = asyncio.create_task(
asyncio.to_thread(
build_session_last_message_table,
account_output_dir,
rebuild=True,
include_hidden=True,
include_official=True,
)
source_info = account_sources.get(account, {})
source_db_storage_path = str(source_info.get("db_storage_path") or p)
wxid_dir = str(source_info.get("wxid_dir") or "")
(account_output_dir / "_source.json").write_text(
json.dumps(
{"db_storage_path": source_db_storage_path, "wxid_dir": wxid_dir},
ensure_ascii=False,
indent=2,
),
encoding="utf-8",
)
except Exception:
pass
account_success = 0
account_processed: list[str] = []
account_failed: list[str] = []
account_db_diagnostics: dict[str, dict] = {}
account_diagnostic_warning_count = 0
for db_info in dbs:
if await request.is_disconnected():
return
overall_current += 1
db_path = str(db_info.get("path") or "")
db_name = str(db_info.get("name") or "")
current_file = f"{account}/{db_name}" if account else db_name
# Emit a "processing" event so UI updates immediately for large db files.
yield _sse(
{
"type": "progress",
"current": overall_current,
"total": total_databases,
"success_count": success_count,
"fail_count": fail_count,
"current_file": current_file,
"status": "processing",
"message": "解密中...",
}
)
output_path = account_output_dir / db_name
task = asyncio.create_task(asyncio.to_thread(decryptor.decrypt_database, db_path, str(output_path)))
# Wait with heartbeat (can't yield while awaiting the thread directly).
last_heartbeat = time.time()
while not task.done():
if await request.is_disconnected():
@@ -331,34 +345,142 @@ async def decrypt_databases_stream(
now = time.time()
if now - last_heartbeat > 15:
last_heartbeat = now
# SSE comment heartbeat; browsers ignore but keeps proxies alive.
yield ": ping\n\n"
await asyncio.sleep(0.6)
account_results[account]["session_last_message"] = task.result()
except Exception as e:
account_results[account]["session_last_message"] = {"status": "error", "message": str(e)}
try:
ok = bool(task.result())
except Exception:
ok = False
db_diagnostic = dict(getattr(decryptor, "last_result", {}) or {})
if not db_diagnostic:
db_diagnostic = {
"db_path": str(db_path),
"db_name": str(db_name),
"output_path": str(output_path),
"success": bool(ok),
}
db_diagnostic["account"] = str(account)
account_db_diagnostics[db_name] = db_diagnostic
status = "completed" if success_count > 0 else "failed"
result = {
"status": status,
"total_databases": total_databases,
"success_count": success_count,
"failure_count": total_databases - success_count,
"output_directory": str(base_output_dir.absolute()),
"message": f"解密完成: 成功 {success_count}/{total_databases}",
"processed_files": processed_files,
"failed_files": failed_files,
"account_results": account_results,
"diagnostic_warning_count": int(diagnostic_warning_count),
}
if (
(not bool(db_diagnostic.get("success", ok)))
or int(db_diagnostic.get("failed_pages") or 0) > 0
or str(db_diagnostic.get("diagnostic_status") or "") != "ok"
):
account_diagnostic_warning_count += 1
# Save db key for frontend autofill.
try:
for account in (account_results or {}).keys():
upsert_account_keys_in_store(str(account), db_key=k)
except Exception:
pass
if ok:
account_success += 1
success_count += 1
account_processed.append(str(output_path))
processed_files.append(str(output_path))
status = "success"
msg = "解密成功"
else:
account_failed.append(db_path)
failed_files.append(db_path)
fail_count += 1
status = "fail"
msg = "解密失败"
yield _sse({"type": "complete", **result})
payload = {
"type": "progress",
"current": overall_current,
"total": total_databases,
"success_count": success_count,
"fail_count": fail_count,
"current_file": current_file,
"status": status,
"message": msg,
}
if db_diagnostic:
payload["diagnostic_status"] = str(db_diagnostic.get("diagnostic_status") or "")
payload["page_failures"] = int(db_diagnostic.get("failed_pages") or 0)
if db_diagnostic.get("failed_page_samples"):
payload["failed_page_samples"] = db_diagnostic.get("failed_page_samples")
if db_diagnostic.get("diagnostics"):
payload["diagnostics"] = db_diagnostic.get("diagnostics")
yield _sse(payload)
if overall_current % 5 == 0:
await asyncio.sleep(0)
account_results[account] = {
"total": len(dbs),
"success": account_success,
"failed": len(dbs) - account_success,
"output_dir": str(account_output_dir),
"processed_files": account_processed,
"failed_files": account_failed,
"db_diagnostics": account_db_diagnostics,
"diagnostic_warning_count": int(account_diagnostic_warning_count),
}
diagnostic_warning_count += int(account_diagnostic_warning_count)
# Build cache table (keep behavior consistent with the POST endpoint).
if os.environ.get("WECHAT_TOOL_BUILD_SESSION_LAST_MESSAGE", "1") != "0":
yield _sse(
{
"type": "phase",
"phase": "session_last_message",
"account": account,
"message": "正在构建会话缓存（最后一条消息）...",
}
)
await asyncio.sleep(0)
try:
from ..session_last_message import build_session_last_message_table
task = asyncio.create_task(
asyncio.to_thread(
build_session_last_message_table,
account_output_dir,
rebuild=True,
include_hidden=True,
include_official=True,
)
)
last_heartbeat = time.time()
while not task.done():
if await request.is_disconnected():
return
now = time.time()
if now - last_heartbeat > 15:
last_heartbeat = now
yield ": ping\n\n"
await asyncio.sleep(0.6)
account_results[account]["session_last_message"] = task.result()
except Exception as e:
account_results[account]["session_last_message"] = {"status": "error", "message": str(e)}
status = "completed" if success_count > 0 else "failed"
result = {
"status": status,
"total_databases": total_databases,
"success_count": success_count,
"failure_count": total_databases - success_count,
"output_directory": str(base_output_dir.absolute()),
"message": f"解密完成: 成功 {success_count}/{total_databases}",
"processed_files": processed_files,
"failed_files": failed_files,
"account_results": account_results,
"diagnostic_warning_count": int(diagnostic_warning_count),
}
# Save db key for frontend autofill.
try:
for account in (account_results or {}).keys():
upsert_account_keys_in_store(str(account), db_key=k)
except Exception:
pass
yield _sse({"type": "complete", **result})
finally:
if decrypt_guards:
await asyncio.to_thread(_release_decrypt_account_guards, decrypt_guards, reason="decrypt:sse")
headers = {"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}
return StreamingResponse(generate_progress(), media_type="text/event-stream", headers=headers)