Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add large file upload support #33

Merged
merged 15 commits into from
Jan 9, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion LICENSE
Original file line number Diff line number Diff line change
@@ -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
Expand Down
6 changes: 6 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
18 changes: 18 additions & 0 deletions examples/scan_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion virustotal_python/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
from virustotal_python.virustotal import Virustotal
name = "virustotal-python"
from virustotal_python.virustotal import VirustotalError
name = "virustotal-python"
70 changes: 63 additions & 7 deletions virustotal_python/tests.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import virustotal_python
import pytest
import os.path
import subprocess
from time import sleep
from base64 import urlsafe_b64encode

Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"
9 changes: 7 additions & 2 deletions virustotal_python/virustotal.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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'."
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down