From 4ec3c44732e94013fe58049ae29dfc2bd68eaae0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=B6hling=2C=20Maximilian?= Date: Mon, 11 Mar 2024 16:51:24 +0100 Subject: [PATCH 1/2] python: Add libmagic & detections in response --- python/src/vaas/vaas.py | 50 ++++++++++++++++----------------------- python/tests/test_vaas.py | 16 ++++++++----- 2 files changed, 31 insertions(+), 35 deletions(-) diff --git a/python/src/vaas/vaas.py b/python/src/vaas/vaas.py index 6e3ef177..046d8e2b 100644 --- a/python/src/vaas/vaas.py +++ b/python/src/vaas/vaas.py @@ -85,6 +85,7 @@ def get_ssl_context(url, verify): return ssl.create_default_context() return ssl._create_unverified_context() # pylint: disable=W0212 + def problem_details_to_error(problem_details): type = problem_details.get("type") details = problem_details.get("details") @@ -92,6 +93,17 @@ def problem_details_to_error(problem_details): return VaasClientError(details) return VaasServerError(details) + +def map_response(verdict_response): + return { + "Sha256": verdict_response.get("sha256"), + "Guid": verdict_response.get("guid"), + "Verdict": verdict_response.get("verdict"), + "LibMagic": verdict_response.get("libMagic"), + "Detections": verdict_response.get("detections"), + } + + class Vaas: """Verdict-as-a-Service client""" @@ -167,11 +179,7 @@ async def for_sha256(self, sha256, verdict_request_attributes=None, guid=None): verdict_response = await self.__for_sha256( sha256, verdict_request_attributes, guid ) - return { - "Sha256": verdict_response.get("sha256"), - "Guid": verdict_response.get("guid"), - "Verdict": verdict_response.get("verdict"), - } + return map_response(verdict_response) async def __for_stream(self, verdict_request_attributes=None, guid=None): if verdict_request_attributes is not None and not isinstance( @@ -254,7 +262,7 @@ async def __receive_loop(self): future.set_result(vaas_message) if message_kind == "Error": problem_details = vaas_message.get("problem_details") - if guid is None or problem_details is None : + if guid is None or problem_details is None: # Error: Server sent guid we are not waiting for, or problem details are null, ignore it continue future = self.results.get(guid) @@ -281,12 +289,8 @@ async def for_buffer(self, buffer, verdict_request_attributes=None, guid=None): verdict_response, buffer, len(buffer) ) - return { - "Sha256": verdict_response.get("sha256"), - "Guid": verdict_response.get("guid"), - "Verdict": verdict_response.get("verdict"), - } - + return map_response(verdict_response) + async def _for_unknown_buffer(self, response, buffer, buffer_len): start = time.time() guid = response.get("guid") @@ -315,10 +319,10 @@ async def for_stream(self, asyncBufferedReader, len, verdict_request_attributes= if verdict != "Unknown": raise VaasServerError("server returned verdict without receiving content") - + if token == None: raise VaasServerError("VerdictResponse missing UploadToken for stream upload") - + if url == None: raise VaasServerError("VerdictResponse missing URL for stream upload") @@ -332,11 +336,7 @@ async def for_stream(self, asyncBufferedReader, len, verdict_request_attributes= raise VaasTimeoutError() from ex self.tracing.trace_upload_request(time.time() - start, len) - return { - "Sha256": verdict_response.get("sha256"), - "Guid": verdict_response.get("guid"), - "Verdict": verdict_response.get("verdict"), - } + return map_response(verdict_response) async def for_file(self, path, verdict_request_attributes=None, guid=None): """Returns the verdict for a file""" @@ -356,11 +356,7 @@ async def for_file(self, path, verdict_request_attributes=None, guid=None): verdict_response, buffer, len(buffer) ) - return { - "Sha256": verdict_response.get("sha256"), - "Guid": verdict_response.get("guid"), - "Verdict": verdict_response.get("verdict"), - } + return map_response(verdict_response) async def __upload(self, token, upload_uri, buffer_or_file, content_length): jwt = PyJWT() @@ -411,8 +407,4 @@ async def for_url(self, url, verdict_request_attributes=None, guid=None): self.tracing.trace_url_request(time.time() - start) - return { - "Sha256": result.get("sha256"), - "Guid": result.get("guid"), - "Verdict": result.get("verdict"), - } + return map_response(result) diff --git a/python/tests/test_vaas.py b/python/tests/test_vaas.py index db2b5055..bd708193 100644 --- a/python/tests/test_vaas.py +++ b/python/tests/test_vaas.py @@ -161,7 +161,6 @@ async def test_for_buffer_returns_unknown_for_random_buffer(self): self.assertEqual(verdict["Sha256"].casefold(), sha256.casefold()) self.assertEqual(verdict["Guid"].casefold(), guid) - async def test_for_file_returns_verdict(self): async with await create_and_connect() as vaas: with open("eicar.txt", "wb") as f: @@ -173,7 +172,6 @@ async def test_for_file_returns_verdict(self): self.assertEqual(verdict["Sha256"].casefold(), sha256.casefold()) self.assertEqual(verdict["Guid"].casefold(), guid) - async def test_for_file_returns_verdict_if_no_cache_or_shed(self): options = get_disabled_options() @@ -187,7 +185,6 @@ async def test_for_file_returns_verdict_if_no_cache_or_shed(self): self.assertEqual(verdict["Sha256"].casefold(), sha256.casefold()) self.assertEqual(verdict["Guid"].casefold(), guid) - async def test_for_url_returns_malicious_for_eicar(self): options = get_disabled_options() async with await create_and_connect(options=options) as vaas: @@ -196,7 +193,6 @@ async def test_for_url_returns_malicious_for_eicar(self): self.assertEqual(verdict["Verdict"], "Malicious") self.assertEqual(verdict["Guid"].casefold(), guid) - async def test_for_url_without_shed_and_cache_returns_clean_for_robots_txt(self): options = get_disabled_options() async with await create_and_connect(options=options) as vaas: @@ -205,7 +201,6 @@ async def test_for_url_without_shed_and_cache_returns_clean_for_robots_txt(self) self.assertEqual(verdict["Verdict"], "Clean") self.assertEqual(verdict["Guid"].casefold(), guid) - async def test_for_url_without_cache_returns_clean_for_robots_txt(self): options = VaasOptions() options.use_cache = False @@ -238,7 +233,6 @@ async def test_for_buffer_traces(self): tracing.trace_hash_request.assert_called_with(ANY) tracing.trace_upload_request.assert_called_with(ANY, 1024) - async def test_for_empty_buffer_returns_clean(self): async with await create_and_connect() as vaas: buffer = bytes("", "utf-8") @@ -249,6 +243,16 @@ async def test_for_empty_buffer_returns_clean(self): self.assertEqual(verdict["Sha256"].casefold(), sha256.casefold()) self.assertEqual(verdict["Guid"].casefold(), guid) + async def test_for_url_returns_detections_and_mime_type(self): + options = get_disabled_options() + async with await create_and_connect(options=options) as vaas: + guid = str(uuid.uuid4()) + verdict = await vaas.for_url("https://secure.eicar.org/eicar.com.txt", guid=guid) + self.assertEqual(verdict["Verdict"], "Malicious") + self.assertIsNotNone(verdict["Detections"]) + self.assertEqual(verdict["LibMagic"]['fileType'], "EICAR virus test files") + self.assertEqual(verdict["LibMagic"]['mimeType'], "text/plain") + if __name__ == "__main__": unittest.main() From 4871ecbba7b68ae62fcf591632ebdb6688baf155 Mon Sep 17 00:00:00 2001 From: unglaublicherdude Date: Wed, 13 Mar 2024 14:18:30 +0100 Subject: [PATCH 2/2] changes the fieldnames to snake_case --- python/src/vaas/vaas.py | 2 +- python/tests/test_vaas.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/src/vaas/vaas.py b/python/src/vaas/vaas.py index 046d8e2b..e59dd26c 100644 --- a/python/src/vaas/vaas.py +++ b/python/src/vaas/vaas.py @@ -99,7 +99,7 @@ def map_response(verdict_response): "Sha256": verdict_response.get("sha256"), "Guid": verdict_response.get("guid"), "Verdict": verdict_response.get("verdict"), - "LibMagic": verdict_response.get("libMagic"), + "LibMagic": verdict_response.get("lib_magic"), "Detections": verdict_response.get("detections"), } diff --git a/python/tests/test_vaas.py b/python/tests/test_vaas.py index bd708193..f7083dfb 100644 --- a/python/tests/test_vaas.py +++ b/python/tests/test_vaas.py @@ -250,8 +250,8 @@ async def test_for_url_returns_detections_and_mime_type(self): verdict = await vaas.for_url("https://secure.eicar.org/eicar.com.txt", guid=guid) self.assertEqual(verdict["Verdict"], "Malicious") self.assertIsNotNone(verdict["Detections"]) - self.assertEqual(verdict["LibMagic"]['fileType'], "EICAR virus test files") - self.assertEqual(verdict["LibMagic"]['mimeType'], "text/plain") + self.assertEqual(verdict["LibMagic"]['file_type'], "EICAR virus test files") + self.assertEqual(verdict["LibMagic"]['mime_type'], "text/plain") if __name__ == "__main__":