implements for stream for python (#378)

python/.devcontainer/devcontainer.json

@@ -1,48 +1,22 @@
-// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
-// https://github.com/microsoft/vscode-dev-containers/tree/v0.233.0/containers/python-3
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/python
 {
 	"name": "Python 3",
-	"runArgs": [
-		"--network",
-		"host"
-	],
-	"build": {
-		"dockerfile": "Dockerfile",
-		"context": "..",
-		"args": {
-			// Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6
-			// Append -bullseye or -buster to pin to an OS version.
-			// Use -bullseye variants on local on arm64/Apple Silicon.
-			"VARIANT": "3.10-bullseye",
-			// Options
-			"NODE_VERSION": "lts/*"
-		}
-	},
-	// Set *default* container specific settings.json values on container create.
-	"settings": {
-		"python.defaultInterpreterPath": "/usr/local/bin/python",
-		"python.linting.enabled": true,
-		"python.linting.pylintEnabled": true,
-		"python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8",
-		"python.formatting.blackPath": "/usr/local/py-utils/bin/black",
-		"python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf",
-		"python.linting.banditPath": "/usr/local/py-utils/bin/bandit",
-		"python.linting.flake8Path": "/usr/local/py-utils/bin/flake8",
-		"python.linting.mypyPath": "/usr/local/py-utils/bin/mypy",
-		"python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle",
-		"python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle",
-		"python.linting.pylintPath": "/usr/local/py-utils/bin/pylint",
-		"python.formatting.provider": "black"
-	},
-	// Add the IDs of extensions you want installed when the container is created.
-	"extensions": [
-		"ms-python.python",
-		"ms-python.vscode-pylance"
-	],
+	// Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
+	"image": "mcr.microsoft.com/devcontainers/python:1-3.12-bullseye"
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+
 	// Use 'forwardPorts' to make a list of ports inside the container available locally.
 	// "forwardPorts": [],
+
 	// Use 'postCreateCommand' to run commands after the container is created.
-	"postCreateCommand": "pip3 install --user -r requirements.txt",
-	// Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
-	"remoteUser": "vscode"
-}
+	// "postCreateCommand": "pip3 install --user -r requirements.txt",
+
+	// Configure tool-specific properties.
+	// "customizations": {},
+
+	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "root"
+}

python/.github/dependabot.yml

@@ -0,0 +1,12 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for more information:
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+# https://containers.dev/guide/dependabot
+
+version: 2
+updates:
+ - package-ecosystem: "devcontainers"
+   directory: "/"
+   schedule:
+     interval: weekly

python/.vscode/settings.json

@@ -2,9 +2,9 @@
     "python.testing.unittestArgs": [
         "-v",
         "-s",
-        "./tests",
+        ".",
         "-p",
-        "test*.py"
+        "test_*.py"
     ],
     "python.testing.pytestEnabled": false,
     "python.testing.unittestEnabled": true

python/src/vaas/vaas.py

@@ -53,7 +53,7 @@ class VaasOptions:
 
     def __init__(self):
         self.use_cache = True
-        self.use_shed = True
+        self.use_hash_lookup = True
 
 
 def hash_file(filename):
@@ -173,6 +173,35 @@ async def for_sha256(self, sha256, verdict_request_attributes=None, guid=None):
             "Verdict": verdict_response.get("verdict"),
         }
 
+    async def __for_stream(self, verdict_request_attributes=None, guid=None):
+        if verdict_request_attributes is not None and not isinstance(
+            verdict_request_attributes, dict
+        ):
+            raise TypeError("verdict_request_attributes has to be dict(str, str)")
+
+        websocket = self.get_authenticated_websocket()
+        start = time.time()
+        guid = guid or str(uuid.uuid4())
+        verdict_request = {
+            "kind": "VerdictRequestForStream",
+            "session_id": self.session_id,
+            "guid": guid,
+            "use_hash_lookup": self.options.use_hash_lookup,
+            "use_cache": self.options.use_cache,
+            "verdict_request_attributes": verdict_request_attributes,
+        }
+        response_message = self.__response_message_for_guid(guid)
+        await websocket.send(json.dumps(verdict_request))
+
+        try:
+            result = await asyncio.wait_for(response_message, timeout=TIMEOUT)
+        except asyncio.TimeoutError as ex:
+            self.tracing.trace_hash_request_timeout()
+            raise VaasTimeoutError() from ex
+
+        self.tracing.trace_hash_request(time.time() - start)
+        return result
+
     async def __for_sha256(self, sha256, verdict_request_attributes=None, guid=None):
         if verdict_request_attributes is not None and not isinstance(
             verdict_request_attributes, dict
@@ -187,7 +216,7 @@ async def __for_sha256(self, sha256, verdict_request_attributes=None, guid=None)
             "sha256": sha256,
             "session_id": self.session_id,
             "guid": guid,
-            "use_shed": self.options.use_shed,
+            "use_hash_lookup": self.options.use_hash_lookup,
             "use_cache": self.options.use_cache,
             "verdict_request_attributes": verdict_request_attributes,
         }
@@ -257,7 +286,7 @@ async def for_buffer(self, buffer, verdict_request_attributes=None, guid=None):
             "Guid": verdict_response.get("guid"),
             "Verdict": verdict_response.get("verdict"),
         }
-
+    
     async def _for_unknown_buffer(self, response, buffer, buffer_len):
         start = time.time()
         guid = response.get("guid")
@@ -273,6 +302,42 @@ async def _for_unknown_buffer(self, response, buffer, buffer_len):
         self.tracing.trace_upload_request(time.time() - start, buffer_len)
         return verdict_response
 
+    async def for_stream(self, asyncBufferedReader, len, verdict_request_attributes=None, guid=None):
+        """Returns the verdict for a file"""
+
+        verdict_response = await self.__for_stream(
+            verdict_request_attributes, guid
+        )
+        guid = verdict_response.get("guid")
+        token = verdict_response.get("upload_token")
+        url = verdict_response.get("url")
+        verdict = verdict_response.get("verdict")
+
+        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")
+
+        start = time.time()
+        response_message = self.__response_message_for_guid(guid)
+        await self.__upload(token, url, asyncBufferedReader, len)
+        try:
+            verdict_response = await asyncio.wait_for(response_message, timeout=TIMEOUT)
+        except asyncio.TimeoutError as ex:
+            self.tracing.trace_upload_result_timeout(len)
+            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"),
+        }
+
     async def for_file(self, path, verdict_request_attributes=None, guid=None):
         """Returns the verdict for a file"""
 
@@ -331,7 +396,7 @@ async def for_url(self, url, verdict_request_attributes=None, guid=None):
             "url": url,
             "session_id": self.session_id,
             "guid": guid,
-            "use_shed": self.options.use_shed,
+            "use_hash_lookup": self.options.use_hash_lookup,
             "use_cache": self.options.use_cache,
             "verdict_request_attributes": verdict_request_attributes,
         }

python/tests/test_vaas.py

@@ -13,6 +13,7 @@
 from src.vaas import get_ssl_context
 from src.vaas.vaas import hash_file
 from src.vaas.vaas_errors import VaasConnectionClosedError, VaasInvalidStateError, VaasClientError
+import httpx
 
 load_dotenv()
 
@@ -39,7 +40,7 @@ async def create_and_connect(tracing=VaasTracing(), options=VaasOptions()):
 def get_disabled_options():
     options = VaasOptions()
     options.use_cache = False
-    options.use_shed = False
+    options.use_hash_lookup = False
     return options
 
 
@@ -90,6 +91,20 @@ async def test_use_for_sha256_if_not_connected(self):
                 "275a021bbfb6489e54d471899f7db9d1663fc695ec2fe2a2c4538aabf651fd0f"
             )
 
+    async def test_for_stream_eicar_form_url_returns_malicious(self):
+        async with await create_and_connect() as vaas:
+            guid = str(uuid.uuid4())
+            async with httpx.AsyncClient() as client:
+                response = await client.get("https://secure.eicar.org/eicar.com")
+                content_length = response.headers["Content-Length"]
+                verdict = await vaas.for_stream(
+                    response.aiter_bytes(),
+                    content_length,
+                    guid=guid
+                )
+                self.assertEqual(verdict["Verdict"], "Malicious")
+                self.assertEqual(verdict["Guid"].casefold(), guid)
+
     async def test_for_sha256_returns_malicious_for_eicar(self):
         async with await create_and_connect() as vaas:
             guid = str(uuid.uuid4())
@@ -118,7 +133,6 @@ async def test_for_sha256_returns_malicious_for_eicar(self):
     #         )
     #         self.assertEqual(verdict["Guid"].casefold(), guid)
 
-
     async def test_for_buffer_returns_malicious_for_eicar(self):
         async with await create_and_connect() as vaas:
             buffer = base64.b64decode(EICAR_BASE64)
@@ -129,6 +143,13 @@ async def test_for_buffer_returns_malicious_for_eicar(self):
             self.assertEqual(verdict["Sha256"].casefold(), sha256.casefold())
             self.assertEqual(verdict["Guid"].casefold(), guid)
 
+    async def test_for_stream_returns_malicious_for_eicar(self):
+        async with await create_and_connect() as vaas:
+            buffer = base64.b64decode(EICAR_BASE64)
+            guid = str(uuid.uuid4())
+            verdict = await vaas.for_buffer(buffer, guid=guid)
+            self.assertEqual(verdict["Verdict"], "Malicious")
+            self.assertEqual(verdict["Guid"].casefold(), guid)
 
     async def test_for_buffer_returns_unknown_for_random_buffer(self):
         async with await create_and_connect() as vaas:
@@ -188,7 +209,7 @@ async def test_for_url_without_shed_and_cache_returns_clean_for_random_beer(self
     async def test_for_url_without_cache_returns_clean_for_random_beer(self):
         options = VaasOptions()
         options.use_cache = False
-        options.use_shed = True
+        options.use_hash_lookup = True
         async with await create_and_connect(options=options) as vaas:
             guid = str(uuid.uuid4())
             verdict = await vaas.for_url("https://random-data-api.com/api/v2/beers", guid=guid)