def setUp(self):
self.frontend = InterComFrontEndBinding(config=self.config)
self.backend = None
class TestInterComTaskCommunication(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.tmp_dir = TemporaryDirectory(prefix='fact_test_')
cls.config = get_config_for_testing(temp_dir=cls.tmp_dir)
cls.config.set('ExpertSettings', 'communication_timeout', '1')
cls.mongo_server = MongoMgr(config=cls.config)
def setUp(self):
self.frontend = InterComFrontEndBinding(config=self.config)
self.backend = None
def tearDown(self):
for item in self.frontend.connections.keys():
self.frontend.client.drop_database(
self.frontend.connections[item]['name'])
if self.backend:
self.backend.shutdown()
self.frontend.shutdown()
gc.collect()
@classmethod
def tearDownClass(cls):
cls.mongo_server.shutdown()
cls.tmp_dir.cleanup()
def test_analysis_task(self):
self.backend = InterComBackEndAnalysisTask(config=self.config)
test_fw = create_test_firmware()
test_fw.file_path = None
self.frontend.add_analysis_task(test_fw)
task = self.backend.get_next_task()
self.assertEqual(task.get_uid(), test_fw.get_uid(), 'uid not correct')
self.assertIsNotNone(task.file_path, 'file_path not set')
self.assertTrue(os.path.exists(task.file_path), 'file does not exist')
def test_single_file_task(self):
self.backend = InterComBackEndSingleFileTask(config=self.config)
test_fw = create_test_firmware()
test_fw.file_path = None
test_fw.scheduled_analysis = ['binwalk']
self.frontend.add_single_file_task(test_fw)
task = self.backend.get_next_task()
assert task.get_uid() == test_fw.get_uid(
), 'uid not transported correctly'
assert task.scheduled_analysis
def test_re_analyze_task(self):
self.backend = InterComBackEndReAnalyzeTask(config=self.config)
fs_organizer = FS_Organizer(config=self.config)
test_fw = create_test_firmware()
fs_organizer.store_file(test_fw)
original_file_path = test_fw.file_path
original_binary = test_fw.binary
test_fw.file_path = None
test_fw.binary = None
self.frontend.add_re_analyze_task(test_fw)
task = self.backend.get_next_task()
self.assertEqual(task.get_uid(), test_fw.get_uid(), 'uid not correct')
self.assertIsNotNone(task.file_path, 'file path not set')
self.assertEqual(task.file_path, original_file_path)
self.assertIsNotNone(task.binary, 'binary not set')
self.assertEqual(task.binary, original_binary,
'binary content not correct')
def test_compare_task(self):
self.backend = InterComBackEndCompareTask(config=self.config)
self.frontend.add_compare_task('valid_id', force=False)
result = self.backend.get_next_task()
self.assertEqual(result, ('valid_id', False))
def test_analysis_plugin_publication(self):
self.backend = InterComBackEndAnalysisPlugInsPublisher(
config=self.config, analysis_service=AnalysisServiceMock())
plugins = self.frontend.get_available_analysis_plugins()
self.assertEqual(len(plugins), 1, 'Not all plug-ins found')
self.assertEqual(plugins, {'dummy': 'dummy description'},
'content not correct')
def test_analysis_plugin_publication_not_available(self):
with self.assertRaises(Exception):
self.frontend.get_available_analysis_plugins()
@mock.patch('intercom.front_end_binding.generate_task_id')
@mock.patch('intercom.back_end_binding.BinaryService')
def test_raw_download_task(self, binary_service_mock,
generate_task_id_mock):
binary_service_mock().get_binary_and_file_name.return_value = (
b'test', 'test.txt')
generate_task_id_mock.return_value = 'valid_uid_0.0'
result = self.frontend.get_binary_and_filename('valid_uid')
self.assertIsNone(result, 'should be none because of timeout')
self.backend = InterComBackEndRawDownloadTask(config=self.config)
task = self.backend.get_next_task()
self.assertEqual(task, 'valid_uid', 'task not correct')
result = self.frontend.get_binary_and_filename('valid_uid_0.0')
self.assertEqual(result, (b'test', 'test.txt'),
'retrieved binary not correct')
@mock.patch('intercom.front_end_binding.generate_task_id')
@mock.patch('intercom.back_end_binding.BinaryService')
def test_tar_repack_task(self, binary_service_mock, generate_task_id_mock):
binary_service_mock(
).get_repacked_binary_and_file_name.return_value = (b'test',
'test.tar')
generate_task_id_mock.return_value = 'valid_uid_0.0'
result = self.frontend.get_repacked_binary_and_file_name('valid_uid')
self.assertIsNone(result, 'should be none because of timeout')
self.backend = InterComBackEndTarRepackTask(config=self.config)
task = self.backend.get_next_task()
self.assertEqual(task, 'valid_uid', 'task not correct')
result = self.frontend.get_repacked_binary_and_file_name(
'valid_uid_0.0')
self.assertEqual(result, (b'test', 'test.tar'),
'retrieved binary not correct')
def __init__(self, config=None):
super().__init__(config=config)
self.intercom = InterComFrontEndBinding(config=config)
class AdminDbInterface(MongoInterfaceCommon):
READ_ONLY = False
def __init__(self, config=None):
super().__init__(config=config)
self.intercom = InterComFrontEndBinding(config=config)
def shutdown(self):
self.intercom.shutdown()
super().shutdown()
def remove_object_field(self, uid, field):
current_db = self.firmwares if self.is_firmware(
uid) else self.file_objects
current_db.find_one_and_update({'_id': uid}, {'$unset': {field: ''}})
def remove_from_object_array(self, uid, field, value):
current_db = self.firmwares if self.is_firmware(
uid) else self.file_objects
current_db.find_one_and_update({'_id': uid}, {'$pull': {field: value}})
def delete_firmware(self, uid, delete_root_file=True):
removed_fp, deleted = 0, 1
fw = self.firmwares.find_one(uid)
if fw:
for included_file_uid in fw['files_included']:
child_removed_fp, child_deleted = self._remove_virtual_path_entries(
uid, included_file_uid)
removed_fp += child_removed_fp
deleted += child_deleted
if delete_root_file:
self.intercom.delete_file(fw)
self._delete_swapped_analysis_entries(fw)
self.firmwares.delete_one({'_id': uid})
else:
logging.error('Firmware not found in Database: {}'.format(uid))
return removed_fp, deleted
def _delete_swapped_analysis_entries(self, fo_entry):
for key in fo_entry['processed_analysis']:
try:
if fo_entry['processed_analysis'][key]['file_system_flag']:
for analysis_key in fo_entry['processed_analysis'][
key].keys():
if analysis_key != 'file_system_flag' and isinstance(
fo_entry['processed_analysis'][key]
[analysis_key], str):
sanitize_id = fo_entry['processed_analysis'][key][
analysis_key]
entry = self.sanitize_fs.find_one(
{'filename': sanitize_id})
self.sanitize_fs.delete(entry._id)
except KeyError:
logging.warning(
'key error while deleting analysis for {}:{}'.format(
fo_entry['_id'], key))
def _remove_virtual_path_entries(self, root_uid, fo_uid):
'''
Recursively checks if the provided root uid is the only entry in the virtual path of the file object specified \
by fo uid. If this is the case, the file object is deleted from the database. Otherwise, only the entry from \
the virtual path is removed.
:param root_uid: the uid of the root firmware
:param fo_uid: he uid of the current file object
:return: tuple with numbers of recursively removed virtual file path entries and deleted files
'''
removed_fp, deleted = 0, 0
fo = self.file_objects.find_one(fo_uid)
if fo is not None:
for child_uid in fo['files_included']:
child_removed_fp, child_deleted = self._remove_virtual_path_entries(
root_uid, child_uid)
removed_fp += child_removed_fp
deleted += child_deleted
if any(
[root != root_uid for root in fo['virtual_file_path'].keys()]):
# there are more roots in the virtual path, meaning this file is included in other firmwares
self.remove_object_field(
fo_uid, 'virtual_file_path.{}'.format(root_uid))
if 'parent_firmware_uids' in fo:
self.remove_from_object_array(fo_uid,
'parent_firmware_uids',
root_uid)
removed_fp += 1
else:
self._delete_swapped_analysis_entries(fo)
self._delete_file_object(fo)
deleted += 1
return removed_fp, deleted
def _delete_file_object(self, fo_entry):
self.intercom.delete_file(fo_entry)
self.file_objects.delete_one({'_id': fo_entry['_id']})
def setUp(self):
self.config = get_config_for_testing(temp_dir=TMP_DIR)
self.config.set('ExpertSettings', 'communication_timeout', '1')
self.mongo_server = MongoMgr(config=self.config)
self.frontend = InterComFrontEndBinding(config=self.config)
self.backend = None