285 lines
11 KiB
Python
285 lines
11 KiB
Python
# -*- 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()
|