# -*- coding: utf-8 -*- """现有企业微信解密库 -> 新聊天归档的独立增量适配器。 不改 ``wxwork_db.py`` 的客服监听行为；这里直接读取它生成的明文 SQLite 副本， 保留群聊、自己发送、非文本类型等归档数据。每个账号按 ``send_time + rowid`` 保存水位，中断后可继续。 """ from __future__ import annotations import argparse import base64 import hashlib import json import sqlite3 from pathlib import Path from typing import Any import admin_backend from archive_store import ArchiveStore from wxwork_db import ( connect_sqlite, decrypt_with_keys, get_msg_type_name, load_keys, parse_content, ) def _safe_raw(value: Any) -> Any: """保留原表字段，但不把大二进制素材塞进数据库。""" if isinstance(value, bytes): if len(value) <= 64 * 1024: return { "encoding": "base64", "data": base64.b64encode(value).decode("ascii"), } return { "encoding": "external-binary", "size": len(value), "sha256": hashlib.sha256(value).hexdigest(), } if value is None or isinstance(value, (bool, float, int, str)): return value return str(value) class WxworkArchiveAdapter: def __init__( self, store: ArchiveStore, source_root: Path | str, *, keys_map: dict[str, str] | None = None, cache_dir: Path | str | None = None, corp_scopes: dict[str, str] | None = None, ): self.store = store self.source_root = Path(source_root).resolve() self.keys_map = keys_map if keys_map is not None else load_keys() self.cache_dir = Path(cache_dir or (self.source_root.parent / "wxwork_archive_cache")) self.corp_scopes = corp_scopes or {} @staticmethod def _columns(connection: sqlite3.Connection, table: str) -> list[str]: return [str(row[1]) for row in connection.execute(f'PRAGMA table_info("{table}")')] @staticmethod def _tables(connection: sqlite3.Connection) -> set[str]: return { str(row[0]) for row in connection.execute( "SELECT name FROM sqlite_master WHERE type='table'" ) } def _metadata( self, decrypted: list[tuple[str, str, str]] ) -> tuple[dict[tuple[str, str], str], dict[tuple[str, str], str]]: users: dict[tuple[str, str], str] = {} conversations: dict[tuple[str, str], str] = {} for path, name, account in decrypted: if name not in {"user.db", "session.db"}: continue connection = connect_sqlite(path) try: tables = self._tables(connection) if name == "user.db" and "user_table" in tables: columns = set(self._columns(connection, "user_table")) wanted = [ key for key in ("id", "name", "real_name", "account") if key in columns ] if "id" in wanted: for row in connection.execute( f"SELECT {','.join(wanted)} FROM user_table" ): item = dict(zip(wanted, row)) uid = str(item.get("id") or "") label = str( item.get("name") or item.get("real_name") or item.get("account") or uid ) if uid: users[(account, uid)] = label if name == "session.db" and "conversation_table" in tables: columns = set(self._columns(connection, "conversation_table")) wanted = [ key for key in ("id", "name", "roomname_remark", "session_id") if key in columns ] if "id" in wanted: for row in connection.execute( f"SELECT {','.join(wanted)} FROM conversation_table" ): item = dict(zip(wanted, row)) cid = str(item.get("id") or "") label = str( item.get("roomname_remark") or item.get("name") or item.get("session_id") or cid ) if cid: conversations[(account, cid)] = label finally: connection.close() return users, conversations @staticmethod def _value(row: dict[str, Any], *names: str) -> Any: for name in names: if row.get(name) not in (None, ""): return row[name] return None def _normalized_message( self, account: str, row: dict[str, Any], users: dict[tuple[str, str], str], conversations: dict[tuple[str, str], str], ) -> dict[str, Any]: conversation_id = str(row.get("conversation_id") or "unknown") sender_id = str(row.get("sender_id") or "") content_type = self._value(row, "content_type", "msg_type", "type") status = "normal" if self._value(row, "is_revoke", "revoke_status"): status = "revoked" elif self._value(row, "is_deleted", "delete_status"): status = "deleted" return { "source_table": "message_table", "source_message_id": str( self._value(row, "server_id", "client_id") or f"rowid:{row['__archive_rowid']}" ), "server_id": str(row.get("server_id") or ""), "client_id": str(row.get("client_id") or ""), "sequence_no": self._value( row, "message_seq", "sequence_no", "seq", "local_id" ), "conversation": { "external_id": conversation_id, "name": conversations.get((account, conversation_id), conversation_id), }, "sender": { "external_id": sender_id, "display_name": users.get((account, sender_id), sender_id), "identity_type": "wecom_local_uid", "scope_id": self.corp_scopes.get(account) or account, "source": "wxwork_db", }, "sent_at": row.get("send_time"), "message_type": get_msg_type_name(content_type), "direction": "outbound" if sender_id and sender_id == account else "inbound", "status": status, "content": parse_content(row.get("content")), "raw_fields": {key: _safe_raw(value) for key, value in row.items()}, } def import_all( self, user_id: int, ip: str = "local", *, batch_size: int = 1000 ) -> dict[str, Any]: batch_size = max(1, min(int(batch_size), 5000)) decrypted = decrypt_with_keys( str(self.source_root), str(self.cache_dir), self.keys_map, use_cache=True ) users, conversations = self._metadata(decrypted) result = {"accounts": 0, "batches": 0, "inserted": 0, "duplicates": 0} for path, name, account in decrypted: if name != "message.db": continue connection = connect_sqlite(path) try: if "message_table" not in self._tables(connection): continue columns = self._columns(connection, "message_table") if "send_time" not in columns: continue checkpoint = self.store.source_checkpoint(account, "message_table") sent_at = float(checkpoint.get("send_time") or 0) rowid = int(checkpoint.get("rowid") or 0) result["accounts"] += 1 while True: cursor = connection.execute( """SELECT rowid AS __archive_rowid,* FROM message_table WHERE send_time>? OR (send_time=? AND rowid>?) ORDER BY send_time,rowid LIMIT ?""", (sent_at, sent_at, rowid, batch_size), ) names = [item[0] for item in cursor.description] rows = [dict(zip(names, values)) for values in cursor.fetchall()] if not rows: break last = rows[-1] batch = self.store.import_messages( { "source_account": { "external_account_id": account, "display_name": users.get((account, account), account), "corp_scope_id": self.corp_scopes.get(account) or account, }, "source_table": "message_table", "checkpoint": { "send_time": float(last["send_time"]), "rowid": int(last["__archive_rowid"]), }, "messages": [ self._normalized_message( account, row, users, conversations ) for row in rows ], }, user_id, ip, ) result["batches"] += 1 result["inserted"] += int(batch["inserted"]) result["duplicates"] += int(batch["duplicates"]) sent_at = float(last["send_time"]) rowid = int(last["__archive_rowid"]) finally: connection.close() return result def main() -> None: parser = argparse.ArgumentParser(description="将现有企微解密库增量导入聊天归档") parser.add_argument("--source-root", required=True, help="企微账号数据根目录") parser.add_argument("--backend-db", default="backend.db") parser.add_argument("--cache-dir", default="") parser.add_argument("--admin-user", default="admin") parser.add_argument("--batch-size", type=int, default=1000) args = parser.parse_args() database = admin_backend.Database(Path(args.backend_db).resolve()) database.migrate() with database.connect() as connection: user = connection.execute( "SELECT id FROM users WHERE username=? AND active=1", (args.admin_user,) ).fetchone() if user is None: raise SystemExit("找不到可用的审计用户") store = ArchiveStore(database) store.initialize() adapter = WxworkArchiveAdapter( store, args.source_root, cache_dir=args.cache_dir or None ) print( json.dumps( adapter.import_all(int(user["id"]), batch_size=args.batch_size), ensure_ascii=False, ) ) if __name__ == "__main__": main()