117 lines
4.7 KiB
Python
117 lines
4.7 KiB
Python
"""安装版的本机数据库初始化；密钥和聊天数据均不随程序分发。"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
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(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, "获取密钥失败，请导入当前电脑的密钥文件"))
|