def test_archiver_in_results(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['simple_archiver'])
response = s.scan(
self.generic_content, request_meta=RequestMeta(archive_payloads=True)
)
self.assertIn('simple_archiver', response.results[0].archivers)
self.assertIn('file_save_id', response.results[0].archivers['simple_archiver'])
def test_dont_dest_archive_yara(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['dummy_archiver'])
response = s.scan(
self.generic_content, request_meta=RequestMeta(archive_payloads=True)
)
# The yara rule 'similar_simple_rule' should set save = False
self.assertNotIn('dummy_archiver', response.results[0].archivers)
def test_worker_in_results(self):
s = Stoq(base_dir=utils.get_data_dir())
response = s.scan(self.generic_content, add_start_dispatch=['simple_worker'])
self.assertIn('simple_worker', response.results[0].workers[0])
self.assertIn(
'valuable_insight', response.results[0].workers[0]['simple_worker']
)
def test_start_dispatch(self):
s = Stoq(base_dir=utils.get_data_dir())
response = s.scan(self.generic_content, add_start_dispatch=['extract_random'])
self.assertIn('extract_random', response.results[0].plugins_run['workers'][0])
self.assertNotIn(
'extract_random', response.results[1].plugins_run['workers'][0]
)
def test_dispatcher_exception(self):
s = Stoq(base_dir=utils.get_data_dir(), dispatchers=['simple_dispatcher'])
simple_dispatcher = s.load_plugin('simple_dispatcher')
simple_dispatcher.RAISE_EXCEPTION = True
with self.assertRaises(Exception) as context:
simple_dispatcher.get_dispatches(task)
self.assertTrue('Test exception', context.exception)
def test_decorator_exception(self):
s = Stoq(base_dir=utils.get_data_dir(), decorators=['simple_decorator'])
simple_decorator = s.load_plugin('simple_decorator')
simple_decorator.RAISE_EXCEPTION = True
response = s.scan(self.generic_content)
self.assertEqual(len(response.errors), 1)
self.assertIn('Test exception', response.errors['simple_decorator'][0])
def test_decorator(self):
s = Stoq(base_dir=utils.get_data_dir(), decorators=['simple_decorator'])
_ = s.load_plugin('simple_decorator')
response = s.scan(self.generic_content)
self.assertIn('simple_decorator', response.decorators)
self.assertIn('simple_decoration', response.decorators['simple_decorator'])
self.assertEqual(len(response.errors), 0)
def test_deep_dispatch_nonexistent_plugin(self):
s = Stoq(base_dir=utils.get_data_dir())
response = s.scan(self.generic_content,
add_start_deep_dispatch=['this_plugin_doesnt_exist'])
self.assertNotIn('this_plugin_doesnt_exist',
response.results[0].plugins_run['workers'][0])
self.assertEqual(len(response.errors), 1)
def test_dispatch_from_worker(self):
s = Stoq(base_dir=utils.get_data_dir())
simple_worker = s.load_plugin('simple_worker')
simple_worker.DISPATCH_TO = ['extract_random']
response = s.scan(self.generic_content, add_start_dispatch=['simple_worker'])
self.assertIn('simple_worker', response.results[0].plugins_run['workers'][0])
self.assertIn('extract_random', response.results[1].plugins_run['workers'][0])
self.assertEqual('extract_random', response.results[2].extracted_by)
def test_source_archive(self):
s = Stoq(base_dir=utils.get_data_dir(), source_archivers=['simple_archiver'])
simple_archiver = s.load_plugin('simple_archiver')
simple_archiver.PAYLOAD = b'This is a payload'
task = ArchiverResponse(results={'path': '/tmp/123'})
payload = simple_archiver.get(task)
self.assertEqual('/tmp/123', payload.payload_meta.extra_data['path'])
self.assertEqual(payload.content, simple_archiver.PAYLOAD)
def test_dispatch_multiple_rules(self):
s = Stoq(base_dir=utils.get_data_dir(), dispatchers=['simple_dispatcher'])
s.load_plugin('simple_dispatcher').WORKERS = ['simple_worker', 'simple_worker']
simple_worker = s.load_plugin('simple_worker')
simple_worker.scan = create_autospec(simple_worker.scan, return_value=None)
s.scan(self.generic_content)
self.assertEqual(simple_worker.scan.call_count, 1)
self.assertEqual(len(simple_worker.scan.call_args[0]), 2)
def test_source_archiver_exception(self):
s = Stoq(base_dir=utils.get_data_dir(), source_archivers=['simple_archiver'])
simple_archiver = s.load_plugin('simple_archiver')
simple_archiver.RAISE_EXCEPTION = True
task = "This will fail"
with self.assertRaises(Exception) as context:
simple_archiver.get(task)
self.assertTrue('Test exception', context.exception)
def test_dest_archiver_exception(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['simple_archiver'])
simple_archiver = s.load_plugin('simple_archiver')
simple_archiver.RAISE_EXCEPTION = True
response = s.scan(self.generic_content)
self.assertIn('simple_archiver', response.results[0].plugins_run['archivers'])
self.assertEqual(len(response.errors), 1)
self.assertIn('Test exception', response.errors['simple_archiver'][0])
def test_connector_exception(self):
s = Stoq(base_dir=utils.get_data_dir(), connectors=['dummy_connector'])
dummy_connector = s.load_plugin('dummy_connector')
dummy_connector.save = create_autospec(
dummy_connector.save,
side_effect=RuntimeError('Unexpected exception'))
with self.assertRaises(Exception):
s.scan(self.generic_content)
def test_worker_errors(self):
s = Stoq(base_dir=utils.get_data_dir())
simple_worker = s.load_plugin('simple_worker')
simple_worker.RETURN_ERRORS = True
response = s.scan(self.generic_content, add_start_dispatch=['simple_worker'])
self.assertIn('simple_worker', response.results[0].plugins_run['workers'][0])
self.assertIn('simple_worker', response.results[0].workers[0])
self.assertEqual(len(response.errors), 1)
self.assertIn('Test error', response.errors['simple_worker'][0])
async def test_dest_archiver_errors(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['simple_archiver'])
simple_archiver = s.load_plugin('simple_archiver')
simple_archiver.RETURN_ERRORS = True
response = await s.scan(self.generic_content)
self.assertIn('simple_archiver', response.results[0].plugins_run['archivers'])
self.assertIn('simple_archiver', response.results[0].archivers)
self.assertEqual(len(response.errors), 1)
self.assertIn('Test error', response.errors[0].error)
def test_decorator_errors(self):
s = Stoq(base_dir=utils.get_data_dir(), decorators=['simple_decorator'])
simple_decorator = s.load_plugin('simple_decorator')
simple_decorator.RETURN_ERRORS = True
response = s.scan(self.generic_content)
self.assertIn('simple_decorator', response.decorators)
self.assertIn('simple_decoration', response.decorators['simple_decorator'])
self.assertEqual(len(response.errors), 1)
self.assertIn('Test error', response.errors['simple_decorator'][0])
def test_always_dispatch(self):
s = Stoq(base_dir=utils.get_data_dir(),
always_dispatch=['simple_worker'])
response = s.scan(self.generic_content)
self.assertIn('simple_worker', s._loaded_plugins)
self.assertIn('simple_worker',
response.results[0].plugins_run['workers'][0])
self.assertIn('simple_worker',
response.results[1].plugins_run['workers'][0])
async def test_dispatch_duplicate(self):
s = Stoq(base_dir=utils.get_data_dir(), dispatchers=['simple_dispatcher'])
s.load_plugin('simple_dispatcher').WORKERS = ['simple_worker', 'simple_worker']
simple_worker = s.load_plugin('simple_worker')
simple_worker.scan = asynctest.create_autospec(
simple_worker.scan, return_value=None
)
await s.scan(self.generic_content)
self.assertEqual(simple_worker.scan.await_count, 1)
self.assertEqual(len(simple_worker.scan.await_args[0]), 2)
async def test_worker_exception(self):
s = Stoq(base_dir=utils.get_data_dir())
simple_worker = s.load_plugin('simple_worker')
simple_worker.RAISE_EXCEPTION = True
response = await s.scan(
self.generic_content, add_start_dispatch=['simple_worker']
)
self.assertIn('simple_worker', response.results[0].plugins_run['workers'])
self.assertEqual(len(response.errors), 1)
self.assertIn('Test exception', response.errors[0].error)
def test_multi_providers(self):
s = Stoq(
base_dir=utils.get_data_dir(),
providers=['simple_provider', 'simple_provider2'],
connectors=['dummy_connector'],
)
dummy_connector = s.load_plugin('dummy_connector')
dummy_connector.save = create_autospec(dummy_connector.save)
s.run()
self.assertEqual(dummy_connector.save.call_count, 2)
def test_provider(self):
s = Stoq(
base_dir=utils.get_data_dir(),
providers=['simple_provider'],
connectors=['dummy_connector'],
)
dummy_connector = s.load_plugin('dummy_connector')
dummy_connector.save = create_autospec(dummy_connector.save)
s.run()
dummy_connector.save.assert_called_once()
async def test_provider_exception(self):
s = Stoq(base_dir=utils.get_data_dir(), providers=['simple_provider'])
simple_provider = s.load_plugin('simple_provider')
simple_provider.RAISE_EXCEPTION = True
logging.disable(logging.NOTSET)
with self.assertLogs(level='ERROR') as cm:
await s.run()
self.assertTrue(
cm.output[0].startswith('ERROR:stoq:Test exception, please ignore')
)
logging.disable(logging.CRITICAL)
def test_dispatch(self):
s = Stoq(base_dir=utils.get_data_dir(), dispatchers=['simple_dispatcher'])
dummy_worker = s.load_plugin('dummy_worker')
dummy_worker.scan = create_autospec(dummy_worker.scan, return_value=None)
response = s.scan(self.generic_content)
self.assertEqual(len(dummy_worker.scan.call_args[0]), 2)
self.assertEqual(
dummy_worker.scan.call_args[0][0].dispatch_meta['simple_dispatcher'],
{'test_key': 'Useful metadata info'},
)
self.assertIn('dummy_worker', response.results[0].plugins_run['workers'][0])
def test_dest_archive(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['dummy_archiver'])
dummy_archiver = s.load_plugin('dummy_archiver')
dummy_archiver.archive = create_autospec(
dummy_archiver.archive, return_value=None
)
response = s.scan(
self.generic_content, request_meta=RequestMeta(archive_payloads=True)
)
dummy_archiver.archive.assert_called_once()
self.assertIn('dummy_archiver', response.results[0].plugins_run['archivers'])
async def test_worker_should_scan(self):
s = Stoq(base_dir=utils.get_data_dir())
simple_worker = s.load_plugin('simple_worker')
simple_worker.SHOULD_SCAN = False
simple_worker.EXTRACTED_DISPATCH_TO = ['dummy_worker']
response = await s.scan(
self.generic_content, add_start_dispatch=['simple_worker']
)
# TODO: Extracted payloads are not being dispatched
self.assertIn('simple_worker', response.results[0].plugins_run['workers'][0])
self.assertEqual(2, len(response.results))
self.assertNotIn('dummy_worker', response.results[1].plugins_run['workers'])
async def test_scan_with_required_plugin_circular_reference(self):
s = Stoq(base_dir=utils.get_data_dir(), max_required_worker_depth=2000)
simple_worker = s.load_plugin('simple_worker')
simple_worker.EXTRACTED_DISPATCH_TO = ['simple_worker']
simple_worker.required_workers.add('dummy_worker')
dummy_worker = s.load_plugin('dummy_worker')
dummy_worker.required_workers.add('simple_worker')
response = await s.scan(
self.generic_content, add_start_dispatch=['simple_worker']
)
self.assertEqual(1, len(response.results))
self.assertIn('Circular', response.errors[0].error)
def test_split_results(self):
s = Stoq(base_dir=utils.get_data_dir())
response = s.scan(self.generic_content, add_start_dispatch=['multiclass_plugin', 'simple_worker'])
split_response = response.split()
self.assertEqual(len(split_response), 2)
for r in split_response:
if 'simple_worker' in r['results'][0]['workers'][0]:
self.assertNotIn('multiclass_plugin', r['results'][0]['workers'][0])
elif 'multiclass_plugin' in r['results'][0]['workers'][0]:
self.assertNotIn('simple_worker', r['results'][0]['workers'][0])
else:
raise Exception('required plugin not found in results')
def test_dont_dest_archive_request(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['dummy_archiver'])
dummy_archiver = s.load_plugin('dummy_archiver')
dummy_archiver.archive = Mock(return_value=None)
response = s.scan(
self.generic_content,
add_start_dispatch=['extract_random'],
request_meta=RequestMeta(archive_payloads=False),
)
dummy_archiver.archive.assert_not_called()
self.assertNotIn('dummy_archiver', response.results[0].plugins_run['archivers'])
self.assertNotIn('dummy_archiver', response.results[1].plugins_run['archivers'])
def test_dispatch_multiple_plugins2(self):
again_multi_plugin_content = b'again-multi-plugin-space-content'
s = Stoq(base_dir=utils.get_data_dir(), dispatchers=['simple_dispatcher'])
s.load_plugin('simple_dispatcher').WORKERS = ['simple_worker', 'dummy_worker']
simple_worker = s.load_plugin('simple_worker')
simple_worker.scan = create_autospec(simple_worker.scan, return_value=None)
dummy_worker = s.load_plugin('dummy_worker')
dummy_worker.scan = create_autospec(dummy_worker.scan, return_value=None)
s.scan(again_multi_plugin_content)
simple_worker.scan.assert_called_once()
self.assertEqual(len(simple_worker.scan.call_args[0]), 2)
dummy_worker.scan.assert_called_once()
self.assertEqual(len(dummy_worker.scan.call_args[0]), 2)