Skip to content

Commit a729191

Browse files
smk762dbrennand
andauthored
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>
1 parent 7351473 commit a729191

7 files changed

Lines changed: 98 additions & 12 deletions

File tree

LICENSE

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
MIT License
22

3-
Copyright (c) 2021 dbrennand
3+
Copyright (c) 2022 dbrennand
44

55
Permission is hereby granted, free of charge, to any person obtaining a copy
66
of this software and associated documentation files (the "Software"), to deal

README.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ A Python library to interact with the public VirusTotal v2 and v3 APIs.
66
> [!NOTE]
77
>
88
> 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.
9+
>
10+
> It is highly recommended that you use the VirusTotal v3 API as it is the "default and encouraged way to programmatically interact with VirusTotal".
911
1012
# Dependencies and installation
1113

@@ -220,6 +222,8 @@ To run the tests, perform the following steps:
220222

221223
## Changelog
222224

225+
* 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.
226+
223227
* 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).
224228

225229
* 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:
250254

251255
* [**dbrennand**](https://github.com/dbrennand) - *Author*
252256

257+
* [**smk762**](https://github.com/smk762) - *Contributor*
258+
253259
## License
254260
This project is licensed under the MIT License - see the [LICENSE](LICENSE) for details.

examples/scan_file.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
* v2 documentation - https://developers.virustotal.com/reference#file-scan
99
1010
* v3 documentation - https://developers.virustotal.com/v3.0/reference#files-scan
11+
12+
* https://developers.virustotal.com/reference/files-upload-url
1113
"""
1214
from virustotal_python import Virustotal
1315
import os.path
@@ -35,3 +37,19 @@
3537
resp = vtotal.request("files", files=files, method="POST")
3638

3739
pprint(resp.data)
40+
41+
# v3 example for uploading a file larger than 32MB in size
42+
vtotal = Virustotal(API_KEY=API_KEY, API_VERSION="v3")
43+
44+
# Create dictionary containing the large file to send for multipart encoding upload
45+
large_file = {
46+
"file": (
47+
os.path.basename("/path/to/file/larger/than/32MB"),
48+
open(os.path.abspath("/path/to/file/larger/than/32MB"), "rb"),
49+
)
50+
}
51+
# Get URL to send a large file
52+
upload_url = vtotal.request("files/upload_url").data
53+
# Submit large file to VirusTotal v3 API for analysis
54+
resp = vtotal.request(upload_url, files=large_file, method="POST", large_file=True)
55+
pprint(resp.data)

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
setup(
77
name="virustotal-python",
8-
version="0.1.3",
8+
version="0.2.0",
99
author="dbrennand",
1010
description="A Python library to interact with the public VirusTotal v2 and v3 APIs.",
1111
long_description=long_description,

virustotal_python/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
from virustotal_python.virustotal import Virustotal
2-
name = "virustotal-python"
2+
from virustotal_python.virustotal import VirustotalError
3+
name = "virustotal-python"

virustotal_python/tests.py

Lines changed: 63 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import virustotal_python
22
import pytest
33
import os.path
4+
import subprocess
45
from time import sleep
56
from base64 import urlsafe_b64encode
67

@@ -25,6 +26,29 @@
2526
COMMENT_ID = "f-9f101483662fc071b7c10f81c64bb34491ca4a877191d464ff46fd94c7247115-07457619"
2627

2728

29+
@pytest.fixture()
30+
def large_file_fixture(request):
31+
"""Setup and teardown fixture for `test_large_file_v2` and `test_large_file_v3`."""
32+
# Create a large file of 33MB to submit to the VirusTotal API for analysis
33+
subprocess.run(
34+
["dd", "if=/dev/urandom", "of=dummy.dat", "bs=33M", "count=1"], check=True
35+
)
36+
37+
def teardown():
38+
"""Delete the large file created by the fixture."""
39+
subprocess.run(["rm", "dummy.dat"], check=True)
40+
41+
# Add finalizer function
42+
request.addfinalizer(teardown)
43+
44+
return {
45+
"file": (
46+
os.path.basename("dummy.dat"),
47+
open(os.path.abspath("dummy.dat"), "rb"),
48+
)
49+
}
50+
51+
2852
@pytest.fixture()
2953
def vtotal_v2(request):
3054
yield virustotal_python.Virustotal()
@@ -57,13 +81,6 @@ def test_file_scan_v2(vtotal_v2):
5781
"""
5882
Test for sending a file to the VirusTotal v2 API for analysis.
5983
"""
60-
# Create dictionary containing the file to send for multipart encoding upload
61-
files = {
62-
"file": (
63-
os.path.basename("virustotal_python/oldexamples.py"),
64-
open(os.path.abspath("virustotal_python/oldexamples.py"), "rb"),
65-
)
66-
}
6784
resp = vtotal_v2.request("file/scan", files=FILES, method="POST")
6885
data = resp.json()
6986
assert resp.response_code == 1
@@ -322,3 +339,42 @@ def test_contextmanager_v3():
322339
assert data["id"] == IP
323340
assert data["attributes"]["as_owner"] == "GOOGLE"
324341
assert data["attributes"]["country"] == "US"
342+
343+
344+
def test_large_file_v2(vtotal_v2, large_file_fixture):
345+
"""Test sending a large file to the VirusTotal v2 API for analysis.
346+
347+
https://developers.virustotal.com/v2.0/reference/file-scan-upload-url
348+
349+
NOTE: Currently this test does not work and returns a HTTP 500 internal server error.
350+
351+
Please see: https://github.com/dbrennand/virustotal-python/pull/33#issuecomment-1008307393
352+
"""
353+
# Get URL to send large file
354+
upload_url = vtotal_v2.request("file/scan/upload_url").json()["upload_url"]
355+
# Expect VirustotalError due to HTTP 500 internal server error
356+
with pytest.raises(virustotal_python.VirustotalError):
357+
# Submit large file to VirusTotal v2 API for analysis
358+
resp = vtotal_v2.request(
359+
upload_url, files=large_file_fixture, method="POST", large_file=True
360+
)
361+
assert resp.status_code == 200
362+
data = resp.json()
363+
assert data["scan_id"]
364+
365+
366+
def test_large_file_v3(vtotal_v3, large_file_fixture):
367+
"""Test sending a large file to the VirusTotal v3 API for analysis.
368+
369+
https://developers.virustotal.com/reference/files-upload-url
370+
"""
371+
# Get URL to send large file
372+
upload_url = vtotal_v3.request("files/upload_url").data
373+
# Submit large file to VirusTotal v3 API for analysis
374+
resp = vtotal_v3.request(
375+
upload_url, files=large_file_fixture, method="POST", large_file=True
376+
)
377+
assert resp.status_code == 200
378+
data = resp.data
379+
assert data["id"]
380+
assert data["type"] == "analysis"

virustotal_python/virustotal.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
"""
22
MIT License
33
4-
Copyright (c) 2021 dbrennand
4+
Copyright (c) 2022 dbrennand
55
66
Permission is hereby granted, free of charge, to any person obtaining a copy
77
of this software and associated documentation files (the "Software"), to deal
@@ -243,7 +243,7 @@ def __init__(
243243
:param TIMEOUT: A float for the amount of time to wait in seconds for the HTTP request before timing out.
244244
:raises ValueError: Raises ValueError when no API_KEY is provided or the API_VERSION is invalid.
245245
"""
246-
self.VERSION = "0.1.3"
246+
self.VERSION = "0.2.0"
247247
if API_KEY is None:
248248
raise ValueError(
249249
"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(
294294
json: dict = None,
295295
files: dict = None,
296296
method: str = "GET",
297+
large_file: bool = False,
297298
) -> Tuple[dict, VirustotalResponse]:
298299
"""
299300
Make a request to the VirusTotal API.
@@ -304,12 +305,16 @@ def request(
304305
:param json: A dictionary containing the JSON payload to send with the request.
305306
:param files: A dictionary containing the file for multipart encoding upload. (E.g: {'file': ('filename', open('filename.txt', 'rb'))})
306307
:param method: The request method to use.
308+
:param large_file: If a file is larger than 32MB, a custom generated upload URL is required.
309+
If this param is set to `True`, this URL can be set via the resource param.
307310
:returns: A dictionary containing the HTTP response code (resp_code) and JSON response (json_resp) if self.COMPATIBILITY_ENABLED is True.
308311
Otherwise, a VirustotalResponse class object is returned. If a HTTP status not equal to 200 occurs. Then a VirustotalError class object is returned.
309312
:raises Exception: Raise Exception when an unsupported method is provided.
310313
"""
311314
# Create API endpoint
312315
endpoint = f"{self.BASEURL}{resource}"
316+
if large_file:
317+
endpoint = resource
313318
# If API version being used is v2, add the API key to params
314319
if self.API_VERSION == "v2":
315320
params["apikey"] = self.API_KEY

0 commit comments

Comments
 (0)