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_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_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_scan(self) -> None:
s = Stoq(plugin_dir_list=[self.plugin_dir])
plugin = s.load_plugin(self.plugin_name)
payload = Payload(self.generic_data)
response = plugin.scan(payload, RequestMeta())
self.assertIsInstance(response, WorkerResponse)
self.assertEqual('3:hMCE7pr3Kn:huJ6', response.results['ssdeep'])
async def test_scan(self) -> None:
s = Stoq(plugin_dir_list=[self.plugin_dir])
plugin = s.load_plugin(self.plugin_name)
payload = Payload(self.generic_data)
response = await plugin.scan(payload, Request())
self.assertIsInstance(response, WorkerResponse)
self.assertEqual('text/plain', response.results['mimetype'])
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_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_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_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_dispatcher(self) -> None:
s = Stoq(
plugin_dir_list=[self.plugin_dir],
plugin_opts={
self.plugin_name: {
'dispatch_rules': f'{self.data_dir}/dispatch_rules.yar'
}
},
)
plugin = s.load_plugin(self.plugin_name)
payload = Payload(self.generic_data)
response = plugin.get_dispatches(payload, RequestMeta())
self.assertIsInstance(response, DispatcherResponse)
self.assertIn('test_dispatch_plugin', response.plugin_names)
self.assertEqual(
'test_dispatch_rule', response.meta['test_dispatch_plugin']['rule']
)
self.assertIn(
'test_dispatch_plugin',
response.meta['test_dispatch_plugin']['meta']['plugin'],
)
self.assertIn('True', response.meta['test_dispatch_plugin']['meta']['save'])
self.assertEqual(
['tag1', 'tag2'], response.meta['test_dispatch_plugin']['tags']
)
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)
async def test_scan_notnil(self) -> None:
s = Stoq(plugin_dir_list=[self.plugin_dir])
plugin = s.load_plugin(self.plugin_name)
payload = Payload(self.entropy_not_nil)
response = await plugin.scan(payload, RequestMeta())
self.assertIsInstance(response, WorkerResponse)
self.assertEqual(1.584962500721156, response.results['entropy'])
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)
async def test_scan(self) -> None:
s = Stoq(plugin_dir_list=[str(self.plugin_dir)])
plugin = s.load_plugin(self.plugin_name)
payload = Payload(base64.b64encode(self.generic_data))
response = await plugin.scan(payload, RequestMeta())
self.assertIsInstance(response, WorkerResponse)
self.assertEqual(1, len(response.extracted))
self.assertEqual(self.generic_data, response.extracted[0].content)
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_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_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_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_scan_invalid_rule_file(self) -> None:
s = Stoq(
plugin_dir_list=[self.plugin_dir],
plugin_opts={
self.plugin_name: {'worker_rules': f'{self.data_dir}/nonexistent.yar'}
},
)
with self.assertRaises(StoqPluginException):
s.load_plugin(self.plugin_name)
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_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])
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_scan_invalid_rules(self) -> None:
s = Stoq(
plugin_dir_list=[self.plugin_dir],
plugin_opts={
self.plugin_name: {'worker_rules': f'{self.data_dir}/invalid_rules.yar'}
},
)
with self.assertRaises(yara.SyntaxError):
s.load_plugin(self.plugin_name)
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])
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()
def test_scan(self) -> None:
s = Stoq(plugin_dir_list=[self.plugin_dir])
plugin = s.load_plugin(self.plugin_name)
with open(f'{self.data_dir}/sample.pdf', 'rb') as f:
payload = Payload(f.read())
response = plugin.scan(payload, RequestMeta())
self.assertIsInstance(response, WorkerResponse)
self.assertIn('FileType', response.results)
self.assertEqual('PDF', response.results['FileType'])
self.assertEqual(6, response.results['PageCount'])
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)
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)
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'])
