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