kefu/wechat_rpa/wxwork_local_setup.py

"""安装版的本机数据库初始化；密钥和聊天数据均不随程序分发。"""
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)