def test_payloadresults_to_str(self):
payload = Payload(self.generic_content)
response = PayloadResults.from_payload(payload)
response_str = str(response)
response_dict = json.loads(response_str)
self.assertIsInstance(response_str, str)
self.assertIsInstance(response_dict, dict)
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
initial_response = StoqResponse(
results=[
PayloadResults(
payload_id="A.zip",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result1"}],
plugins_run={"workers": [["fake"]]},
),
PayloadResults(
payload_id="B.txt",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result2"}],
plugins_run={"workers": [["fake"]]},
extracted_from="A.zip",
extracted_by="fake",
),
PayloadResults(
payload_id="C.zip",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result3"}],
plugins_run={"workers": [["fake"]]},
extracted_from="A.zip",
extracted_by="fake",
),
PayloadResults(
payload_id="D.txt",
size=0,
payload_meta=PayloadMeta(),
workers=[{"fake": "result4"}],
plugins_run={"workers": [["fake"]]},
extracted_from="C.zip",
extracted_by="fake",
),
],
request_meta=RequestMeta(extra_data={"check": "me"}),
errors={},
)
s = Stoq(base_dir=utils.get_data_dir(), decorators=["simple_decorator"])
all_subresponses = list(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[0]["fake"]
for stoq_response in all_subresponses
],
["result1", "result2", "result3", "result4"],
)
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 _single_scan(
self,
payload: Payload,
add_dispatch: List[str],
add_deep_dispatch: List[str],
request_meta: RequestMeta,
) -> Tuple[PayloadResults, List[Payload], DefaultDict[str, List[str]]]:
extracted = []
errors: DefaultDict[str, List[str]] = defaultdict(list)
dispatch_pass = 0
dispatches, dispatch_errors = self._get_dispatches(
payload, add_dispatch, request_meta)
if dispatch_errors:
errors = helpers.merge_dicts(errors, dispatch_errors)
for plugin_name in dispatches:
try:
plugin = self.load_plugin(plugin_name)
except Exception as e:
msg = 'worker:failed to load'
self.log.exception(msg)
errors[plugin_name].append(helpers.format_exc(e, msg=msg))
continue
# Normal dispatches are the "1st round" of scanning
payload.plugins_run['workers'][0].append(plugin_name)
try:
worker_response = plugin.scan(payload,
request_meta) # pyre-ignore[16]
except Exception as e:
msg = 'worker:failed to scan'
self.log.exception(msg)
errors[plugin_name].append(helpers.format_exc(e, msg=msg))
continue
if worker_response is None:
continue
if worker_response.results is not None:
# Normal dispatches are the "1st round" of scanning
payload.worker_results[0][
plugin_name] = worker_response.results
extracted.extend([
Payload(ex.content, ex.payload_meta, plugin_name,
payload.payload_id) for ex in worker_response.extracted
])
if worker_response.errors:
errors[plugin_name].extend(worker_response.errors)
while dispatch_pass < self.max_dispatch_passes:
dispatch_pass += 1
deep_dispatches, deep_dispatch_errors = self._get_deep_dispatches(
payload, add_deep_dispatch, request_meta)
if deep_dispatch_errors:
errors = helpers.merge_dicts(errors, deep_dispatch_errors)
if deep_dispatches:
# Add another entry for this round
payload.plugins_run['workers'].append([])
payload.worker_results.append({})
else:
break
for plugin_name in deep_dispatches:
try:
plugin = self.load_plugin(plugin_name)
except Exception as e:
msg = f'deep dispatch:failed to load (pass {dispatch_pass}/{self.max_dispatch_passes})'
self.log.exception(msg)
errors[plugin_name].append(helpers.format_exc(e, msg=msg))
continue
payload.plugins_run['workers'][dispatch_pass].append(
plugin_name)
try:
worker_response = plugin.scan( # pyre-ignore[16]
payload, request_meta)
except Exception as e:
msg = f'deep dispatch:failed to scan (pass {dispatch_pass}/{self.max_dispatch_passes})'
self.log.exception(msg)
errors[plugin_name].append(helpers.format_exc(e, msg=msg))
continue
if worker_response is None:
continue
if worker_response.results is not None:
payload.worker_results[dispatch_pass][
plugin_name] = worker_response.results
extracted.extend([
Payload(ex.content, ex.payload_meta, plugin_name,
payload.payload_id)
for ex in worker_response.extracted
])
if worker_response.errors:
errors[plugin_name].extend(worker_response.errors)
payload_results = PayloadResults.from_payload(payload)
if request_meta.archive_payloads and payload.payload_meta.should_archive:
for plugin_name, archiver in self._loaded_dest_archiver_plugins.items(
):
payload.plugins_run['archivers'].append(plugin_name)
try:
archiver_response = archiver.archive(payload, request_meta)
except Exception as e:
msg = 'archiver:failed to archive'
self.log.exception(msg)
errors[plugin_name].append(helpers.format_exc(e, msg=msg))
continue
if archiver_response is None:
continue
if archiver_response.results is not None:
payload_results.archivers[
plugin_name] = archiver_response.results
if archiver_response.errors:
errors[plugin_name].extend(archiver_response.errors)
return (payload_results, extracted, errors)