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
def test_get_sha256(self):
h = helpers.get_sha256(self.generic_content)
self.assertEqual(
h,
'5cac4f980fedc3d3f1f99b4be3472c9b30d56523e632d151237ec9309048bda9')
async def scan_request(
self, request: Request, add_start_dispatch: Optional[List[str]] = None
) -> StoqResponse:
"""
Scan an individual payload
:param request: ``Request`` object of payload(s) to be scanned
:param add_start_dispatch: Force first round of scanning to use specified plugins
"""
self.log.debug(
f'Request received: RequestMeta: {helpers.dumps(request.request_meta, indent=0)}, '
f'start_dispatches: {helpers.dumps(add_start_dispatch, indent=0)}'
)
add_dispatches: Set[Tuple[Payload, str]] = set()
hashes_seen: DefaultDict[str, List] = defaultdict(list)
for idx, payload in enumerate(request.payloads):
if payload.results.payload_meta.should_scan and add_start_dispatch:
for plugin_name in add_start_dispatch:
add_dispatches.add((payload, plugin_name))
sha = helpers.get_sha256(payload.content)
hashes_seen[sha].append(idx)
for _recursion_level in range(1, self.max_recursion + 1):
self.log.debug(f'Beginning worker round {_recursion_level}')
scan_result = await self._execute_scan_round(request, add_dispatches)
if scan_result is None:
self.log.debug('No more plugins to run, completing scan')
break
extracted_payloads, add_dispatches = scan_result
# TODO: Add option for no-dedup
for extracted_payload in extracted_payloads:
payload_hash = helpers.get_sha256(extracted_payload.content)
if payload_hash not in hashes_seen:
self.log.debug(
f'Extracted payload {extracted_payload.results.payload_id} with '
f'PayloadMeta: {extracted_payload.results.payload_meta}'
)
request.payloads.append(extracted_payload)
hashes_seen[payload_hash].append(len(request.payloads) - 1)
payload_meta = extracted_payload.results.payload_meta
if _recursion_level >= self.max_recursion:
request.errors.append(
Error(
error=f'Final worker round ({_recursion_level}) reached, unable to process payload',
payload_id=extracted_payload.results.payload_id,
)
)
elif payload_meta.should_scan and payload_meta.dispatch_to:
add_dispatches.update(
(extracted_payload, add_dispatch)
for add_dispatch in payload_meta.dispatch_to
)
else:
payload_idx = hashes_seen[payload_hash]
for idx in payload_idx:
request.payloads[idx].results.extracted_by.extend(
extracted_payload.results.extracted_by
)
request.payloads[idx].results.extracted_from.extend(
extracted_payload.results.extracted_from
)
archive_tasks: List = []
if request.request_meta.archive_payloads:
for payload in request.payloads:
if not payload.results.payload_meta.should_archive:
continue
for archiver in self._loaded_dest_archiver_plugins.values():
archive_tasks.append(
self._apply_archiver(archiver, payload, request)
)
await asyncio.gather(*archive_tasks)
response = StoqResponse(request=request)
decorator_tasks = []
for decorator in self._loaded_decorator_plugins.values():
decorator_tasks.append(self._apply_decorator(decorator, response))
await asyncio.gather(*decorator_tasks)
connector_tasks = []
for connector in self._loaded_connector_plugins:
connector_tasks.append(self._apply_connector(connector, response))
await asyncio.gather(*connector_tasks)
return response