"""安装版的本机数据库初始化;密钥和聊天数据均不随程序分发。""" from __future__ import annotations import json import os import subprocess import sys import tempfile import threading from concurrent.futures import Future, TimeoutError as FutureTimeoutError from pathlib import Path import wxwork_db from runtime_paths import application_data_dir, is_frozen def _write_json(path: Path, data: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) name = "" try: with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=path.parent, delete=False) as f: name = f.name json.dump(data, f, ensure_ascii=False, indent=2) os.replace(name, path) finally: if name and os.path.exists(name): os.unlink(name) def select_source_directory(value: str) -> str: selected = Path(value).expanduser().resolve() # 文件选择器允许选择 WXWork、账号目录或账号的 Data 目录。 candidates = [selected, selected.parent, selected.parent.parent] source = next((item for item in candidates if item.is_dir() and any(item.glob("*/Data/message.db"))), None) if source is None: raise ValueError("这个目录没有企业微信数据库,请选择包含账号 Data/message.db 的 WXWork 目录") path = application_data_dir() / "wxwork_gui_config.json" try: settings = json.loads(path.read_text(encoding="utf-8-sig")) except (OSError, ValueError): settings = {} if not isinstance(settings, dict): settings = {} settings["db_dir"] = str(source) _write_json(path, settings) return str(source) def _save_verified_keys(keys: dict, source: str) -> int: databases = sorted(Path(source).glob("*/Data/message.db")) verified = {} for db in databases: for key in set(keys.values()): if wxwork_db.verify_key(bytes.fromhex(key), str(db)): verified[db.parent.parent.name] = key break if not verified: raise ValueError("密钥无法解密当前目录的 message.db,请使用当前电脑、当前企业微信账号对应的密钥") path = application_data_dir() / "wxwork_keys.json" try: existing = wxwork_db.load_keys(path) except ValueError: existing = {} existing.update(verified) _write_json(path, {"keys": existing}) return len(verified) def import_keys_file(filename: str) -> int: source = wxwork_db.detect_wxwork_dir() if not source: raise ValueError("请先选择本机企业微信数据目录,再导入密钥") keys = wxwork_db.load_keys(filename) if not keys: raise ValueError("所选文件没有数据库密钥,请选择当前电脑的 wxwork_keys.json") return _save_verified_keys(keys, source) def key_worker_main() -> int: # 独立进程只读本机企微进程;不登录、不重启企微、不输出密钥。 from wxwork_key import extract_keys_from_running, find_wxwork_pids source = wxwork_db.detect_wxwork_dir() if not source: return 10 if not find_wxwork_pids(): return 11 databases = [str(p) for p in Path(source).glob("*/Data/message.db")] found = extract_keys_from_running(databases, timeout=50.0) keys = {str(index): key for index, key in enumerate({key for values in found.values() for key in values})} if not keys: return 12 try: _save_verified_keys(keys, source) except (OSError, ValueError): return 13 return 0 def _acquire_local_keys_once(timeout: float = 60.0) -> None: # 冻结程序用同一个 EXE 的专用入口,避免尝试运行用户机器上不存在的 python。 command = [sys.executable] if not is_frozen(): command.append(str(Path(__file__).with_name("app_main.py"))) command.append("--wechat-db-key-worker") try: result = subprocess.run(command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=timeout, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0)) except subprocess.TimeoutExpired as exc: raise ValueError("获取密钥超时,已停止读取。请在企业微信登录后重试,或导入已有密钥文件") from exc errors = { 10: "未找到企业微信数据库,请先选择数据目录", 11: "请先启动并登录企业微信,再点击获取本机密钥", 12: "未获取到可用密钥。请在企业微信登录后重试,或导入当前电脑的密钥文件", 13: "密钥校验或保存失败,请检查数据目录后重试", } if result.returncode: raise ValueError(errors.get(result.returncode, "获取密钥失败,请导入当前电脑的密钥文件")) _key_acquisition_lock = threading.Lock() _key_acquisition_inflight = {} def acquire_local_keys(timeout: float = 60.0) -> None: """Share concurrent initialization of the same local data directory only. There is no completed-result cache: a later explicit request scans again. Every caller still validates its own account database after acquisition. """ root = application_data_dir() try: config = json.loads((root / "wxwork_gui_config.json").read_text(encoding="utf-8-sig")) source = str(config.get("db_dir") or "") if isinstance(config, dict) else "" except (OSError, ValueError): source = "" identity = (str(root.resolve()).casefold(), source.casefold(), os.environ.get("WECOM_ARCHIVE_SOURCE_DIR", "").casefold()) with _key_acquisition_lock: future = _key_acquisition_inflight.get(identity) owner = future is None if owner: future = Future() _key_acquisition_inflight[identity] = future if not owner: try: return future.result(timeout=max(0.0, float(timeout))) except FutureTimeoutError as exc: raise ValueError("正在获取本机密钥,等待超时。请稍后重试,或导入已有密钥文件") from exc try: result = _acquire_local_keys_once(timeout) except BaseException as exc: future.set_exception(exc) raise else: future.set_result(result) return result finally: with _key_acquisition_lock: if _key_acquisition_inflight.get(identity) is future: _key_acquisition_inflight.pop(identity, None)