"""Fresh-install database initialization with synthetic files only.""" import json import os import sqlite3 import subprocess import sys import tempfile import unittest from pathlib import Path from unittest import mock import wxwork_db as db import wxwork_local_setup as setup import wxwork_message_browser as browser from test_crypto_roundtrip import PAGE, encrypt_page class LocalDatabaseSetupTest(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) self.data = self.root / "user-data" self.source = self.root / "custom-drive" / "WXWork" self.database = self.source / "100" / "Data" / "message.db" self.database.parent.mkdir(parents=True) connection = sqlite3.connect(self.database) connection.execute("PRAGMA page_size=4096") connection.execute("CREATE TABLE message_table(sender_id TEXT, conversation_id TEXT, content_type INT, send_time INT, content TEXT)") connection.execute("INSERT INTO message_table VALUES('200','M:200',2,123,'test message')") connection.commit() connection.close() for patcher in (mock.patch.object(db, "_SCRIPT_DIR", str(self.data)), mock.patch.object(db, "KEYS_FILE", str(self.data / "wxwork_keys.json")), mock.patch.object(setup, "application_data_dir", return_value=self.data), mock.patch.object(browser, "application_data_dir", return_value=self.data), mock.patch.dict(os.environ, {"WECOM_ARCHIVE_SOURCE_DIR": ""})): patcher.start() self.addCleanup(patcher.stop) setup.select_source_directory(str(self.source)) def encrypt(self): key = bytes(range(1, 17)) raw = self.database.read_bytes() self.database.write_bytes(b"".join(encrypt_page(key, i // PAGE + 1, raw[i:i+PAGE]) for i in range(0, len(raw), PAGE))) return key.hex() def snapshot(self): return browser.load_browser_snapshot(cache_roots=[self.data / "wxwork_decrypted"], refresh_cache=True) def test_saved_custom_path_is_used_after_install(self): self.assertEqual(db.detect_wxwork_dir(), str(self.source)) state = self.snapshot() self.assertEqual(state["error"], "") self.assertEqual(state["messageCount"], 1) self.assertEqual(state["accountCount"], 1) self.assertTrue((self.data / "wxwork_decrypted/100/message.db").exists()) self.assertFalse((self.source / "wxwork_keys.json").exists()) def test_missing_key_is_not_reported_as_missing_database(self): self.encrypt() state = self.snapshot() self.assertIn("已找到企业微信数据库", state["error"]) self.assertIn("密钥", state["error"]) self.assertEqual(state["accountCount"], 0) def test_import_then_refresh_reads_encrypted_message(self): key = self.encrypt() before = self.database.read_bytes() imported = self.root / "keys.json" imported.write_text(json.dumps({"global_key": key}), encoding="utf-8-sig") self.assertEqual(setup.import_keys_file(str(imported)), 1) state = self.snapshot() self.assertEqual(state["error"], "") self.assertEqual(state["messages"][0]["content"], "test message") self.assertEqual(self.database.read_bytes(), before) self.assertEqual(db.load_keys(), {"100": key}) def test_wrong_key_does_not_overwrite_saved_key(self): key = self.encrypt() saved = self.data / "wxwork_keys.json" saved.write_text(json.dumps({"keys": {"100": key}}), encoding="utf-8") before = saved.read_bytes() imported = self.root / "wrong.json" imported.write_text(json.dumps({"global_key": "00" * 16}), encoding="utf-8") with self.assertRaisesRegex(ValueError, "无法解密"): setup.import_keys_file(str(imported)) self.assertEqual(saved.read_bytes(), before) def test_invalid_keys_produce_actionable_error_without_secret(self): for payload in ([], {"keys": []}, {"keys": {"100": "private-invalid-value"}}): with self.subTest(payload=payload): filename = self.root / "bad.json" filename.write_text(json.dumps(payload), encoding="utf-8") with self.assertRaisesRegex(ValueError, "密钥文件无效") as caught: db.load_keys(filename) self.assertNotIn("private-invalid-value", str(caught.exception)) def test_directory_choice_preserves_other_settings(self): config = self.data / "wxwork_gui_config.json" config.write_text(json.dumps({"other": True}), encoding="utf-8") selected = setup.select_source_directory(str(self.database.parent)) self.assertEqual(selected, str(self.source)) self.assertTrue(json.loads(config.read_text())["other"]) with self.assertRaises(ValueError): setup.select_source_directory(str(self.root / "absent")) self.assertEqual(db.detect_wxwork_dir(), str(self.source)) def test_frozen_key_worker_uses_executable_without_python_script(self): with mock.patch.object(setup, "is_frozen", return_value=True), mock.patch.object(setup.subprocess, "run", return_value=mock.Mock(returncode=0)) as run: setup.acquire_local_keys() self.assertEqual(run.call_args.args[0], [sys.executable, "--wechat-db-key-worker"]) self.assertEqual(run.call_args.kwargs["timeout"], 60) self.assertEqual(run.call_args.kwargs["stdout"], subprocess.DEVNULL) def test_key_worker_timeout_and_no_process_are_actionable(self): for failure, message in ((subprocess.TimeoutExpired("test", 60), "超时"), (mock.Mock(returncode=11), "登录企业微信")): with mock.patch.object(setup.subprocess, "run", side_effect=[failure]): with self.assertRaisesRegex(ValueError, message): setup.acquire_local_keys() def test_key_worker_only_persists_verified_keys(self): import wxwork_key key = self.encrypt() with mock.patch.object(wxwork_key, "find_wxwork_pids", return_value=[123]), mock.patch.object(wxwork_key, "extract_keys_from_running", return_value={str(self.database): [key]}): self.assertEqual(setup.key_worker_main(), 0) self.assertEqual(db.load_keys(), {"100": key}) def test_database_notice_keeps_specific_failure(self): with mock.patch.object(browser, "_refresh_cache", return_value=("", "请先选择数据目录")): state = browser.load_browser_snapshot(cache_roots=[self.root / "empty"]) self.assertEqual(state["error"], "请先选择数据目录") def auto_snapshot(self): return browser.load_browser_snapshot(cache_roots=[self.data / "wxwork_decrypted"], auto_initialize=True) def test_first_install_auto_discovers_and_decrypts_without_manual_input(self): key = self.encrypt() before = self.database.read_bytes() (self.data / "wxwork_gui_config.json").unlink() with mock.patch.object(db, "_windows_documents_dirs", return_value=[str(self.source.parent)]), \ mock.patch.object(setup, "acquire_local_keys", side_effect=lambda: setup._save_verified_keys({"auto": key}, str(self.source))) as acquire: state = self.auto_snapshot() acquire.assert_called_once_with() self.assertEqual(state["error"], "") self.assertFalse(state["manualSetupRequired"]) self.assertEqual(state["messages"][0]["content"], "test message") self.assertEqual(self.database.read_bytes(), before) self.assertEqual(db.detect_wxwork_dir(), str(self.source)) def test_existing_key_and_plain_database_do_not_scan_processes(self): for encrypted in (False, True): with self.subTest(encrypted=encrypted): if encrypted: key = self.encrypt() setup._save_verified_keys({"saved": key}, str(self.source)) with mock.patch.object(setup, "acquire_local_keys") as acquire: state = self.auto_snapshot() acquire.assert_not_called() self.assertEqual(state["error"], "") self.assertFalse(state["manualSetupRequired"]) def test_stale_or_malformed_key_automatically_recovers(self): key = self.encrypt() for payload in ({"keys": {"100": "00" * 16}}, {"keys": "invalid"}): with self.subTest(payload=payload): saved = self.data / "wxwork_keys.json" saved.write_text(json.dumps(payload), encoding="utf-8") cached = self.data / "wxwork_decrypted/100/message.db" if cached.exists(): cached.unlink() with mock.patch.object(setup, "acquire_local_keys", side_effect=lambda: setup._save_verified_keys({"auto": key}, str(self.source))) as acquire: state = self.auto_snapshot() acquire.assert_called_once_with() self.assertEqual(state["error"], "") self.assertFalse(state["manualSetupRequired"]) def test_auto_failure_exposes_manual_fallback_and_refresh_retries(self): key = self.encrypt() with mock.patch.object(setup, "acquire_local_keys", side_effect=ValueError("请先登录企业微信")) as acquire: state = self.auto_snapshot() acquire.assert_called_once_with() self.assertTrue(state["manualSetupRequired"]) self.assertIn("登录企业微信", state["error"]) self.assertEqual(state["accountCount"], 0) with mock.patch.object(setup, "acquire_local_keys", side_effect=lambda: setup._save_verified_keys({"auto": key}, str(self.source))): state = self.auto_snapshot() self.assertFalse(state["manualSetupRequired"]) self.assertEqual(state["error"], "") def test_missing_directory_falls_back_without_scanning_processes(self): with mock.patch.object(browser, "detect_wxwork_dir", return_value=None), mock.patch.object(setup, "acquire_local_keys") as acquire: state = self.auto_snapshot() acquire.assert_not_called() self.assertTrue(state["manualSetupRequired"]) self.assertIn("数据目录", state["error"]) def test_scanner_success_without_usable_key_is_still_a_failure(self): self.encrypt() with mock.patch.object(setup, "acquire_local_keys") as acquire: state = self.auto_snapshot() acquire.assert_called_once_with() self.assertTrue(state["manualSetupRequired"]) self.assertIn("仍有数据库无法解密", state["error"]) def test_corrupt_fresh_cache_is_rebuilt_without_scanning(self): key = self.encrypt() setup._save_verified_keys({"saved": key}, str(self.source)) cached = self.data / "wxwork_decrypted/100/message.db" cached.parent.mkdir(parents=True) cached.write_bytes(b"bad-cache" * 1024) with mock.patch.object(setup, "acquire_local_keys") as acquire: state = self.auto_snapshot() acquire.assert_not_called() self.assertEqual(state["error"], "") self.assertEqual(state["messageCount"], 1) def test_partial_accounts_trigger_auto_acquisition_and_manual_fallback(self): key = self.encrypt() second = self.source / "101/Data/message.db" second.parent.mkdir(parents=True) second.write_bytes(self.database.read_bytes()) # 第一账号已解密，第二账号仍未取得密钥。 setup._save_verified_keys({"saved": key}, str(self.source)) self.snapshot() cached = self.data / "wxwork_decrypted/101/message.db" cached.unlink() (self.data / "wxwork_keys.json").write_text('{"keys":{}}') with mock.patch.object(setup, "acquire_local_keys", side_effect=ValueError("未获取到可用密钥")) as acquire: state = self.auto_snapshot() acquire.assert_called_once_with() self.assertTrue(state["manualSetupRequired"]) self.assertEqual(state["accountCount"], 1) self.assertIn("未获取到可用密钥", state["warning"]) def test_onedrive_directory_is_found_without_saved_config(self): (self.data / "wxwork_gui_config.json").unlink() cloud_source = self.root / "OneDrive/Documents/WXWork/100/Data/message.db" cloud_source.parent.mkdir(parents=True) cloud_source.write_bytes(self.database.read_bytes()) with mock.patch.object(db, "_windows_documents_dirs", return_value=[]), \ mock.patch.dict(os.environ, {"OneDrive": str(self.root / "OneDrive")}): self.assertEqual(db.detect_wxwork_dir(), str(cloud_source.parents[2])) def test_registry_documents_expands_environment_variables(self): import winreg with mock.patch.object(winreg, "OpenKey"), \ mock.patch.object(winreg, "QueryValueEx", return_value=("%WECOM_TEST_DOCS%", winreg.REG_EXPAND_SZ)), \ mock.patch.dict(os.environ, {"WECOM_TEST_DOCS": str(self.source.parent)}): self.assertEqual(db._windows_documents_dirs(), [str(self.source.parent)] * 2)