Merge pull request #23 from GDATASoftwareAG/add-tracing-in-python-sdk

Add file_size to tracing, fixes, tests

python/.gitignore

@@ -1,2 +1,3 @@
 test-reports/
-*.egg-info
+*.egg-info
+eicar.txt

python/src/vaas/__init__.py

@@ -3,8 +3,8 @@
 :mod:`vaas` is a Python library for the VaaS-API."""
 
 __version__ = "0.0.1"
-__all__ = ["Vaas"]
+__all__ = ["Vaas", "VaasTracing"]
 
 __author__ = "G DATA CyberDefense AG <[email protected]>"
 
-from .vaas import Vaas
+from .vaas import Vaas, VaasTracing

python/src/vaas/vaas.py

@@ -19,7 +19,7 @@ class VaasTracing:
     def trace_hash_request(self, elapsed_in_seconds):
         """Trace hash request in seconds"""
 
-    def trace_upload_request(self, elapsed_in_seconds):
+    def trace_upload_request(self, elapsed_in_seconds, file_size):
         """Trace upload request in seconds"""
 
 
@@ -83,10 +83,9 @@ async def __for_sha256(self, sha256):
         }
         response_message = self.__response_message_for_guid(guid)
         await self.websocket.send(json.dumps(verdict_request))
-        response_message.add_done_callback(
-            lambda _: self.tracing.trace_hash_request(time.time() - start)
-        )
-        return await response_message
+        result = await response_message
+        self.tracing.trace_hash_request(time.time() - start)
+        return result
 
     def __response_message_for_guid(self, guid):
         result = Future()
@@ -116,9 +115,7 @@ async def for_buffer(self, buffer):
             response_message = self.__response_message_for_guid(guid)
             self.__upload(token, url, buffer)
             verdict = (await response_message).get("verdict")
-            response_message.add_done_callback(
-                lambda _: self.tracing.trace_upload_request(time.time() - start)
-            )
+            self.tracing.trace_upload_request(time.time() - start, len(buffer))
 
         return verdict
 

python/tests/test_vaas.py

@@ -2,10 +2,9 @@
 import base64
 import os
 import unittest
+from unittest.mock import MagicMock, ANY
 from dotenv import load_dotenv
-
-from src.vaas import Vaas
-
+from src.vaas import Vaas, VaasTracing
 
 load_dotenv()
 
@@ -62,6 +61,17 @@ async def test_for_file_returns_verdict(self):
             verdict = await vaas.for_file("eicar.txt")
             self.assertEqual(verdict, "Malicious")
 
+    async def test_for_buffer_traces(self):
+        tracing = VaasTracing()
+        tracing.trace_hash_request = MagicMock()
+        tracing.trace_upload_request = MagicMock()
+        async with Vaas(tracing=tracing) as vaas:
+            await vaas.connect(TOKEN)
+            buffer = os.urandom(1024)
+            verdict = await vaas.for_buffer(buffer)
+            self.assertEqual(verdict, "Clean")
+            tracing.trace_hash_request.assert_called_with(ANY)
+            tracing.trace_upload_request.assert_called_with(ANY, 1024)
 
 if __name__ == "__main__":
     unittest.main()