class TestStorageDbInterfaceBackend(unittest.TestCase): def setUp(self): self._config = get_config_for_testing(TMP_DIR) self.mongo_server = MongoMgr(config=self._config) 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() self.mongo_server.shutdown() TMP_DIR.cleanup() gc.collect() 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.get_uid()) self.assertIsNotNone(result_backend.binary, 'binary not set in backend result') result_common = self.db_interface.get_firmware( self.test_firmware.get_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.get_uid()) self.assertIsNotNone(result_backend.binary, 'binary not set in backend result') result_common = self.db_interface.get_file_object( self.test_fo.get_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.get_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.get_uid()).processed_analysis['stub_plugin'] ['result']) self.assertIn( 'other_plugin', self.db_interface.get_object( self.test_firmware.get_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.get_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.get_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'])
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())