439 lines
16 KiB
Python
439 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""桌面启动自动归档桥接器的增量与素材关联测试。"""
|
|
|
|
import json
|
|
import sqlite3
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest import TestCase, main, mock
|
|
|
|
from archive_content_parser import (
|
|
decode_hex_protobuf_text,
|
|
parse_file_message_metadata,
|
|
parse_mini_program_metadata,
|
|
)
|
|
from archive_auto_backup import (
|
|
ArchiveApiClient,
|
|
AutoBackupConfig,
|
|
IncrementalArchiveImporter,
|
|
_load_exporter_keys,
|
|
)
|
|
|
|
|
|
class _FakeExporter:
|
|
@staticmethod
|
|
def connect_sqlite(path):
|
|
connection = sqlite3.connect(path)
|
|
connection.text_factory = lambda value: value.decode("utf-8", errors="replace")
|
|
return connection
|
|
|
|
@staticmethod
|
|
def parse_content(value):
|
|
if isinstance(value, bytes):
|
|
return value.decode("utf-8", errors="replace")
|
|
return str(value or "")
|
|
|
|
@staticmethod
|
|
def get_msg_type_name(value):
|
|
return {2: "文本", 3: "图片", 20: "合并转发", 78: "修改群名"}.get(
|
|
int(value), "未知"
|
|
)
|
|
|
|
|
|
class _FakeMediaExporter:
|
|
def __init__(self, media_path: Path):
|
|
self.media_path = media_path
|
|
|
|
@staticmethod
|
|
def extract_media_refs(content):
|
|
if isinstance(content, bytes) and b"media-ref" in content:
|
|
return {"uuids": ["media-ref"], "urls": [], "filenames": []}
|
|
return {"uuids": [], "urls": [], "filenames": []}
|
|
|
|
@staticmethod
|
|
def build_cache_index(_account_dir):
|
|
return {"by_uuid": {}, "by_name": {}}
|
|
|
|
def match_media(self, refs, _index):
|
|
if refs["uuids"]:
|
|
return str(self.media_path), "Image", "uuid"
|
|
return None, None, None
|
|
|
|
|
|
class _FakeApi:
|
|
def __init__(self):
|
|
self.cursor = {}
|
|
self.payloads = []
|
|
self.uploads = []
|
|
self.metadata_payloads = []
|
|
|
|
def checkpoint(self, _account):
|
|
return dict(self.cursor)
|
|
|
|
def advance_checkpoint(
|
|
self, _account, checkpoint, *, display_name="", corp_scope_id=""
|
|
):
|
|
self.cursor = dict(checkpoint)
|
|
return dict(checkpoint)
|
|
|
|
def pending_attachments(self, _account, limit=500):
|
|
return []
|
|
|
|
def upload_media(self, path):
|
|
self.uploads.append(Path(path))
|
|
return "media-001"
|
|
|
|
def sync_metadata(self, payload):
|
|
self.metadata_payloads.append(payload)
|
|
return {
|
|
"people_synced": len(payload.get("people") or []),
|
|
"conversations_updated": len(payload.get("conversations") or []),
|
|
}
|
|
|
|
def import_messages(self, payload):
|
|
self.payloads.append(payload)
|
|
self.cursor = dict(payload["checkpoint"])
|
|
return {
|
|
"received": len(payload["messages"]),
|
|
"inserted": len(payload["messages"]),
|
|
"duplicates": 0,
|
|
}
|
|
|
|
|
|
class _FailingMediaApi(_FakeApi):
|
|
def upload_media(self, path):
|
|
raise RuntimeError(f"COS unavailable: {path}")
|
|
|
|
|
|
class ArchiveAutoBackupTest(TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp.name)
|
|
self.account = "16880001"
|
|
self.source = self.root / "WXWork"
|
|
(self.source / self.account).mkdir(parents=True)
|
|
self.decrypted = self.root / "decrypted" / self.account
|
|
self.decrypted.mkdir(parents=True)
|
|
self.message_db = self.decrypted / "message.db"
|
|
database = sqlite3.connect(self.message_db)
|
|
try:
|
|
database.execute(
|
|
"""CREATE TABLE message_table(
|
|
send_time INTEGER,conversation_id TEXT,sender_id TEXT,
|
|
content BLOB,content_type INTEGER,server_id TEXT,
|
|
client_id TEXT,sequence INTEGER)"""
|
|
)
|
|
database.executemany(
|
|
"INSERT INTO message_table VALUES (?,?,?,?,?,?,?,?)",
|
|
[
|
|
(1000, "S:16880001_20001", "20001", b"hello", 2,
|
|
"server-1", "client-1", 1),
|
|
(1001, "R:room-1", self.account, b"media-ref", 3,
|
|
"server-2", "client-2", 2),
|
|
],
|
|
)
|
|
database.commit()
|
|
finally:
|
|
database.close()
|
|
self.media = self.root / "source-image.png"
|
|
self.media.write_bytes(b"png-data")
|
|
self.config = AutoBackupConfig(
|
|
exporter_root=self.root,
|
|
source_root=self.source,
|
|
work_root=self.root / "work",
|
|
api_urls=("http://127.0.0.1:8766",),
|
|
corp_scope_id="corp-one",
|
|
batch_size=1,
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.temp.cleanup()
|
|
|
|
def test_incremental_batches_upload_media_before_advancing_checkpoint(self):
|
|
api = _FakeApi()
|
|
importer = IncrementalArchiveImporter(
|
|
self.config,
|
|
api,
|
|
_FakeExporter(),
|
|
_FakeMediaExporter(self.media),
|
|
)
|
|
decrypted = [(str(self.message_db), "message.db", self.account)]
|
|
first = importer.run(decrypted)
|
|
second = importer.run(decrypted)
|
|
|
|
self.assertEqual(first["inserted"], 2)
|
|
self.assertEqual(first["batches"], 2)
|
|
self.assertEqual(first["media"], 1)
|
|
self.assertEqual(second["inserted"], 0)
|
|
self.assertEqual(len(api.uploads), 1)
|
|
self.assertEqual(api.payloads[-1]["messages"][0]["media_ids"], ["media-001"])
|
|
self.assertEqual(api.payloads[-1]["messages"][0]["direction"], "outbound")
|
|
self.assertEqual(
|
|
api.payloads[0]["messages"][0]["sender"]["scope_id"], "corp-one"
|
|
)
|
|
self.assertEqual(api.cursor, {"send_time": 1001.0, "rowid": 2})
|
|
|
|
def test_metadata_resolves_direct_chat_peer_nickname(self):
|
|
user_db = self.decrypted / "user.db"
|
|
with sqlite3.connect(user_db) as database:
|
|
database.execute(
|
|
"CREATE TABLE user_table(id TEXT,name TEXT,real_name TEXT,account TEXT)"
|
|
)
|
|
database.executemany(
|
|
"INSERT INTO user_table VALUES (?,?,?,?)",
|
|
[
|
|
(self.account, "归档账号", "", ""),
|
|
("20001", "客户昵称", "", ""),
|
|
],
|
|
)
|
|
database.close()
|
|
session_db = self.decrypted / "session.db"
|
|
with sqlite3.connect(session_db) as database:
|
|
database.execute(
|
|
"""CREATE TABLE conversation_table(
|
|
id TEXT,name TEXT,roomname_remark TEXT,session_id TEXT)"""
|
|
)
|
|
database.execute(
|
|
"INSERT INTO conversation_table VALUES (?,?,?,?)",
|
|
("S:16880001_20001", "", "", ""),
|
|
)
|
|
database.close()
|
|
api = _FakeApi()
|
|
importer = IncrementalArchiveImporter(
|
|
self.config,
|
|
api,
|
|
_FakeExporter(),
|
|
_FakeMediaExporter(self.media),
|
|
)
|
|
importer.run(
|
|
[
|
|
(str(user_db), "user.db", self.account),
|
|
(str(session_db), "session.db", self.account),
|
|
(str(self.message_db), "message.db", self.account),
|
|
]
|
|
)
|
|
self.assertEqual(
|
|
api.payloads[0]["messages"][0]["conversation"]["name"], "客户昵称"
|
|
)
|
|
self.assertEqual(
|
|
api.metadata_payloads[0]["conversations"][0]["name"], "客户昵称"
|
|
)
|
|
|
|
def test_exporter_key_file_is_always_read_as_utf8(self):
|
|
key = "ab" * 16
|
|
(self.root / "wxwork_keys.json").write_text(
|
|
json.dumps({"备注": "中文", "keys": {self.account: key}}, ensure_ascii=False),
|
|
encoding="utf-8",
|
|
)
|
|
self.assertEqual(_load_exporter_keys(self.root), {self.account: key})
|
|
|
|
def test_nested_short_text_protobuf_keeps_single_character_and_emoji(self):
|
|
cases = {
|
|
"0a07080012030a0131": "1",
|
|
"0a09080012050a03e5a5bd": "好",
|
|
"0a0a080012060a04f09f918c": "👌",
|
|
"0a08080012040a023131": "11",
|
|
}
|
|
for encoded, expected in cases.items():
|
|
with self.subTest(encoded=encoded):
|
|
self.assertEqual(
|
|
decode_hex_protobuf_text(encoded, "文本"), expected
|
|
)
|
|
self.assertEqual(decode_hex_protobuf_text("1688857886854158", "安全通知"), "")
|
|
|
|
def test_file_card_is_structured_and_marked_for_cache_retry(self):
|
|
def varint(value):
|
|
output = bytearray()
|
|
while value >= 0x80:
|
|
output.append((value & 0x7F) | 0x80)
|
|
value >>= 7
|
|
output.append(value)
|
|
return bytes(output)
|
|
|
|
def bytes_field(number, value):
|
|
return varint((number << 3) | 2) + varint(len(value)) + value
|
|
|
|
file_size = 162_862_195
|
|
filename = "DoctorWorkstation-Setup-Windows-x64-1.0.0.exe"
|
|
payload = b"".join(
|
|
[
|
|
bytes_field(1, b"opaque-source-reference"),
|
|
bytes_field(2, filename.encode()),
|
|
varint(4 << 3) + varint(file_size),
|
|
bytes_field(10, b"9D4295ED638B290925A5A960D644934E"),
|
|
]
|
|
)
|
|
metadata = parse_file_message_metadata(payload, 20)
|
|
self.assertEqual(metadata["original_filename"], filename)
|
|
self.assertEqual(metadata["size_bytes"], file_size)
|
|
|
|
database = sqlite3.connect(self.message_db)
|
|
try:
|
|
database.execute(
|
|
"INSERT INTO message_table VALUES (?,?,?,?,?,?,?,?)",
|
|
(1002, "S:16880001_20001", self.account, payload, 20,
|
|
"server-file", "client-file", 3),
|
|
)
|
|
database.commit()
|
|
finally:
|
|
database.close()
|
|
api = _FakeApi()
|
|
importer = IncrementalArchiveImporter(
|
|
self.config, api, _FakeExporter(), _FakeMediaExporter(self.media)
|
|
)
|
|
result = importer.run([(str(self.message_db), "message.db", self.account)])
|
|
message = api.payloads[-1]["messages"][0]
|
|
self.assertEqual(result["inserted"], 3)
|
|
self.assertEqual(message["message_type"], "文件")
|
|
self.assertIn(filename, message["content"])
|
|
self.assertIn("源文件未缓存", message["content"])
|
|
self.assertEqual(
|
|
message["attachment_metadata"][0]["status"], "source_not_cached"
|
|
)
|
|
|
|
def test_application_conversation_is_skipped_and_checkpoint_advances(self):
|
|
database = sqlite3.connect(self.message_db)
|
|
try:
|
|
database.execute(
|
|
"INSERT INTO message_table VALUES (?,?,?,?,?,?,?,?)",
|
|
(1002, "Y:10011", self.account, b"approval notice", 2,
|
|
"server-app", "client-app", 3),
|
|
)
|
|
database.commit()
|
|
finally:
|
|
database.close()
|
|
api = _FakeApi()
|
|
importer = IncrementalArchiveImporter(
|
|
self.config, api, _FakeExporter(), _FakeMediaExporter(self.media)
|
|
)
|
|
result = importer.run([(str(self.message_db), "message.db", self.account)])
|
|
imported_conversations = {
|
|
message["conversation"]["external_id"]
|
|
for payload in api.payloads
|
|
for message in payload["messages"]
|
|
}
|
|
self.assertNotIn("Y:10011", imported_conversations)
|
|
self.assertEqual(result["application_messages_skipped"], 1)
|
|
self.assertEqual(result["inserted"], 2)
|
|
self.assertEqual(api.cursor, {"send_time": 1002.0, "rowid": 3})
|
|
|
|
def test_mini_program_card_is_decoded_from_protobuf(self):
|
|
def varint(value):
|
|
output = bytearray()
|
|
while value >= 0x80:
|
|
output.append((value & 0x7F) | 0x80)
|
|
value >>= 7
|
|
output.append(value)
|
|
return bytes(output)
|
|
|
|
def bytes_field(number, value):
|
|
return varint((number << 3) | 2) + varint(len(value)) + value
|
|
|
|
page_path = b"pages/order/monad/monad.html?id=11&doctor_id=117"
|
|
nested = b"".join(
|
|
[
|
|
bytes_field(1, b"gh_example@app"),
|
|
bytes_field(2, b"wx79b9a0bfbfe7cbcd"),
|
|
bytes_field(3, page_path),
|
|
bytes_field(6, b"https://example.test/cover.png"),
|
|
bytes_field(7, "点击进入诊室".encode()),
|
|
bytes_field(10, "甄养堂互联网医院".encode()),
|
|
]
|
|
)
|
|
raw = bytes_field(3, "点击进入诊室".encode()) + bytes_field(107, nested)
|
|
metadata = parse_mini_program_metadata(raw, 78)
|
|
self.assertEqual(metadata["title"], "点击进入诊室")
|
|
self.assertEqual(metadata["app_name"], "甄养堂互联网医院")
|
|
self.assertEqual(metadata["page_path"], page_path.decode())
|
|
|
|
database = sqlite3.connect(self.message_db)
|
|
try:
|
|
database.execute(
|
|
"INSERT INTO message_table VALUES (?,?,?,?,?,?,?,?)",
|
|
(1002, "S:16880001_20001", self.account, raw, 78,
|
|
"server-mini", "client-mini", 3),
|
|
)
|
|
database.commit()
|
|
finally:
|
|
database.close()
|
|
api = _FakeApi()
|
|
importer = IncrementalArchiveImporter(
|
|
self.config, api, _FakeExporter(), _FakeMediaExporter(self.media)
|
|
)
|
|
importer.run([(str(self.message_db), "message.db", self.account)])
|
|
message = api.payloads[-1]["messages"][0]
|
|
self.assertEqual(message["message_type"], "小程序")
|
|
self.assertIn("甄养堂互联网医院", message["content"])
|
|
self.assertIn(page_path.decode(), message["content"])
|
|
|
|
def test_media_upload_failure_does_not_advance_that_batch_checkpoint(self):
|
|
api = _FailingMediaApi()
|
|
importer = IncrementalArchiveImporter(
|
|
AutoBackupConfig(
|
|
exporter_root=self.config.exporter_root,
|
|
source_root=self.config.source_root,
|
|
work_root=self.config.work_root,
|
|
api_urls=self.config.api_urls,
|
|
corp_scope_id=self.config.corp_scope_id,
|
|
batch_size=2,
|
|
),
|
|
api,
|
|
_FakeExporter(),
|
|
_FakeMediaExporter(self.media),
|
|
)
|
|
decrypted = [(str(self.message_db), "message.db", self.account)]
|
|
with self.assertRaisesRegex(RuntimeError, "COS unavailable"):
|
|
importer.run(decrypted)
|
|
self.assertEqual(api.payloads, [])
|
|
self.assertEqual(api.cursor, {})
|
|
|
|
def test_client_uploads_large_media_parts_and_reports_etags(self):
|
|
large = self.root / "video.mp4"
|
|
large.write_bytes(b"0123456789abcdefghij")
|
|
client = ArchiveApiClient("http://127.0.0.1:8766", "test-key")
|
|
prepared = {
|
|
"media": {"id": "media-large"},
|
|
"reused": False,
|
|
"upload_mode": "multipart",
|
|
"multipart": {
|
|
"upload_id": "upload-1",
|
|
"part_size": 10,
|
|
"parts": [
|
|
{"part_number": 1, "size_bytes": 10, "upload_url": "https://cos/1"},
|
|
{"part_number": 2, "size_bytes": 10, "upload_url": "https://cos/2"},
|
|
],
|
|
"completed_parts": [
|
|
{"part_number": 1, "size_bytes": 10, "etag": "etag-1"}
|
|
],
|
|
},
|
|
}
|
|
client._json = mock.Mock(side_effect=[prepared, {"media": {"status": "ready"}}])
|
|
|
|
def uploaded(url, data, timeout):
|
|
response = mock.Mock()
|
|
response.headers = {"ETag": f'"etag-{url[-1]}"'}
|
|
response.raise_for_status.return_value = None
|
|
return response
|
|
|
|
try:
|
|
with mock.patch(
|
|
"archive_auto_backup.requests.put", side_effect=uploaded
|
|
) as put:
|
|
self.assertEqual(client.upload_media(large), "media-large")
|
|
self.assertEqual(put.call_count, 1)
|
|
finally:
|
|
client.close()
|
|
completed = client._json.call_args_list[-1]
|
|
self.assertTrue(completed.args[1].endswith("/multipart-complete"))
|
|
self.assertEqual(
|
|
completed.kwargs["payload"]["parts"],
|
|
[
|
|
{"part_number": 1, "etag": "etag-1"},
|
|
{"part_number": 2, "etag": "etag-2"},
|
|
],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|