212 lines
11 KiB
Python
212 lines
11 KiB
Python
"""Bounded probe work; synthetic files/process identities, no native calls or messages."""
|
|
import hashlib
|
|
from contextlib import closing
|
|
import os
|
|
from pathlib import Path
|
|
import sqlite3
|
|
import tempfile
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import wecom_native_sender as sender
|
|
import wxwork_protocol_probe as probe
|
|
import wxwork_db as db
|
|
import wecom_environment as environment
|
|
import reply_database
|
|
from wecom_native_builds import BUILDS
|
|
|
|
|
|
class BuildHashReuseTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp=tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.path=Path(self.temp.name)/'synthetic-WXWork.exe'
|
|
self.path.write_bytes(b'synthetic pinned executable'*65536)
|
|
self.creation=self.enterContext(mock.patch.object(sender,'_process_creation_id',return_value=100))
|
|
self.hash_file=self.enterContext(mock.patch.object(sender,'_hash_client_file',wraps=sender._hash_client_file))
|
|
|
|
def digest(self,pid=7,readonly=True):
|
|
return sender._client_file_digest(self.path,pid,None,readonly=readonly)
|
|
|
|
def test_one_probe_reads_same_process_build_once_but_next_probe_rechecks(self):
|
|
expected=hashlib.sha256(self.path.read_bytes()).hexdigest()
|
|
with sender.readonly_build_validation_scope() as metrics:
|
|
self.assertEqual(self.digest(),expected)
|
|
self.assertEqual(self.digest(),expected)
|
|
self.assertEqual(self.hash_file.call_count,1)
|
|
self.assertEqual(metrics['reusedFiles'],1)
|
|
self.assertEqual(metrics['bytesHashed'],self.path.stat().st_size)
|
|
with sender.readonly_build_validation_scope():self.digest()
|
|
self.assertEqual(self.hash_file.call_count,2)
|
|
|
|
def test_new_pid_and_reused_pid_with_new_creation_time_rehash(self):
|
|
with sender.readonly_build_validation_scope():
|
|
self.digest(7)
|
|
self.digest(8)
|
|
self.creation.return_value=101
|
|
self.digest(7)
|
|
self.assertEqual(self.hash_file.call_count,3)
|
|
|
|
def test_changed_binary_is_not_reused_with_same_process(self):
|
|
with sender.readonly_build_validation_scope():
|
|
old=self.digest()
|
|
self.path.write_bytes(b'different current executable')
|
|
current=self.digest()
|
|
self.assertNotEqual(old,current)
|
|
self.assertEqual(self.hash_file.call_count,2)
|
|
|
|
def test_unavailable_process_start_time_does_not_enable_cache(self):
|
|
self.creation.return_value=None
|
|
with sender.readonly_build_validation_scope():
|
|
self.digest();self.digest()
|
|
self.assertEqual(self.hash_file.call_count,2)
|
|
|
|
def test_actual_send_mode_never_uses_probe_hash_cache(self):
|
|
with sender.readonly_build_validation_scope():
|
|
self.digest();self.digest(readonly=False);self.digest(readonly=False)
|
|
self.assertEqual(self.hash_file.call_count,3)
|
|
|
|
def test_file_change_during_hash_fails_closed(self):
|
|
before=self.path.stat()
|
|
changed=SimpleNamespace(st_dev=before.st_dev,st_ino=before.st_ino,st_size=before.st_size,
|
|
st_mtime_ns=before.st_mtime_ns+1,st_ctime_ns=before.st_ctime_ns)
|
|
with mock.patch.object(sender.os,'fstat',side_effect=[before,changed]):
|
|
with self.assertRaisesRegex(sender.NativeUnavailable,'发生变化'):
|
|
sender._hash_client_file(self.path)
|
|
|
|
def test_native_client_still_validates_live_identity_each_time(self):
|
|
path=self.path
|
|
class Kernel:
|
|
def CreateToolhelp32Snapshot(self,*_):return 11
|
|
def Module32FirstW(self,_snapshot,pointer):
|
|
value=pointer._obj
|
|
value.szModule='wxwork.exe';value.modBaseAddr=0x400000;value.szExePath=str(path)
|
|
return True
|
|
def CloseHandle(self,*_):return True
|
|
def OpenProcess(self,*_):return 12
|
|
accounts=iter(['100','101'])
|
|
def validate(client):
|
|
client.account=next(accounts)
|
|
return client.account
|
|
with mock.patch.object(sender,'kernel',return_value=Kernel()), \
|
|
mock.patch.object(sender,'build_for_hash',return_value=BUILDS[0]), \
|
|
mock.patch.object(sender.NativeClient,'validate',autospec=True,side_effect=validate) as checks, \
|
|
sender.readonly_build_validation_scope():
|
|
with sender.NativeClient(7) as first:self.assertEqual(first.account,'100')
|
|
with sender.NativeClient(7) as second:self.assertEqual(second.account,'101')
|
|
self.assertEqual(checks.call_count,2)
|
|
self.assertEqual(self.hash_file.call_count,1)
|
|
|
|
|
|
class ProbeFreshnessTests(unittest.TestCase):
|
|
def setUp(self):
|
|
build=BUILDS[0]
|
|
self.client=dict(pid=123,path=str(Path('D:/synthetic/WXWork.exe')),version=build.version,architecture='x86',
|
|
size=123,machine='0x14c',sections=5,buildTimestamp=1,isDll=False)
|
|
self.live=dict(pid=123,accountId='100',clientVersion=build.version,clientSha256=build.sha256)
|
|
self.enterContext(mock.patch.object(probe,'running_clients',return_value=[self.client]))
|
|
self.inspect=self.enterContext(mock.patch.object(probe,'inspect_binary',side_effect=AssertionError('duplicate static scan')))
|
|
self.enterContext(mock.patch.object(probe,'bundled_agent_info',return_value={}))
|
|
self.discover=self.enterContext(mock.patch.object(sender,'discover',return_value=self.live))
|
|
self.database=self.enterContext(mock.patch.object(environment,'protocol_database_status',return_value={'ready':True}))
|
|
|
|
def test_repeated_probes_revalidate_account_and_database(self):
|
|
for _ in range(2):
|
|
report=probe.probe_installed_client()
|
|
self.assertTrue(report['canSend'])
|
|
self.assertEqual(report['binaries'][0]['size'],123)
|
|
self.assertTrue({'processes','clientValidation','databaseValidation','identityRecheck','total'} <= report['timingsMs'].keys())
|
|
self.assertEqual(self.discover.call_count,4)
|
|
self.assertEqual(self.database.call_count,2)
|
|
self.inspect.assert_not_called()
|
|
|
|
def test_change_after_database_check_is_never_hidden_by_static_reuse(self):
|
|
for field,value in [('accountId','101'),('pid',124),('clientSha256','changed-build')]:
|
|
with self.subTest(field=field):
|
|
self.discover.side_effect=[self.live,{**self.live,field:value}]
|
|
report=probe.probe_installed_client(initialize_database=True)
|
|
self.assertFalse(report['canSend'])
|
|
self.assertFalse(report['loginVerified'])
|
|
self.assertEqual(report['status'],'client_not_ready')
|
|
|
|
|
|
class ScopedDatabaseTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp=tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.root=Path(self.temp.name)
|
|
self.source=self.root/'WXWork'
|
|
self.appdata=self.root/'appdata'
|
|
self.appdata.mkdir()
|
|
for account in ('100','101'):
|
|
data=self.source/account/'Data';data.mkdir(parents=True)
|
|
with closing(sqlite3.connect(data/'message.db')) as connection:
|
|
connection.execute('CREATE TABLE message_table(sender_id TEXT,conversation_id TEXT,content_type INT,send_time INT,content TEXT)')
|
|
with closing(sqlite3.connect(data/'user.db')) as connection:
|
|
connection.execute('CREATE TABLE user_table(id TEXT,name TEXT,real_name TEXT,account TEXT)')
|
|
connection.execute('INSERT INTO user_table VALUES(?,?,?,?)',('200','合成客户'+account,'',account))
|
|
connection.commit()
|
|
with closing(sqlite3.connect(data/'session.db')) as connection:
|
|
connection.execute('CREATE TABLE conversation_table(id TEXT,name TEXT,roomname_remark TEXT,session_id TEXT)')
|
|
connection.execute('INSERT INTO conversation_table VALUES(?,?,?,?)',('S:'+account+'_200','合成会话','',''))
|
|
connection.commit()
|
|
self.enterContext(mock.patch.object(reply_database,'application_data_dir',return_value=self.appdata))
|
|
self.enterContext(mock.patch.object(reply_database,'load_keys',return_value={}))
|
|
self.acquire=self.enterContext(mock.patch('wxwork_local_setup.acquire_local_keys',side_effect=AssertionError('no real key worker')))
|
|
|
|
def test_probe_prepares_only_current_message_database_without_metadata(self):
|
|
with mock.patch.object(db,'detect_wxwork_dir',return_value=str(self.source)) as detect, \
|
|
mock.patch.object(reply_database,'detect_wxwork_dir',side_effect=AssertionError('directory detected twice')), \
|
|
mock.patch.object(db.WXWorkDB,'_load_metadata',side_effect=AssertionError('metadata not needed')):
|
|
report=environment.protocol_database_status('100',initialize=False)
|
|
self.assertTrue(report['ready'])
|
|
detect.assert_called_once_with(account='100')
|
|
cached=sorted(str(path.relative_to(self.appdata/'wxwork_reply_cache')).replace('\\','/') for path in (self.appdata/'wxwork_reply_cache').rglob('*.db'))
|
|
self.assertEqual(cached,['100/message.db'])
|
|
self.acquire.assert_not_called()
|
|
|
|
def test_actual_reply_startup_loads_current_account_metadata(self):
|
|
database=reply_database.LiveReplyDatabase._open_database(account='100',initialize=False,source_directory=str(self.source))
|
|
try:
|
|
self.assertTrue(database.health_check(account='100'))
|
|
self.assertFalse(database.health_check(account='101'))
|
|
self.assertEqual(database.user_cache[('100','200')],'合成客户100')
|
|
self.assertEqual(set(database._conns),{'100'})
|
|
finally:database.close()
|
|
self.assertFalse((self.appdata/'wxwork_reply_cache'/'101').exists())
|
|
|
|
def test_default_database_keeps_all_account_and_metadata_behavior(self):
|
|
database=db.WXWorkDB(str(self.source),{},str(self.root/'all-cache'))
|
|
try:
|
|
self.assertEqual(set(database._conns),{'100','101'})
|
|
self.assertEqual(database.user_cache[('101','200')],'合成客户101')
|
|
finally:database.close()
|
|
|
|
def test_other_account_cache_cannot_prove_current_account_ready(self):
|
|
cached=db.WXWorkDB(str(self.source),{},str(self.root/'cache'),account='101',load_metadata=False)
|
|
cached.close()
|
|
database=db.WXWorkDB(str(self.source),{},str(self.root/'cache'),account='999',load_metadata=False)
|
|
try:
|
|
self.assertFalse(database.health_check(account='999'))
|
|
self.assertFalse(database.health_check(account='101'))
|
|
self.assertFalse(database._conns)
|
|
finally:database.close()
|
|
|
|
def test_schema_check_is_preserved_in_lightweight_probe(self):
|
|
with closing(sqlite3.connect(self.source/'100'/'Data'/'message.db')) as connection:
|
|
connection.execute('DROP TABLE message_table')
|
|
with mock.patch.object(db,'detect_wxwork_dir',return_value=str(self.source)):
|
|
report=environment.protocol_database_status('100',initialize=False)
|
|
self.assertFalse(report['ready'])
|
|
self.acquire.assert_not_called()
|
|
|
|
def test_unsafe_account_or_database_name_is_rejected(self):
|
|
for options in ({'account':'../100'},{'database_names':['../../secret.db']}):
|
|
with self.subTest(options=options),self.assertRaises(ValueError):
|
|
db.decrypt_with_keys(str(self.source),str(self.root/'cache'),{},**options)
|
|
|
|
|
|
if __name__=='__main__':
|
|
unittest.main()
|