async def scan(
self,
content: bytes,
payload_meta: Optional[PayloadMeta] = None,
request_meta: Optional[RequestMeta] = None,
add_start_dispatch: Optional[List[str]] = None,
ratelimit: Optional[str] = None,
) -> StoqResponse:
"""
Wrapper for `scan_request` that creates a `Payload` object from bytes
:param content: Raw bytes to be scanned
:param payload_meta: Metadata pertaining to originating source
:param request_meta: Metadata pertaining to the originating request
:param add_start_dispatch: Force first round of scanning to use specified plugins
:param ratelimit: Rate limit calls to scan
"""
self.log.debug(
f'Content received ({len(content)} bytes): '
f'PayloadMeta: {helpers.dumps(payload_meta, indent=0)}, '
f'RequestMeta: {helpers.dumps(request_meta, indent=0)}'
)
payload_meta = payload_meta or PayloadMeta()
payload = Payload(content, payload_meta)
request_meta = request_meta or RequestMeta()
request = Request(payloads=[payload], request_meta=request_meta)
return await self.scan_request(request, add_start_dispatch)
async def test_dont_dest_archive_yara(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['dummy_archiver'])
response = await 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 scan_payload(
self,
payload: Payload,
request_meta: Optional[RequestMeta] = None,
add_start_dispatch: Optional[List[str]] = None,
add_start_deep_dispatch: Optional[List[str]] = None,
) -> StoqResponse:
"""
Scan an individual payload
:param payload: ``Payload`` object of data to be scanned
:param request_meta: Metadata pertaining to the originating request
:param add_start_dispatch: Force first round of scanning to use specified plugins
:param add_start_deep_dispatch: Force second round of scanning to use specified plugins
:return: Complete scan results
:rtype: StoqResponse
"""
request_meta = RequestMeta() if request_meta is None else request_meta
add_start_dispatch = [] if add_start_dispatch is None else add_start_dispatch
add_start_deep_dispatch = ([] if add_start_deep_dispatch is None else
add_start_deep_dispatch)
scan_results: List = []
errors: DefaultDict[str, List[str]] = defaultdict(list)
scan_queue = [(payload, add_start_dispatch, add_start_deep_dispatch)]
hashes_seen: Set[str] = set(helpers.get_sha256(payload.content))
for _recursion_level in range(self.max_recursion + 1):
next_scan_queue: List[Tuple[Payload, List[str], List[str]]] = []
for payload, add_dispatch, add_deep_dispatch in scan_queue:
payload_results, extracted, p_errors = self._single_scan(
payload, add_dispatch, add_deep_dispatch, request_meta)
scan_results.append(payload_results)
# TODO: Add option for no-dedup
for ex in extracted:
ex_hash = helpers.get_sha256(ex.content)
if ex_hash not in hashes_seen:
hashes_seen.add(ex_hash)
next_scan_queue.append(
(ex, ex.payload_meta.dispatch_to, []))
errors = helpers.merge_dicts(errors, p_errors)
scan_queue = next_scan_queue
response = StoqResponse(results=scan_results,
request_meta=request_meta,
errors=errors)
self._apply_decorators(response)
for connector in self._loaded_connector_plugins:
try:
connector.save(response)
except Exception:
self.log.exception(
f'Failed to save results using {connector.__module__}: {response}'
)
return response
async def test_archiver_in_results(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['simple_archiver'])
response = await 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'])
self.assertEqual(len(response.errors), 0)
async 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 = asynctest.create_autospec(
dummy_archiver.archive, return_value=None
)
response = await s.scan(
self.generic_content, request_meta=RequestMeta(archive_payloads=True)
)
dummy_archiver.archive.assert_awaited_once()
self.assertIn('dummy_archiver', response.results[0].plugins_run['archivers'])
async 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 = asynctest.CoroutineMock(return_value=None)
response = await s.scan(
self.generic_content,
add_start_dispatch=['extract_payload'],
request_meta=RequestMeta(archive_payloads=False),
)
dummy_archiver.archive.assert_not_awaited()
self.assertNotIn('dummy_archiver', response.results[0].plugins_run['archivers'])
self.assertNotIn('dummy_archiver', response.results[1].plugins_run['archivers'])
async def test_dont_dest_archive_payload(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['dummy_archiver'])
dummy_archiver = s.load_plugin('dummy_archiver')
dummy_archiver.archive = asynctest.create_autospec(
dummy_archiver.archive, return_value=None
)
response = await s.scan(
self.generic_content,
payload_meta=PayloadMeta(should_archive=False),
add_start_dispatch=['extract_payload'],
request_meta=RequestMeta(archive_payloads=True),
)
dummy_archiver.archive.assert_awaited_once()
self.assertNotIn('dummy_archiver', response.results[0].plugins_run['archivers'])
self.assertIn('dummy_archiver', response.results[1].plugins_run['archivers'])
async def test_reconstruct_all_subresponses(self):
# Construct a fake stoq_response as if it were generated from a file
# A.zip that contains two files, B.txt and C.zip, where C.zip contains D.txt
results = [
Payload(content=b'', payload_id='A.zip', payload_meta=PayloadMeta()),
Payload(
content=b'',
payload_id='B.txt',
payload_meta=PayloadMeta(),
extracted_from='A.zip',
extracted_by='fake',
),
Payload(
content=b'',
payload_id='C.zip',
payload_meta=PayloadMeta(),
extracted_from='A.zip',
extracted_by='fake',
),
Payload(
content=b'',
payload_id='D.txt',
payload_meta=PayloadMeta(),
extracted_from='C.zip',
extracted_by='fake',
),
]
request = Request(request_meta=RequestMeta(extra_data={'check': 'me'}))
payload_count = 1
for result in results:
result.results.workers['fake'] = f'result-{payload_count}'
result.results.plugins_run['workers'].append('fake')
request.payloads.append(result)
payload_count += 1
initial_response = StoqResponse(request)
s = Stoq(base_dir=utils.get_data_dir(), decorators=['simple_decorator'])
all_subresponses = [
r async for r in s.reconstruct_all_subresponses(initial_response)
]
# We expect there to be four "artificial" responses generated, one for
# each payload as the root.
self.assertEqual(len(all_subresponses), 4)
# We expect the first response to have all 4 payloads, the second response
# to have just the second payload, the third response to have the third
# and fourth payload, and the fourth response to have just the fourth payload
self.assertEqual(
[len(stoq_response.results) for stoq_response in all_subresponses],
[4, 1, 2, 1],
)
self.assertEqual(
[
stoq_response.results[0].workers['fake']
for stoq_response in all_subresponses
],
['result-1', 'result-2', 'result-3', 'result-4'],
)
self.assertTrue(
all(
'simple_decorator' in stoq_response.decorators
for stoq_response in all_subresponses
)
)
# Assert that they all have the same scan ID
self.assertEqual(
len({stoq_response.scan_id for stoq_response in all_subresponses}), 1
)
def test_requestmeta_to_str(self):
response = RequestMeta()
response_str = str(response)
response_dict = json.loads(response_str)
self.assertIsInstance(response_str, str)
self.assertIsInstance(response_dict, dict)
async def test_archiver_not_in_results(self):
s = Stoq(base_dir=utils.get_data_dir(), dest_archivers=['dummy_archiver'])
response = await s.scan(
self.generic_content, request_meta=RequestMeta(archive_payloads=True)
)
self.assertNotIn('dummy_archiver', response.results[0].archivers)