From a7291910bb9846ae073848cef8657c487e9b5e32 Mon Sep 17 00:00:00 2001 From: smk762 <35845239+smk762@users.noreply.github.com> Date: Mon, 10 Jan 2022 05:31:52 +0800 Subject: [PATCH] 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 <52419383+dbrennand@users.noreply.github.com> --- LICENSE | 2 +- README.md | 6 +++ examples/scan_file.py | 18 +++++++++ setup.py | 2 +- virustotal_python/__init__.py | 3 +- virustotal_python/tests.py | 70 +++++++++++++++++++++++++++++---- virustotal_python/virustotal.py | 9 ++++- 7 files changed, 98 insertions(+), 12 deletions(-) diff --git a/LICENSE b/LICENSE index 9028fcc..e7f6d7e 100644 --- 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 index 724947b..1a9f366 100644 --- 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 index 469564c..10fb606 100644 --- 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 index 125c230..d05cfa4 100644 --- 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 index 848e22b..32823b0 100644 --- a/virustotal_python/__init__.py +++ b/virustotal_python/__init__.py @@ -1,2 +1,3 @@ from virustotal_python.virustotal import Virustotal -name = "virustotal-python" \ No newline at end of file +from virustotal_python.virustotal import VirustotalError +name = "virustotal-python" diff --git a/virustotal_python/tests.py b/virustotal_python/tests.py index 07c2536..d9c03f2 100644 --- 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 index 1a01dc4..0ba9bc4 100644 --- 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