|
1 | 1 | import virustotal_python |
2 | 2 | import pytest |
3 | 3 | import os.path |
| 4 | +import subprocess |
4 | 5 | from time import sleep |
5 | 6 | from base64 import urlsafe_b64encode |
6 | 7 |
|
|
25 | 26 | COMMENT_ID = "f-9f101483662fc071b7c10f81c64bb34491ca4a877191d464ff46fd94c7247115-07457619" |
26 | 27 |
|
27 | 28 |
|
| 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 | + |
28 | 52 | @pytest.fixture() |
29 | 53 | def vtotal_v2(request): |
30 | 54 | yield virustotal_python.Virustotal() |
@@ -57,13 +81,6 @@ def test_file_scan_v2(vtotal_v2): |
57 | 81 | """ |
58 | 82 | Test for sending a file to the VirusTotal v2 API for analysis. |
59 | 83 | """ |
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 | | - } |
67 | 84 | resp = vtotal_v2.request("file/scan", files=FILES, method="POST") |
68 | 85 | data = resp.json() |
69 | 86 | assert resp.response_code == 1 |
@@ -322,3 +339,42 @@ def test_contextmanager_v3(): |
322 | 339 | assert data["id"] == IP |
323 | 340 | assert data["attributes"]["as_owner"] == "GOOGLE" |
324 | 341 | 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" |
0 commit comments