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
