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_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)
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)
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_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_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)
def test_provider_with_task(self):
s = Stoq(
base_dir=utils.get_data_dir(),
source_archivers=['simple_archiver'],
providers=['simple_provider'],
connectors=['dummy_connector'],
)
dummy_connector = s.load_plugin('dummy_connector')
dummy_connector.save = create_autospec(dummy_connector.save)
simple_provider = s.load_plugin('simple_provider')
simple_provider.RETURN_PAYLOAD = False
simple_archiver = s.load_plugin('simple_archiver')
simple_archiver.PAYLOAD = b'This is a payload'
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)
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_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_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_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_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']
)
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_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)
async def test_dispatch_multiple_plugins(self):
multi_plugin_content = b'multi-plugin-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 = asynctest.create_autospec(
simple_worker.scan, return_value=None
)
dummy_worker = s.load_plugin('dummy_worker')
dummy_worker.scan = asynctest.create_autospec(
dummy_worker.scan, return_value=None
)
await s.scan(multi_plugin_content)
simple_worker.scan.assert_awaited_once()
self.assertEqual(len(simple_worker.scan.await_args[0]), 2)
dummy_worker.scan.assert_awaited_once()
self.assertEqual(len(dummy_worker.scan.await_args[0]), 2)
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_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_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_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_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_provider_with_start_deep_dispatch(self):
s = Stoq(
base_dir=utils.get_data_dir(),
source_archivers=['simple_archiver'],
providers=['simple_provider'],
connectors=['dummy_connector'],
)
dummy_connector = s.load_plugin('dummy_connector')
dummy_connector.save = create_autospec(dummy_connector.save)
simple_provider = s.load_plugin('simple_provider')
simple_provider.RETURN_PAYLOAD = True
simple_archiver = s.load_plugin('simple_archiver')
simple_archiver.PAYLOAD = b'This is a payload'
dummy_worker = s.load_plugin('dummy_worker')
dummy_worker.scan = create_autospec(dummy_worker.scan)
s.run(add_start_deep_dispatch=['dummy_worker'])
dummy_worker.scan.assert_called_once()
dummy_connector.save.assert_called_once()
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)
async def test_scan_with_duplicate_extracted_payloads(self):
s = Stoq(base_dir=utils.get_data_dir())
simple_worker = s.load_plugin('simple_worker')
simple_worker.EXTRACTED_DISPATCH_TO = ['extract_payload']
simple_worker.EXTRACTED_PAYLOAD = self.generic_content + b'more data'
extract_worker = s.load_plugin('extract_payload')
extract_worker.EXTRACTED_PAYLOAD = self.generic_content + b'more data'
response = await s.scan(
self.generic_content, add_start_dispatch=['simple_worker']
)
self.assertEqual(2, len(response.results))
self.assertEqual('simple_worker', response.results[0].plugins_run['workers'][0])
self.assertNotIn('extract_payload', response.results[0].plugins_run['workers'])
self.assertNotIn('simple_worker', response.results[1].plugins_run['workers'])
self.assertEqual(
'extract_payload', response.results[1].plugins_run['workers'][0]
)
self.assertEqual('simple_worker', response.results[1].extracted_by[0])
self.assertEqual('extract_payload', response.results[1].extracted_by[1])
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_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'])
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_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)
