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]>

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

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.

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)

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,

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"

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"

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