class TestTagPropagation(unittest.TestCase): def setUp(self): self._tmp_dir = TemporaryDirectory() self._config = initialize_config(self._tmp_dir) self.analysis_finished_event = Event() self.uid_of_key_file = '530bf2f1203b789bfe054d3118ebd29a04013c587efd22235b3b9677cee21c0e_2048' self._mongo_server = MongoMgr(config=self._config, auth=False) self.backend_interface = BackEndDbInterface(config=self._config) self._analysis_scheduler = AnalysisScheduler( config=self._config, post_analysis=self.count_analysis_finished_event) self._tagging_scheduler = TaggingDaemon( analysis_scheduler=self._analysis_scheduler) self._unpack_scheduler = UnpackingScheduler( config=self._config, post_unpack=self._analysis_scheduler.add_task) def count_analysis_finished_event(self, fw_object): self.backend_interface.add_object(fw_object) if fw_object.uid == self.uid_of_key_file: self.analysis_finished_event.set() def _wait_for_empty_tag_queue(self): while not self._analysis_scheduler.tag_queue.empty(): sleep(0.1) def tearDown(self): self._unpack_scheduler.shutdown() self._tagging_scheduler.shutdown() self._analysis_scheduler.shutdown() clean_test_database(self._config, get_database_names(self._config)) self._mongo_server.shutdown() self._tmp_dir.cleanup() gc.collect() def test_run_analysis_with_tag(self): test_fw = Firmware( file_path='{}/container/with_key.7z'.format(get_test_data_dir())) test_fw.release_date = '2017-01-01' test_fw.scheduled_analysis = ['crypto_material'] self._unpack_scheduler.add_task(test_fw) assert self.analysis_finished_event.wait(timeout=20) processed_fo = self.backend_interface.get_object( self.uid_of_key_file, analysis_filter=['crypto_material']) assert processed_fo.processed_analysis['crypto_material'][ 'tags'], 'no tags set in analysis' self._wait_for_empty_tag_queue() processed_fw = self.backend_interface.get_object( test_fw.uid, analysis_filter=['crypto_material']) assert processed_fw.analysis_tags, 'tags not propagated properly' assert processed_fw.analysis_tags['crypto_material'][ 'private_key_inside']
class TestStorageDbInterfaceFrontendEditing(unittest.TestCase): def setUp(self): self._config = get_config_for_testing(TMP_DIR) self.mongo_server = MongoMgr(config=self._config) self.db_frontend_editing = FrontendEditingDbInterface( config=self._config) self.db_frontend_interface = FrontEndDbInterface(config=self._config) self.db_backend_interface = BackEndDbInterface(config=self._config) def tearDown(self): self.db_frontend_editing.shutdown() self.db_frontend_interface.shutdown() self.db_backend_interface.client.drop_database( self._config.get('data_storage', 'main_database')) self.db_backend_interface.shutdown() self.mongo_server.shutdown() TMP_DIR.cleanup() gc.collect() def test_add_comment(self): test_fw = create_test_firmware() self.db_backend_interface.add_object(test_fw) comment, author, uid, time = 'this is a test comment!', 'author', test_fw.get_uid( ), 1234567890 self.db_frontend_editing.add_comment_to_object(uid, comment, author, time) test_fw = self.db_backend_interface.get_object(uid) self.assertEqual(test_fw.comments[0], { 'time': str(time), 'author': author, 'comment': comment }) def test_get_latest_comments(self): comments = [{ 'time': '1234567890', 'author': 'author1', 'comment': 'test comment' }, { 'time': '1234567899', 'author': 'author2', 'comment': 'test comment2' }] test_fw = self._add_test_fw_with_comments_to_db() latest_comments = self.db_frontend_interface.get_latest_comments() comments.sort(key=lambda x: x['time'], reverse=True) for i in range(len(comments)): time, author, comment, uid = comments[i]['time'], comments[i][ 'author'], comments[i]['comment'], test_fw.get_uid() self.assertEqual(latest_comments[i]['time'], time) self.assertEqual(latest_comments[i]['author'], author) self.assertEqual(latest_comments[i]['comment'], comment) self.assertEqual(latest_comments[i]['uid'], uid) def test_remove_element_from_array_in_field(self): test_fw = self._add_test_fw_with_comments_to_db() retrieved_fw = self.db_backend_interface.get_object(test_fw.get_uid()) self.assertEqual(len(retrieved_fw.comments), 2, 'comments were not saved correctly') self.db_frontend_editing.remove_element_from_array_in_field( test_fw.get_uid(), 'comments', {'time': '1234567899'}) retrieved_fw = self.db_backend_interface.get_object(test_fw.get_uid()) self.assertEqual(len(retrieved_fw.comments), 1, 'comment was not deleted') def test_delete_comment(self): test_fw = self._add_test_fw_with_comments_to_db() retrieved_fw = self.db_backend_interface.get_object(test_fw.get_uid()) self.assertEqual(len(retrieved_fw.comments), 2, 'comments were not saved correctly') self.db_frontend_editing.delete_comment(test_fw.get_uid(), '1234567899') retrieved_fw = self.db_backend_interface.get_object(test_fw.get_uid()) self.assertEqual(len(retrieved_fw.comments), 1, 'comment was not deleted') def _add_test_fw_with_comments_to_db(self): test_fw = create_test_firmware() comments = [{ 'time': '1234567890', 'author': 'author1', 'comment': 'test comment' }, { 'time': '1234567899', 'author': 'author2', 'comment': 'test comment2' }] test_fw.comments.extend(comments) self.db_backend_interface.add_object(test_fw) return test_fw
class TestStorageDbInterfaceBackend(unittest.TestCase): @classmethod def setUpClass(cls): cls._config = get_config_for_testing(TMP_DIR) cls.mongo_server = MongoMgr(config=cls._config) def setUp(self): self.db_interface = MongoInterfaceCommon(config=self._config) self.db_interface_backend = BackEndDbInterface(config=self._config) self.test_firmware = create_test_firmware() self.test_yara_match = { 'rule': 'OpenSSH', 'tags': [], 'namespace': 'default', 'strings': [(0, '$a', b'OpenSSH')], 'meta': { 'description': 'SSH library', 'website': 'http://www.openssh.com', 'open_source': True, 'software_name': 'OpenSSH' }, 'matches': True } self.test_fo = create_test_file_object() def tearDown(self): self.db_interface.client.drop_database( self._config.get('data_storage', 'main_database')) self.db_interface_backend.shutdown() self.db_interface.shutdown() gc.collect() @classmethod def tearDownClass(cls): cls.mongo_server.shutdown() TMP_DIR.cleanup() def _get_all_firmware_uids(self): uid_list = [] tmp = self.db_interface.firmwares.find() for item in tmp: uid_list.append(item['_id']) return uid_list def test_add_firmware(self): self.db_interface_backend.add_firmware(self.test_firmware) self.assertGreater(len(self._get_all_firmware_uids()), 0, 'No entry added to DB') recoverd_firmware_entry = self.db_interface_backend.firmwares.find_one( ) self.assertAlmostEqual(recoverd_firmware_entry['submission_date'], time(), msg='submission time not set correctly', delta=5.0) def test_add_and_get_firmware(self): self.db_interface_backend.add_firmware(self.test_firmware) result_backend = self.db_interface_backend.get_firmware( self.test_firmware.uid) self.assertIsNotNone(result_backend.binary, 'binary not set in backend result') result_common = self.db_interface.get_firmware(self.test_firmware.uid) self.assertIsNone(result_common.binary, 'binary set in common result') self.assertEqual(result_common.size, 787, 'file size not correct in common') self.assertIsInstance(result_common.tags, dict, 'tag field type not correct') def test_add_and_get_file_object(self): self.db_interface_backend.add_file_object(self.test_fo) result_backend = self.db_interface_backend.get_file_object( self.test_fo.uid) self.assertIsNotNone(result_backend.binary, 'binary not set in backend result') result_common = self.db_interface.get_file_object(self.test_fo.uid) self.assertIsNone(result_common.binary, 'binary set in common result') self.assertEqual(result_common.size, 62, 'file size not correct in common') def test_update_firmware(self): first_dict = { 'stub_plugin': { 'result': 0 }, 'other_plugin': { 'field': 'day' } } second_dict = {'stub_plugin': {'result': 1}} self.test_firmware.processed_analysis = first_dict self.db_interface_backend.add_firmware(self.test_firmware) self.assertEqual( 0, self.db_interface.get_object( self.test_firmware.uid).processed_analysis['stub_plugin'] ['result']) self.test_firmware.processed_analysis = second_dict self.db_interface_backend.add_firmware(self.test_firmware) self.assertEqual( 1, self.db_interface.get_object( self.test_firmware.uid).processed_analysis['stub_plugin'] ['result']) self.assertIn( 'other_plugin', self.db_interface.get_object( self.test_firmware.uid).processed_analysis.keys()) def test_update_file_object(self): first_dict = {'other_plugin': {'result': 0}} second_dict = {'stub_plugin': {'result': 1}} self.test_fo.processed_analysis = first_dict self.test_fo.files_included = {'file a', 'file b'} self.db_interface_backend.add_file_object(self.test_fo) self.test_fo.processed_analysis = second_dict self.test_fo.files_included = {'file b', 'file c'} self.db_interface_backend.add_file_object(self.test_fo) received_object = self.db_interface.get_object(self.test_fo.uid) self.assertEqual( 0, received_object.processed_analysis['other_plugin']['result']) self.assertEqual( 1, received_object.processed_analysis['stub_plugin']['result']) self.assertEqual(3, len(received_object.files_included)) def test_add_and_get_object_including_comment(self): comment, author, date, uid = 'this is a test comment!', 'author', '1473431685', self.test_fo.uid self.test_fo.comments.append({ 'time': str(date), 'author': author, 'comment': comment }) self.db_interface_backend.add_file_object(self.test_fo) retrieved_comment = self.db_interface.get_object(uid).comments[0] self.assertEqual(author, retrieved_comment['author']) self.assertEqual(comment, retrieved_comment['comment']) self.assertEqual(date, retrieved_comment['time']) def test_update_analysis_tag_no_firmware(self): self.db_interface_backend.add_file_object(self.test_fo) tag = {'value': 'yay', 'color': 'default', 'propagate': True} self.db_interface_backend.update_analysis_tags(self.test_fo.uid, plugin_name='dummy', tag_name='some_tag', tag=tag) processed_fo = self.db_interface_backend.get_object(self.test_fo.uid) assert not processed_fo.analysis_tags def test_update_analysis_tag_uid_not_found(self): self.db_interface_backend.update_analysis_tags(self.test_fo.uid, plugin_name='dummy', tag_name='some_tag', tag='should not matter') assert not self.db_interface_backend.get_object(self.test_fo.uid) def test_update_analysis_tag_bad_tag(self): self.db_interface_backend.add_firmware(self.test_firmware) self.db_interface_backend.update_analysis_tags(self.test_firmware.uid, plugin_name='dummy', tag_name='some_tag', tag='bad_tag') processed_firmware = self.db_interface_backend.get_object( self.test_firmware.uid) assert not processed_firmware.analysis_tags def test_update_analysis_tag_success(self): self.db_interface_backend.add_firmware(self.test_firmware) tag = {'value': 'yay', 'color': 'default', 'propagate': True} self.db_interface_backend.update_analysis_tags(self.test_firmware.uid, plugin_name='dummy', tag_name='some_tag', tag=tag) processed_firmware = self.db_interface_backend.get_object( self.test_firmware.uid) assert processed_firmware.analysis_tags assert processed_firmware.analysis_tags['dummy']['some_tag'] == tag def test_add_analysis_firmware(self): self.db_interface_backend.add_object(self.test_firmware) before = self.db_interface_backend.get_object( self.test_firmware.uid).processed_analysis self.test_firmware.processed_analysis['foo'] = {'bar': 5} self.db_interface_backend.add_analysis(self.test_firmware) after = self.db_interface_backend.get_object( self.test_firmware.uid).processed_analysis assert before != after assert 'foo' not in before assert 'foo' in after assert after['foo'] == {'bar': 5} def test_add_analysis_file_object(self): self.db_interface_backend.add_object(self.test_fo) self.test_fo.processed_analysis['foo'] = {'bar': 5} self.db_interface_backend.add_analysis(self.test_fo) analysis = self.db_interface_backend.get_object( self.test_fo.uid).processed_analysis assert 'foo' in analysis assert analysis['foo'] == {'bar': 5} def test_crash_add_analysis(self): with self.assertRaises(RuntimeError): self.db_interface_backend.add_analysis(dict()) with self.assertRaises(AttributeError): self.db_interface_backend._update_analysis(dict(), 'dummy', dict())
class TestStorageDbInterfaceFrontendEditing(unittest.TestCase): @classmethod def setUpClass(cls): cls._config = get_config_for_testing(TMP_DIR) cls.mongo_server = MongoMgr(config=cls._config) def setUp(self): self.db_frontend_editing = FrontendEditingDbInterface( config=self._config) self.db_frontend_interface = FrontEndDbInterface(config=self._config) self.db_backend_interface = BackEndDbInterface(config=self._config) def tearDown(self): self.db_frontend_editing.shutdown() self.db_frontend_interface.shutdown() self.db_backend_interface.client.drop_database( self._config.get('data_storage', 'main_database')) self.db_backend_interface.shutdown() gc.collect() @classmethod def tearDownClass(cls): cls.mongo_server.shutdown() TMP_DIR.cleanup() def test_add_comment(self): test_fw = create_test_firmware() self.db_backend_interface.add_object(test_fw) comment, author, uid, time = 'this is a test comment!', 'author', test_fw.uid, 1234567890 self.db_frontend_editing.add_comment_to_object(uid, comment, author, time) test_fw = self.db_backend_interface.get_object(uid) self.assertEqual(test_fw.comments[0], { 'time': str(time), 'author': author, 'comment': comment }) def test_get_latest_comments(self): comments = [{ 'time': '1234567890', 'author': 'author1', 'comment': 'test comment' }, { 'time': '1234567899', 'author': 'author2', 'comment': 'test comment2' }] test_fw = self._add_test_fw_with_comments_to_db() latest_comments = self.db_frontend_interface.get_latest_comments() comments.sort(key=lambda x: x['time'], reverse=True) for i, comment in enumerate(comments): assert latest_comments[i]['time'] == comment['time'] assert latest_comments[i]['author'] == comment['author'] assert latest_comments[i]['comment'] == comment['comment'] assert latest_comments[i]['uid'] == test_fw.uid def test_remove_element_from_array_in_field(self): test_fw = self._add_test_fw_with_comments_to_db() retrieved_fw = self.db_backend_interface.get_object(test_fw.uid) self.assertEqual(len(retrieved_fw.comments), 2, 'comments were not saved correctly') self.db_frontend_editing.remove_element_from_array_in_field( test_fw.uid, 'comments', {'time': '1234567899'}) retrieved_fw = self.db_backend_interface.get_object(test_fw.uid) self.assertEqual(len(retrieved_fw.comments), 1, 'comment was not deleted') def test_delete_comment(self): test_fw = self._add_test_fw_with_comments_to_db() retrieved_fw = self.db_backend_interface.get_object(test_fw.uid) self.assertEqual(len(retrieved_fw.comments), 2, 'comments were not saved correctly') self.db_frontend_editing.delete_comment(test_fw.uid, '1234567899') retrieved_fw = self.db_backend_interface.get_object(test_fw.uid) self.assertEqual(len(retrieved_fw.comments), 1, 'comment was not deleted') def _add_test_fw_with_comments_to_db(self): test_fw = create_test_firmware() comments = [{ 'time': '1234567890', 'author': 'author1', 'comment': 'test comment' }, { 'time': '1234567899', 'author': 'author2', 'comment': 'test comment2' }] test_fw.comments.extend(comments) self.db_backend_interface.add_object(test_fw) return test_fw def test_update_object_field(self): test_fw = create_test_firmware(vendor='foo') self.db_backend_interface.add_object(test_fw) result = self.db_frontend_editing.get_object(test_fw.uid) assert result.vendor == 'foo' self.db_frontend_editing.update_object_field(test_fw.uid, 'vendor', 'bar') result = self.db_frontend_editing.get_object(test_fw.uid) assert result.vendor == 'bar' def test_add_to_search_query_cache(self): query = '{"device_class": "Router"}' uid = create_uid(query) assert self.db_frontend_editing.add_to_search_query_cache(query) == uid assert self.db_frontend_editing.search_query_cache.find_one( {'_id': uid})['search_query'] == query # check what happens if search is added again assert self.db_frontend_editing.add_to_search_query_cache(query) == uid assert self.db_frontend_editing.search_query_cache.count_documents( {'_id': uid}) == 1