From 2599093a53819fdccaee804de4d2aa2dfb46cd97 Mon Sep 17 00:00:00 2001 From: noamblitz Date: Wed, 25 Sep 2024 18:07:44 +0200 Subject: [PATCH] use multiple file upload --- boefjes/boefjes/clients/bytes_client.py | 23 +++++++++++++++++++++++ boefjes/boefjes/job_handler.py | 8 ++++---- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/boefjes/boefjes/clients/bytes_client.py b/boefjes/boefjes/clients/bytes_client.py index b96369065c1..a04ca1b70d8 100644 --- a/boefjes/boefjes/clients/bytes_client.py +++ b/boefjes/boefjes/clients/bytes_client.py @@ -122,6 +122,29 @@ def save_raw(self, boefje_meta_id: str, raw: str | bytes, mime_types: Set[str] = return UUID(response.json()[file_name]) + @retry_with_login + def save_raw_files(self, boefje_meta_id: str, boefje_results: list[tuple[set, bytes | str]]) -> UUID: + file_name = "raw" # The name provides a key for all ids returned, so this is arbitrary as we only upload 1 file + + response = self._session.post( + "/bytes/raw", + json={ + "files": [ + { + "name": file_name, + "content": b64encode(raw if isinstance(raw, bytes) else raw.encode()).decode(), + "tags": list(mime_types), + } + for mime_types, raw in boefje_results + ] + }, + headers=self.headers, + params={"boefje_meta_id": str(boefje_meta_id)}, + ) + self._verify_response(response) + + return UUID(response.json()[file_name]) + @retry_with_login def get_raw(self, raw_data_id: str) -> bytes: response = self._session.get(f"/bytes/raw/{raw_data_id}", headers=self.headers) diff --git a/boefjes/boefjes/job_handler.py b/boefjes/boefjes/job_handler.py index 06f8b0eaa39..494879376f7 100644 --- a/boefjes/boefjes/job_handler.py +++ b/boefjes/boefjes/job_handler.py @@ -141,6 +141,7 @@ def handle(self, boefje_meta: BoefjeMeta) -> None: self.bytes_client.save_boefje_meta(boefje_meta) if boefje_results: + boefje_results_with_updated_mime_types = [] for boefje_added_mime_types, output in boefje_results: valid_mimetypes = set() for mimetype in boefje_added_mime_types: @@ -152,10 +153,9 @@ def handle(self, boefje_meta: BoefjeMeta) -> None: ) else: valid_mimetypes.add(mimetype) - raw_file_id = self.bytes_client.save_raw(boefje_meta.id, output, mime_types.union(valid_mimetypes)) - logger.info( - "Saved raw file %s for boefje %s[%s]", raw_file_id, boefje_meta.boefje.id, boefje_meta.id - ) + boefje_results_with_updated_mime_types.append((mime_types.union(valid_mimetypes), output)) + self.bytes_client.save_raw_files(boefje_meta.id, boefje_results) + logger.info("Saved raw file for boefje %s[%s]", boefje_meta.boefje.id, boefje_meta.id) else: logger.info("No results for boefje %s[%s]", boefje_meta.boefje.id, boefje_meta.id)