gengx
This commit is contained in:
Your Name
@@ -0,0 +1,284 @@
|
||||
# -*- 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()
|
||||
Reference in New Issue
Block a user