Add large file upload support (#33)

* Update virustotal.py * add param to docstring * refactor: Remove f string and trailing space. * feat: Add `VirustotalError` class. Add EOF. * feat: Add tests for `large_file` parameter. Add fixture `large_file_fixture` to setup and teardown a large file. * refactor: Remove testing pytest marks. * fix: Typo in test comment. * feat: Add example to upload a large file for analysis. * fix: Formatting using `black`. * chore: Prep for new release. Bump version to `0.2.0`. * docs: Add note to recommend v3 API use. Add changelog for `0.2.0`. * docs: Add link to PR. * docs: Add contributor. * chore: Bump version to `0.2.0`. * chore: Bump license year. Co-authored-by: dbrennand <[email protected]>
dbrennand · Jan 9, 2022 · a729191 · a729191
1 parent 7351473
commit a729191
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 12 deletions.
diff --git a/LICENSE b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021 dbrennand
+Copyright (c) 2022 dbrennand
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

diff --git a/README.md b/README.md
@@ -6,6 +6,8 @@ A Python library to interact with the public VirusTotal v2 and v3 APIs.
 > [!NOTE]
 >
 > This library is intended to be used with the public VirusTotal APIs. However, it *could* be used to interact with premium API endpoints as well.
+>
+> It is highly recommended that you use the VirusTotal v3 API as it is the "default and encouraged way to programmatically interact with VirusTotal".
 
 # Dependencies and installation
 
@@ -220,6 +222,8 @@ To run the tests, perform the following steps:
 
 ## Changelog
 
+* 0.2.0 - Added `large_file` parameter to `request` so a file larger than 32MB can be submitted for analysis. See [#33](https://github.com/dbrennand/virustotal-python/pull/33). Thank you @smk762.
+
 * 0.1.3 - Update urllib3 to 1.26.5 to address [CVE-2021-33503](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-33503).
 
 * 0.1.2 - Update dependencies for security vulnerability. Fixed an issue with some tests failing.
@@ -250,5 +254,7 @@ To run the tests, perform the following steps:
 
 * [**dbrennand**](https://github.com/dbrennand) - *Author*
 
+* [**smk762**](https://github.com/smk762) - *Contributor*
+
 ## License
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) for details.
diff --git a/examples/scan_file.py b/examples/scan_file.py
@@ -8,6 +8,8 @@
     * v2 documentation - https://developers.virustotal.com/reference#file-scan
 
     * v3 documentation - https://developers.virustotal.com/v3.0/reference#files-scan
+
+        * https://developers.virustotal.com/reference/files-upload-url
 """
 from virustotal_python import Virustotal
 import os.path
@@ -35,3 +37,19 @@
 resp = vtotal.request("files", files=files, method="POST")
 
 pprint(resp.data)
+
+# v3 example for uploading a file larger than 32MB in size
+vtotal = Virustotal(API_KEY=API_KEY, API_VERSION="v3")
+
+# Create dictionary containing the large file to send for multipart encoding upload
+large_file = {
+    "file": (
+        os.path.basename("/path/to/file/larger/than/32MB"),
+        open(os.path.abspath("/path/to/file/larger/than/32MB"), "rb"),
+    )
+}
+# Get URL to send a large file
+upload_url = vtotal.request("files/upload_url").data
+# Submit large file to VirusTotal v3 API for analysis
+resp = vtotal.request(upload_url, files=large_file, method="POST", large_file=True)
+pprint(resp.data)
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="virustotal-python",
-    version="0.1.3",
+    version="0.2.0",
     author="dbrennand",
     description="A Python library to interact with the public VirusTotal v2 and v3 APIs.",
     long_description=long_description,

diff --git a/virustotal_python/__init__.py b/virustotal_python/__init__.py
@@ -1,2 +1,3 @@
 from virustotal_python.virustotal import Virustotal
-name = "virustotal-python"
+from virustotal_python.virustotal import VirustotalError
+name = "virustotal-python"
diff --git a/virustotal_python/tests.py b/virustotal_python/tests.py
@@ -1,6 +1,7 @@
 import virustotal_python
 import pytest
 import os.path
+import subprocess
 from time import sleep
 from base64 import urlsafe_b64encode
 
@@ -25,6 +26,29 @@
 COMMENT_ID = "f-9f101483662fc071b7c10f81c64bb34491ca4a877191d464ff46fd94c7247115-07457619"
 
 
+@pytest.fixture()
+def large_file_fixture(request):
+    """Setup and teardown fixture for `test_large_file_v2` and `test_large_file_v3`."""
+    # Create a large file of 33MB to submit to the VirusTotal API for analysis
+    subprocess.run(
+        ["dd", "if=/dev/urandom", "of=dummy.dat", "bs=33M", "count=1"], check=True
+    )
+
+    def teardown():
+        """Delete the large file created by the fixture."""
+        subprocess.run(["rm", "dummy.dat"], check=True)
+
+    # Add finalizer function
+    request.addfinalizer(teardown)
+
+    return {
+        "file": (
+            os.path.basename("dummy.dat"),
+            open(os.path.abspath("dummy.dat"), "rb"),
+        )
+    }
+
+
 @pytest.fixture()
 def vtotal_v2(request):
     yield virustotal_python.Virustotal()
@@ -57,13 +81,6 @@ def test_file_scan_v2(vtotal_v2):
     """
     Test for sending a file to the VirusTotal v2 API for analysis.
     """
-    # Create dictionary containing the file to send for multipart encoding upload
-    files = {
-        "file": (
-            os.path.basename("virustotal_python/oldexamples.py"),
-            open(os.path.abspath("virustotal_python/oldexamples.py"), "rb"),
-        )
-    }
     resp = vtotal_v2.request("file/scan", files=FILES, method="POST")
     data = resp.json()
     assert resp.response_code == 1
@@ -322,3 +339,42 @@ def test_contextmanager_v3():
         assert data["id"] == IP
         assert data["attributes"]["as_owner"] == "GOOGLE"
         assert data["attributes"]["country"] == "US"
+
+
+def test_large_file_v2(vtotal_v2, large_file_fixture):
+    """Test sending a large file to the VirusTotal v2 API for analysis.
+
+    https://developers.virustotal.com/v2.0/reference/file-scan-upload-url
+
+    NOTE: Currently this test does not work and returns a HTTP 500 internal server error.
+
+    Please see: https://github.com/dbrennand/virustotal-python/pull/33#issuecomment-1008307393
+    """
+    # Get URL to send large file
+    upload_url = vtotal_v2.request("file/scan/upload_url").json()["upload_url"]
+    # Expect VirustotalError due to HTTP 500 internal server error
+    with pytest.raises(virustotal_python.VirustotalError):
+        # Submit large file to VirusTotal v2 API for analysis
+        resp = vtotal_v2.request(
+            upload_url, files=large_file_fixture, method="POST", large_file=True
+        )
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["scan_id"]
+
+
+def test_large_file_v3(vtotal_v3, large_file_fixture):
+    """Test sending a large file to the VirusTotal v3 API for analysis.
+
+    https://developers.virustotal.com/reference/files-upload-url
+    """
+    # Get URL to send large file
+    upload_url = vtotal_v3.request("files/upload_url").data
+    # Submit large file to VirusTotal v3 API for analysis
+    resp = vtotal_v3.request(
+        upload_url, files=large_file_fixture, method="POST", large_file=True
+    )
+    assert resp.status_code == 200
+    data = resp.data
+    assert data["id"]
+    assert data["type"] == "analysis"
diff --git a/virustotal_python/virustotal.py b/virustotal_python/virustotal.py
@@ -1,7 +1,7 @@
 """
 MIT License
 
-Copyright (c) 2021 dbrennand
+Copyright (c) 2022 dbrennand
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -243,7 +243,7 @@ def __init__(
         :param TIMEOUT: A float for the amount of time to wait in seconds for the HTTP request before timing out.
         :raises ValueError: Raises ValueError when no API_KEY is provided or the API_VERSION is invalid.
         """
-        self.VERSION = "0.1.3"
+        self.VERSION = "0.2.0"
         if API_KEY is None:
             raise ValueError(
                 "An API key is required to interact with the VirusTotal API.\nProvide one to the API_KEY parameter or by setting the environment variable 'VIRUSTOTAL_API_KEY'."
@@ -294,6 +294,7 @@ def request(
         json: dict = None,
         files: dict = None,
         method: str = "GET",
+        large_file: bool = False,
     ) -> Tuple[dict, VirustotalResponse]:
         """
         Make a request to the VirusTotal API.
@@ -304,12 +305,16 @@ def request(
         :param json: A dictionary containing the JSON payload to send with the request.
         :param files: A dictionary containing the file for multipart encoding upload. (E.g: {'file': ('filename', open('filename.txt', 'rb'))})
         :param method: The request method to use.
+        :param large_file: If a file is larger than 32MB, a custom generated upload URL is required.
+            If this param is set to `True`, this URL can be set via the resource param.
         :returns: A dictionary containing the HTTP response code (resp_code) and JSON response (json_resp) if self.COMPATIBILITY_ENABLED is True.
             Otherwise, a VirustotalResponse class object is returned. If a HTTP status not equal to 200 occurs. Then a VirustotalError class object is returned.
         :raises Exception: Raise Exception when an unsupported method is provided.
         """
         # Create API endpoint
         endpoint = f"{self.BASEURL}{resource}"
+        if large_file:
+            endpoint = resource
         # If API version being used is v2, add the API key to params
         if self.API_VERSION == "v2":
             params["apikey"] = self.API_KEY