270 lines
14 KiB
Python
270 lines
14 KiB
Python
"""Page-driven local Mac connection, with isolated backups and atomic import."""
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import platform
|
||
import re
|
||
import secrets
|
||
import shlex
|
||
import shutil
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
from . import store
|
||
from .cipher import PAGE, decrypt_database, matches
|
||
from .importer import import_rows
|
||
from .mac_keys import scan_keys, wechat_pid
|
||
from .wechat_native import message_rows
|
||
|
||
|
||
def parse_keys(value):
|
||
keys = set()
|
||
def walk(item):
|
||
if isinstance(item, dict):
|
||
for child in item.values():
|
||
walk(child)
|
||
elif isinstance(item, list):
|
||
for child in item:
|
||
walk(child)
|
||
elif isinstance(item, str):
|
||
text = item.strip()
|
||
if text.lower().startswith("x'") and text.endswith("'"):
|
||
text = text[2:-1]
|
||
if re.fullmatch('[a-fA-F0-9]{64}(?:[a-fA-F0-9]{32})?', text):
|
||
keys.add(text[:64].lower())
|
||
walk(value)
|
||
if not keys or len(keys) > 2000:
|
||
raise ValueError('密钥文件应包含 1～2000 个 64/96 位十六进制 key 字段；只在本机内存处理')
|
||
return [bytes.fromhex(key) for key in keys]
|
||
|
||
|
||
class MacConnection:
|
||
def __init__(self, data):
|
||
self.data = Path(data)
|
||
self.accounts, self.keys = {}, []
|
||
self.job = None
|
||
self.lock = threading.RLock()
|
||
deps = self.data / 'mac_deps'
|
||
if deps.is_dir():
|
||
sys.path.insert(0, str(deps))
|
||
|
||
def require_mac(self):
|
||
if platform.system() != 'Darwin':
|
||
raise ValueError('请在目标 Mac 上启动本工具。当前服务不在 Mac，无法读取另一台电脑上的微信。')
|
||
|
||
def status(self):
|
||
with self.lock:
|
||
return {'available': platform.system() == 'Darwin', 'architecture': platform.machine(),
|
||
'accounts': [self.describe(key, path) for key,path in self.accounts.items()],
|
||
'keys_ready': bool(self.keys), 'job': dict(self.job) if self.job else None}
|
||
|
||
def describe(self, key, path):
|
||
files = self.files(path)
|
||
encrypted, ready, total = 0, 0, 0
|
||
for file in files:
|
||
if file.is_symlink():
|
||
raise ValueError('不支持符号链接数据库')
|
||
with file.open('rb') as handle:
|
||
first = handle.read(PAGE)
|
||
total += file.stat().st_size
|
||
if first.startswith(b'SQLite format 3\x00'):
|
||
ready += 1
|
||
else:
|
||
encrypted += 1
|
||
ready += any(matches(first, candidate) for candidate in self.keys)
|
||
return {'id': key, 'name': path.parent.name, 'path': str(path), 'databases': len(files),
|
||
'bytes': total, 'encrypted': encrypted, 'readable': ready}
|
||
|
||
@staticmethod
|
||
def files(root, biz=False):
|
||
paths = [root/'contact/contact.db']
|
||
paths += sorted((root/'message').glob('message_[0-9]*.db'))
|
||
if biz:
|
||
paths += sorted((root/'message').glob('biz_message_[0-9]*.db'))
|
||
return [p for p in paths if p.is_file() and (p.name == 'contact.db' or re.fullmatch(r'(?:biz_)?message_\d+\.db', p.name))]
|
||
|
||
def scan(self, pick=False):
|
||
self.require_mac()
|
||
base = Path.home()/'Library/Containers/com.tencent.xinWeChat/Data/Documents/xwechat_files'
|
||
roots = [base, Path.home()/'Documents/xwechat_files']
|
||
if pick:
|
||
result = subprocess.run(['/usr/bin/osascript', '-e', 'POSIX path of (choose folder with prompt "选择微信 xwechat_files 或账号的 db_storage 目录")'], capture_output=True, text=True, timeout=180)
|
||
if result.returncode:
|
||
raise ValueError('目录选择已取消或系统未授予访问权限')
|
||
roots = [Path(result.stdout.strip())]
|
||
found = {}
|
||
for root in roots:
|
||
candidates = [root] if root.name == 'db_storage' else list(root.glob('*/db_storage'))
|
||
for path in candidates:
|
||
if path.is_dir() and not path.is_symlink() and (path/'message').is_dir() and (path/'contact').is_dir():
|
||
if any(p.is_symlink() for p in path.rglob('*')):
|
||
raise ValueError('选中的数据目录包含符号链接，请选择实际数据库目录')
|
||
resolved = path.resolve()
|
||
found[hashlib.sha256(str(resolved).encode()).hexdigest()[:16]] = resolved
|
||
with self.lock:
|
||
if self.job and self.job['status'] == 'running':
|
||
raise ValueError('正在执行接入任务，请稍后再扫描')
|
||
self.accounts = found
|
||
return self.status()
|
||
|
||
def load_keys(self, value):
|
||
keys = parse_keys(value)
|
||
with self.lock:
|
||
if self.job and self.job['status'] == 'running':
|
||
raise ValueError('正在执行接入任务，暂不能替换密钥')
|
||
self.keys = keys
|
||
return self.status()
|
||
|
||
def start(self, action, body):
|
||
self.require_mac()
|
||
with self.lock:
|
||
if self.job and self.job['status'] == 'running':
|
||
raise ValueError('已有接入任务正在运行')
|
||
if action == 'import' and body.get('account') not in self.accounts:
|
||
raise ValueError('请先扫描并选择账号')
|
||
if action not in ('capture', 'import', 'dependencies'):
|
||
raise ValueError('未知接入操作')
|
||
job_id = secrets.token_hex(12)
|
||
self.job = {'id': job_id, 'action': action, 'status': 'running', 'stage': '准备中', 'error': ''}
|
||
path = self.accounts.get(body.get('account'))
|
||
key_snapshot = list(self.keys)
|
||
threading.Thread(target=self.work, args=(job_id, action, dict(body), path, key_snapshot), daemon=True).start()
|
||
return {'id': job_id}
|
||
|
||
def stage(self, text):
|
||
with self.lock:
|
||
self.job['stage'] = text
|
||
|
||
def work(self, job_id, action, body, path, keys):
|
||
try:
|
||
if action == 'capture':
|
||
self.stage('正在读取已登录微信中的候选密钥；最多扫描 45 秒')
|
||
if body.get('administrator'):
|
||
# No credentials in argv or files. System owns the password dialog.
|
||
script_path = Path(__file__).with_name('mac_keys.py')
|
||
command = shlex.join([sys.executable, '-I', str(script_path)])
|
||
script = 'do shell script ' + json.dumps(command, ensure_ascii=False) + ' with administrator privileges'
|
||
result = subprocess.run(['/usr/bin/osascript', '-e', script], capture_output=True, text=True, timeout=180)
|
||
if result.returncode:
|
||
raise ValueError('管理员授权已取消或系统拒绝访问；没有修改微信签名或系统设置')
|
||
payload = json.loads(result.stdout)
|
||
if payload.get('error'):
|
||
raise ValueError(payload['error'])
|
||
values = payload.get('keys', [])
|
||
else:
|
||
values = scan_keys()
|
||
with self.lock:
|
||
self.keys = parse_keys(values)
|
||
result = {'message': '候选密钥已保存在运行内存。请查看账号的匹配数量；退出微信后点击备份并导入。'}
|
||
elif action == 'dependencies':
|
||
self.stage('从 PyPI 安装压缩文本支持；不会发送聊天或密钥')
|
||
target = self.data/'mac_deps'
|
||
done = subprocess.run([sys.executable, '-m', 'pip', 'install', '--disable-pip-version-check', '--no-input', '--only-binary=:all:',
|
||
'--target', str(target), 'zstandard==0.23.0'], capture_output=True, timeout=180)
|
||
if done.returncode:
|
||
raise ValueError('安装未成功，请检查网络和 Python pip 支持。没有发送聊天或密钥。')
|
||
if str(target) not in sys.path:
|
||
sys.path.insert(0, str(target))
|
||
result = {'message': '文本解压支持已就绪'}
|
||
else:
|
||
result = self.import_account(path, keys, bool(body.get('include_biz')))
|
||
with self.lock:
|
||
self.job.update(status='complete', stage='已完成', result=result)
|
||
except Exception as exc:
|
||
# Never serialize subprocess stderr, process memory or raw key data.
|
||
message = str(exc) if isinstance(exc, ValueError) else '接入失败，请检查文件访问权限、磁盘空间与数据格式。已有记录未被覆盖。'
|
||
with self.lock:
|
||
self.job.update(status='failed', error=message, stage='需要处理')
|
||
|
||
def import_account(self, source, keys, biz):
|
||
if wechat_pid():
|
||
raise ValueError('请先正常退出微信，再点击“备份并导入”。已取得的密钥会保留在本工具运行内存中。')
|
||
files = self.files(source, biz)
|
||
if not (source/'contact/contact.db').is_file() or not any(re.fullmatch(r'message_[0-9]+\.db', p.name) for p in files):
|
||
raise ValueError('联系人或消息库缺失，未开始导入')
|
||
self.data.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||
total_size = sum(p.stat().st_size + (Path(str(p)+'-wal').stat().st_size if Path(str(p)+'-wal').exists() else 0) for p in files)
|
||
if shutil.disk_usage(self.data).free < total_size * 4:
|
||
raise ValueError('磁盘可用空间不足，请预留至少约四倍所选数据库与 WAL 的大小，再开始备份导入')
|
||
job_root = self.data/'mac_sources'/secrets.token_hex(12)
|
||
job_root.mkdir(parents=True, mode=0o700)
|
||
copied, plain = job_root/'backup', job_root/'readable'
|
||
self.stage('正在创建本地数据库副本，保留 WAL；原文件不修改')
|
||
snapshots = {}
|
||
for path in files:
|
||
for suffix in ('', '-wal', '-shm'):
|
||
item = Path(str(path)+suffix)
|
||
if item.exists():
|
||
if item.is_symlink():
|
||
raise ValueError('源目录出现符号链接，已停止备份')
|
||
if suffix != '-shm':
|
||
st = item.stat()
|
||
snapshots[item] = (st.st_size, st.st_mtime_ns)
|
||
target = copied/item.relative_to(source)
|
||
target.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||
shutil.copyfile(item, target)
|
||
target.chmod(0o600)
|
||
if wechat_pid() or any(not p.exists() or (p.stat().st_size,p.stat().st_mtime_ns) != st for p,st in snapshots.items()):
|
||
raise ValueError('备份期间微信数据发生变化，请保持微信退出后重试；没有导入不稳定副本')
|
||
for index, original in enumerate(files, 1):
|
||
self.stage(f'解密并检查数据库 {index}/{len(files)}：{original.name}')
|
||
rel = original.relative_to(source)
|
||
src, dst = copied/rel, plain/rel
|
||
dst.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||
with src.open('rb') as handle:
|
||
first = handle.read(PAGE)
|
||
if first.startswith(b'SQLite format 3\x00'):
|
||
inp = sqlite3.connect(src.resolve().as_uri()+'?mode=ro', uri=True)
|
||
out = sqlite3.connect(dst)
|
||
try:
|
||
inp.backup(out)
|
||
finally:
|
||
out.close()
|
||
inp.close()
|
||
dst.chmod(0o600)
|
||
else:
|
||
key = next((candidate for candidate in keys if matches(first, candidate)), None)
|
||
if not key:
|
||
raise ValueError('部分库缺少匹配密钥：'+original.name+'。请重新登录并打开相关聊天后读取密钥，再退出微信重试。副本已保留。')
|
||
decrypt_database(src, dst, key)
|
||
self.stage('关联联系人、会话和发送人，解析消息；校验完成后统一写入')
|
||
staging_root = job_root/'staging'
|
||
staging = store.connect(staging_root)
|
||
total = 0
|
||
try:
|
||
chunk = []
|
||
for row in message_rows(plain, source.parent.name, biz):
|
||
chunk.append(row)
|
||
if len(chunk) >= 1000:
|
||
result = import_rows(staging, chunk, 'Mac 微信本地接入', {})
|
||
if result['rejected']:
|
||
raise ValueError('部分消息无法标准化，已停止整次导入；请检查版本适配')
|
||
total += len(chunk)
|
||
self.stage(f'已解析 {total} 条消息，仍在校验（尚未写入主库）')
|
||
chunk = []
|
||
if chunk:
|
||
result = import_rows(staging, chunk, 'Mac 微信本地接入', {})
|
||
if result['rejected']:
|
||
raise ValueError('部分消息无法标准化，已停止整次导入')
|
||
total += len(chunk)
|
||
if not total:
|
||
raise ValueError('没有解析出消息；该版本消息表结构需要进一步适配')
|
||
finally:
|
||
staging.close()
|
||
db = store.connect(self.data)
|
||
try:
|
||
db.execute('ATTACH DATABASE ? AS incoming', (str(staging_root/'workbench.sqlite3'),))
|
||
fields = [row[1] for row in db.execute('PRAGMA main.table_info(messages)') if row[1] != 'id']
|
||
columns = ','.join(fields)
|
||
before = db.total_changes
|
||
db.execute('INSERT OR IGNORE INTO main.messages ('+columns+') SELECT '+columns+' FROM incoming.messages')
|
||
added = db.total_changes - before
|
||
db.commit()
|
||
finally:
|
||
db.close()
|
||
return {'message': '导入完成，可进入“发现知识需求”按时间段分析群聊和私聊。',
|
||
'parsed': total, 'added': added, 'duplicates': total-added, 'backup': str(job_root), 'include_biz': biz}
|