195 lines
9.6 KiB
Python
195 lines
9.6 KiB
Python
"""Synthetic contact snapshots; never touches live accounts or sends messages."""
|
|
from pathlib import Path
|
|
import sqlite3
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from contact_relation import read_contact_relation
|
|
from wxwork_db import WXWorkDB, decrypt_with_keys, wal_read_frames
|
|
|
|
|
|
class ContactRelationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.root = Path(self.temp.name)
|
|
self.user = self.root / '100' / 'user.db'
|
|
self.user.parent.mkdir()
|
|
self.conn = sqlite3.connect(self.user)
|
|
self.addCleanup(self.conn.close)
|
|
self.conn.executescript('''
|
|
CREATE TABLE external_user_relation_v3(user_id TEXT,status INTEGER,stranger_type INTEGER,add_customer_time REAL);
|
|
CREATE TABLE delete_external_userV1(user_id TEXT);
|
|
CREATE TABLE blacklist_external_userids(value TEXT);
|
|
INSERT INTO external_user_relation_v3 VALUES('200',2057,0,100);
|
|
''')
|
|
self.conn.commit()
|
|
|
|
def relation(self, account='100', conv='S:100_200'):
|
|
return read_contact_relation(self.user, account, conv)
|
|
|
|
def test_exact_active_relation_retains_original_add_time(self):
|
|
r = self.relation()
|
|
self.assertEqual((r['state'], r['account'], r['peer_id'], r['added_at']), ('active', '100', '200', 100))
|
|
|
|
def test_private_chat_identity_is_supported(self):
|
|
self.assertEqual(self.relation(conv='M:200')['state'], 'active')
|
|
|
|
def test_group_self_and_other_account_conversation_are_unknown(self):
|
|
for conv in ('R:100_200', 'S:101_200', 'S:100_100', 'M:100', 'S:100_200_300'):
|
|
with self.subTest(conv=conv):
|
|
self.assertEqual(self.relation(conv=conv)['state'], 'unknown')
|
|
self.assertEqual(self.relation(account='101', conv='S:101_200')['state'], 'unknown')
|
|
|
|
def test_deleted_or_blocklisted_is_never_active(self):
|
|
for table, column in (('delete_external_userV1', 'user_id'), ('blacklist_external_userids', 'value')):
|
|
with self.subTest(table=table):
|
|
self.conn.execute(f'INSERT INTO {table}({column}) VALUES(?)', ('200',))
|
|
self.conn.commit()
|
|
self.assertEqual(self.relation()['state'], 'blocked')
|
|
self.conn.execute(f'DELETE FROM {table}')
|
|
self.conn.commit()
|
|
|
|
def test_unknown_status_stranger_or_invalid_added_time_not_active(self):
|
|
for status, stranger, stamp in ((9999, 0, 100), (2057, 1, 100), (2057, 0, 0), (2057, 0, 'bad')):
|
|
self.conn.execute('UPDATE external_user_relation_v3 SET status=?,stranger_type=?,add_customer_time=?', (status, stranger, stamp))
|
|
self.conn.commit()
|
|
self.assertNotEqual(self.relation()['state'], 'active')
|
|
|
|
def test_missing_or_duplicate_relation_is_unknown(self):
|
|
self.assertEqual(self.relation(conv='S:100_201')['state'], 'unknown')
|
|
self.conn.execute("INSERT INTO external_user_relation_v3 VALUES('200',2057,0,100)")
|
|
self.conn.commit()
|
|
self.assertEqual(self.relation()['state'], 'unknown')
|
|
|
|
def test_missing_denial_table_is_unknown(self):
|
|
self.conn.execute('DROP TABLE delete_external_userV1')
|
|
self.conn.commit()
|
|
self.assertEqual(self.relation()['state'], 'unknown')
|
|
|
|
def test_wal_only_relation_updates_are_observed(self):
|
|
self.conn.execute('PRAGMA journal_mode=WAL')
|
|
self.conn.execute('PRAGMA wal_autocheckpoint=0')
|
|
self.conn.execute("INSERT INTO delete_external_userV1 VALUES('200')")
|
|
self.conn.commit()
|
|
self.assertEqual(self.relation()['state'], 'blocked')
|
|
self.conn.execute('DELETE FROM delete_external_userV1')
|
|
self.conn.commit()
|
|
self.assertEqual(self.relation()['state'], 'active')
|
|
|
|
def test_missing_path_does_not_create_database(self):
|
|
missing = self.root / '101' / 'user.db'
|
|
self.assertEqual(read_contact_relation(missing, '101', 'S:101_200')['state'], 'unknown')
|
|
self.assertFalse(missing.exists())
|
|
|
|
def reader(self):
|
|
db = WXWorkDB.__new__(WXWorkDB)
|
|
db.cache_dir = str(self.root)
|
|
db._decrypted = [(str(self.user), 'user.db', '100')]
|
|
db._source_signature = ('current',)
|
|
db._source_files_state = mock.Mock(return_value=('current',))
|
|
return db
|
|
|
|
def test_reader_only_uses_its_own_account_snapshot(self):
|
|
db = self.reader()
|
|
self.assertEqual(db._contact_relation('100', 'S:100_200')['state'], 'active')
|
|
db._decrypted = [(str(self.user), 'user.db', '101')]
|
|
self.assertEqual(db._contact_relation('100', 'S:100_200')['state'], 'unknown')
|
|
|
|
def test_reader_rejects_foreign_cache_and_missing_snapshot(self):
|
|
db = self.reader()
|
|
db.cache_dir = str(self.root / 'other')
|
|
self.assertEqual(db._contact_relation('100', 'S:100_200')['state'], 'unknown')
|
|
db = self.reader()
|
|
db._decrypted = []
|
|
self.assertEqual(db._contact_relation('100', 'S:100_200')['state'], 'unknown')
|
|
|
|
def test_reader_rejects_source_changed_before_or_during_read(self):
|
|
for signatures in ((('new',),), (('current',), ('new',))):
|
|
db = self.reader()
|
|
db._source_files_state = mock.Mock(side_effect=signatures)
|
|
self.assertEqual(db._contact_relation('100', 'S:100_200')['state'], 'unknown')
|
|
|
|
def test_encrypted_user_db_applies_wal_like_message_db(self):
|
|
source = self.root / 'source'
|
|
data = source / '100' / 'Data'
|
|
data.mkdir(parents=True)
|
|
for name in ('user.db', 'message.db'):
|
|
(data / name).write_bytes(b'encrypted')
|
|
(data / (name+'-wal')).write_bytes(b'wal')
|
|
with mock.patch('wxwork_db._is_plain_sqlite', return_value=False), \
|
|
mock.patch('wxwork_db.read_page_header', return_value={}), \
|
|
mock.patch('wxwork_db.verify_key', return_value=True), \
|
|
mock.patch('wxwork_db.merge_wal_into_db', return_value=True) as merge, \
|
|
mock.patch('wxwork_db.decrypt_db_to_file') as main:
|
|
decrypted = decrypt_with_keys(str(source), str(self.root / 'output'), {'100':'00'*16}, use_cache=False)
|
|
self.assertEqual({p[1] for p in decrypted}, {'user.db', 'message.db'})
|
|
self.assertEqual(merge.call_count, 2)
|
|
self.assertTrue(next(call for call in merge.call_args_list if call.args[0].endswith("user.db")).kwargs["committed_only"])
|
|
main.assert_not_called()
|
|
|
|
|
|
def test_context_refresh_observes_current_account_relation_wal(self):
|
|
source = self.root / 'source'
|
|
data = source / '100' / 'Data'
|
|
data.mkdir(parents=True)
|
|
user = sqlite3.connect(data / 'user.db')
|
|
self.addCleanup(user.close)
|
|
self.conn.backup(user)
|
|
user.executescript("""
|
|
CREATE TABLE user_table(id TEXT,name TEXT,real_name TEXT,account TEXT);
|
|
INSERT INTO user_table VALUES('200','客户甲','','');
|
|
""")
|
|
user.execute('PRAGMA journal_mode=WAL')
|
|
user.execute('PRAGMA wal_autocheckpoint=0')
|
|
user.commit()
|
|
messages = sqlite3.connect(data / 'message.db')
|
|
messages.executescript("""
|
|
CREATE TABLE message_table(sender_id TEXT,conversation_id TEXT,content_type INT,send_time INT,content TEXT,server_id TEXT);
|
|
INSERT INTO message_table VALUES('200','S:100_200',2,100,'新的问题','1001');
|
|
""")
|
|
messages.close()
|
|
reader = WXWorkDB(str(source), {}, str(self.root / 'current-cache'))
|
|
self.addCleanup(reader.close)
|
|
context = reader.get_conversation_context('客户甲', account='100', conv_id='S:100_200')
|
|
self.assertEqual(context['contact_relation']['state'], 'active')
|
|
user.execute("INSERT INTO delete_external_userV1 VALUES('200')")
|
|
user.commit()
|
|
context = reader.get_conversation_context('客户甲', account='100', conv_id='S:100_200')
|
|
self.assertEqual(context['contact_relation']['state'], 'blocked')
|
|
user.execute('DELETE FROM delete_external_userV1')
|
|
user.commit()
|
|
context = reader.get_conversation_context('客户甲', account='100', conv_id='S:100_200')
|
|
self.assertEqual(context['contact_relation']['state'], 'active')
|
|
|
|
|
|
def test_relationship_wal_ignores_uncommitted_tail(self):
|
|
wal = self.root / 'fixture.wal'
|
|
header = struct.pack('>8I', 0x377f0682, 3007000, 512, 0, 4, 5, 0, 0)
|
|
def frame(page, commit, content):
|
|
return struct.pack('>6I', page, commit, 4, 5, 0, 0) + content * 512
|
|
wal.write_bytes(header + frame(1, 1, b'A') + frame(1, 0, b'B'))
|
|
with mock.patch('wxwork_db.decrypt_page', side_effect=lambda key, n, data, size: data):
|
|
pages, meta = wal_read_frames(str(wal), b'k', committed_only=True)
|
|
legacy, _ = wal_read_frames(str(wal), b'k')
|
|
self.assertEqual(pages, {1: b'A'*512})
|
|
self.assertEqual(meta['db_size'], 1)
|
|
self.assertEqual(legacy, {1: b'B'*512})
|
|
|
|
def test_relationship_wal_commits_multi_page_transaction_together(self):
|
|
wal = self.root / 'fixture.wal'
|
|
header = struct.pack('>8I', 0x377f0682, 3007000, 512, 0, 4, 5, 0, 0)
|
|
def frame(page, commit, content):
|
|
return struct.pack('>6I', page, commit, 4, 5, 0, 0) + content * 512
|
|
wal.write_bytes(header + frame(1, 0, b'A') + frame(2, 2, b'B') + frame(2, 0, b'C'))
|
|
with mock.patch('wxwork_db.decrypt_page', side_effect=lambda key, n, data, size: data):
|
|
pages, meta = wal_read_frames(str(wal), b'k', committed_only=True)
|
|
self.assertEqual(pages, {1: b'A'*512, 2: b'B'*512})
|
|
self.assertEqual(meta['db_size'], 2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|