252 lines
13 KiB
Python
252 lines
13 KiB
Python
"""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)
|