129 lines
6.0 KiB
Python
129 lines
6.0 KiB
Python
"""Adapt readable WeChat 4 tables to the workbench's normalized import format."""
|
||
import hashlib
|
||
import re
|
||
import sqlite3
|
||
import xml.etree.ElementTree as ET
|
||
from pathlib import Path
|
||
|
||
|
||
def quote(name):
|
||
return '"' + name.replace('"', '""') + '"'
|
||
|
||
|
||
def open_db(path):
|
||
db = sqlite3.connect(Path(path).resolve().as_uri() + '?mode=ro&immutable=1', uri=True)
|
||
db.row_factory = sqlite3.Row
|
||
return db
|
||
|
||
|
||
def tables(db):
|
||
return [r[0] for r in db.execute("SELECT name FROM sqlite_master WHERE type='table'")]
|
||
|
||
|
||
def decode(value):
|
||
if value is None:
|
||
return ''
|
||
if not isinstance(value, bytes):
|
||
return str(value)
|
||
if value.startswith(b'\x28\xb5\x2f\xfd'):
|
||
try:
|
||
import zstandard
|
||
value = zstandard.ZstdDecompressor().decompress(value, max_output_size=16*1024**2)
|
||
except ImportError:
|
||
raise ValueError('缺少压缩文本支持，请在连接页面点击“安装文本解压支持”') from None
|
||
try:
|
||
return value.decode('utf-8')
|
||
except UnicodeError:
|
||
raise ValueError('正文含未识别编码，已拒绝把二进制内容当文本') from None
|
||
|
||
|
||
def contacts(root):
|
||
names = {}
|
||
path = root / 'contact/contact.db'
|
||
if not path.exists():
|
||
raise ValueError('缺少联系人库，无法可靠关联会话身份')
|
||
db = open_db(path)
|
||
try:
|
||
for table in tables(db):
|
||
cols = {r[1] for r in db.execute('PRAGMA table_info(' + quote(table) + ')')}
|
||
user = next((v for v in ('username', 'user_name', 'UserName') if v in cols), None)
|
||
if not user:
|
||
continue
|
||
labels = [v for v in ('remark', 'remark_name', 'nick_name', 'nickname', 'NickName') if v in cols]
|
||
for row in db.execute('SELECT ' + ','.join(quote(v) for v in [user]+labels) + ' FROM ' + quote(table)):
|
||
uid = decode(row[0])
|
||
if uid:
|
||
names[uid] = next((decode(row[v]) for v in labels if row[v]), uid)
|
||
finally:
|
||
db.close()
|
||
if not names:
|
||
raise ValueError('联系人库结构未识别；需要为这个版本适配字段')
|
||
return names
|
||
|
||
|
||
def message_rows(root, account, include_biz=False):
|
||
names = contacts(root)
|
||
chats = {hashlib.md5(uid.encode()).hexdigest(): uid for uid in names}
|
||
files = sorted((root / 'message').glob('message_[0-9]*.db'))
|
||
if include_biz:
|
||
files += sorted((root / 'message').glob('biz_message_[0-9]*.db'))
|
||
if not files:
|
||
raise ValueError('没有找到可读消息分片')
|
||
for path in files:
|
||
db = open_db(path)
|
||
try:
|
||
all_tables = tables(db)
|
||
ids = {}
|
||
for table in all_tables:
|
||
if table.lower() == 'name2id':
|
||
cols = {r[1] for r in db.execute('PRAGMA table_info(' + quote(table) + ')')}
|
||
field = next((c for c in ('user_name', 'username') if c in cols), None)
|
||
if field:
|
||
for row in db.execute('SELECT rowid,' + quote(field) + ' FROM ' + quote(table)):
|
||
uid = decode(row[1])
|
||
ids[row[0]] = uid
|
||
chats[hashlib.md5(uid.encode()).hexdigest()] = uid
|
||
for table in all_tables:
|
||
match = re.fullmatch(r'(?:Msg|Chat)_([a-fA-F0-9]{32})', table)
|
||
if not match:
|
||
if table.startswith(('Msg_', 'Chat_')):
|
||
raise ValueError('消息表名称格式未识别，需要版本适配；没有静默跳过')
|
||
continue
|
||
uid = chats.get(match[1].lower())
|
||
if not uid:
|
||
raise ValueError('有消息表无法关联会话身份；请补充完整的联系人/Name2Id 数据后重试')
|
||
cols = {r[1] for r in db.execute('PRAGMA table_info(' + quote(table) + ')')}
|
||
required = {'create_time', 'message_content', 'local_type', 'real_sender_id'}
|
||
if not required <= cols:
|
||
raise ValueError('消息表字段与支持的微信 4.x 格式不同，需要版本适配')
|
||
for raw in db.execute('SELECT rowid AS _rowid_, * FROM ' + quote(table)):
|
||
row = dict(raw)
|
||
typ = int(row['local_type']) & 0xffffffff
|
||
sender = ids.get(row['real_sender_id'])
|
||
if not sender:
|
||
raise ValueError('消息发送人 ID 无法关联，已停止导入，避免错误归属')
|
||
content = decode(row['message_content'])
|
||
if not content and row.get('compress_content'):
|
||
content = decode(row['compress_content'])
|
||
ref = ''
|
||
is_text = typ == 1
|
||
if uid.endswith('@chatroom') and content.startswith(sender + ':\n'):
|
||
content = content[len(sender)+2:]
|
||
if content.startswith('<') and typ == 49:
|
||
try:
|
||
xml = ET.fromstring(content)
|
||
title = xml.findtext('.//appmsg/title') or ''
|
||
ref = xml.findtext('.//refermsg/content') or ''
|
||
is_text = xml.findtext('.//appmsg/type') == '57'
|
||
if title:
|
||
content = title
|
||
except ET.ParseError:
|
||
pass
|
||
source_id = row.get('server_id') or f"{path.name}:{table}:{row.get('local_id',row['_rowid_'])}"
|
||
yield dict(id=str(source_id), conversation=names.get(uid, uid), conversation_id=account+'::'+uid,
|
||
kind='group' if uid.endswith('@chatroom') else 'private', sender=names.get(sender, sender),
|
||
sender_id=sender, timestamp=row['create_time'], content=content or f'［非文字消息 {typ}］',
|
||
message_type='text' if is_text else str(typ), quote=ref)
|
||
finally:
|
||
db.close()
|