feat(client): Download password protected zip files (#91)

tests/test_client.py

@@ -307,6 +307,97 @@ def test_download_file_with_error(httpserver):
   assert e_info.value.args[1] == "Resource not found."
 
 
+def test_download_zip_file(httpserver):
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files',
+      method='POST',
+      headers={'X-Apikey': 'dummy_api_key'},
+      data=json.dumps({'data': {'hashes': ['h1', 'h2'], 'password': 'pass'}})
+  ).respond_with_json({
+        'data': {
+          'id': '1234',
+          'type': 'zip_file',
+          'attributes': {'status': 'starting'}
+        }})
+
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files/1234',
+      method='GET',
+      headers={'x-apikey': 'dummy_api_key'}
+  ).respond_with_json(
+      {'data': {
+        'id': '1234',
+        'type': 'zip_file',
+        'attributes': {'status': 'creating'}
+      }})
+
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files/1234',
+      method='GET',
+      headers={'x-apikey': 'dummy_api_key'}
+  ).respond_with_json({
+      'data': {
+        'id': '1234',
+        'type': 'zip_file',
+        'attributes': {'status': 'finished'}
+      }})
+
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files/1234/download',
+      method='GET',
+      headers={'x-apikey': 'dummy_api_key'}
+  ).respond_with_data('filecontent')
+
+  with new_client(httpserver) as client:
+    with io.BytesIO() as f:
+      client.download_zip_files(['h1', 'h2'], f, 'pass', 1)
+      f.seek(0)
+      assert f.read() == b'filecontent'
+
+
+def test_download_zip_file_error_creating_file(httpserver):
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files',
+      method='POST',
+      headers={'X-Apikey': 'dummy_api_key'},
+      data=json.dumps({'data': {'hashes': ['h1', 'h2'], 'password': 'pass'}})
+  ).respond_with_json({
+        'data': {
+          'id': '1234',
+          'type': 'zip_file',
+          'attributes': {'status': 'starting'}
+        }})
+
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files/1234',
+      method='GET',
+      headers={'x-apikey': 'dummy_api_key'}
+  ).respond_with_json(
+      {'data': {
+        'id': '1234',
+        'type': 'zip_file',
+        'attributes': {'status': 'creating'}
+      }})
+
+  httpserver.expect_ordered_request(
+      '/api/v3/intelligence/zip_files/1234',
+      method='GET',
+      headers={'x-apikey': 'dummy_api_key'}
+  ).respond_with_json({
+      'data': {
+        'id': '1234',
+        'type': 'zip_file',
+        'attributes': {'status': 'timeout'}
+      }})
+
+  with new_client(httpserver) as client:
+    with io.BytesIO() as f:
+      with pytest.raises(APIError) as e_info:
+        client.download_zip_files(['h1', 'h2'], f, 'pass', 1)
+      assert e_info.value.args[0] == 'ServerError'
+      assert e_info.value.args[1] == 'Error when creating zip file: timeout'
+
+
 def test_scan_file(httpserver):
 
   upload_url = (

vt/client.py

@@ -304,9 +304,13 @@ def download_file(self, hash, file):
     """
     return make_sync(self.download_file_async(hash, file))
 
-  async def download_file_async(self, hash, file):
-    """Like :func:`download_file` but returns a coroutine."""
-    response = await self.get_async(f'/files/{hash}/download')
+  async def __download_async(self, endpoint, file):
+    """Downloads a file and writes it to file.
+
+    :param endpoint: endpoint to download the file from.
+    :param file: A file object where the downloaded file will be written to.
+    """
+    response = await self.get_async(endpoint)
     error = await self.get_error_async(response)
     if error:
       raise error
@@ -316,6 +320,60 @@ async def download_file_async(self, hash, file):
         break
       file.write(chunk)
 
+  async def download_file_async(self, hash, file):
+    """Like :func:`download_file` but returns a coroutine."""
+    await self.__download_async(f'/files/{hash}/download', file)
+
+  def download_zip_files(self, hashes, zipfile, password=None, sleep_time=20):
+    """Creates a bundle zip bundle containing one or multiple files.
+
+    The file identified by the hash will be written to the provided file
+    object. The file object must be opened in write binary mode ('wb').
+
+    :param hashes: list of file hashes (SHA-256, SHA-1 or MD5).
+    :param zipfile: A file object where the downloaded zip file
+      will be written to.
+    :param password: optional, a password to protect the zip file.
+    :param sleep_time: optional, seconds to sleep between each request.
+    """
+    return make_sync(
+        self.download_zip_files_async(hashes, zipfile, password, sleep_time))
+
+  async def download_zip_files_async(
+      self, hashes, zipfile, password=None, sleep_time=20):
+
+    data = {'hashes': hashes}
+    if password:
+      data['password'] = password
+
+    response = await self.post_async(
+        '/intelligence/zip_files', data=json.dumps({'data': data}))
+    error = await self.get_error_async(response)
+    if error:
+      raise error
+
+    res_data = (await response.json_async())['data']
+
+    # wait until the zip file is ready
+    while res_data['attributes']['status'] in ('creating', 'starting'):
+      await asyncio.sleep(sleep_time)
+      response = await self.get_async(
+          f'/intelligence/zip_files/{res_data["id"]}')
+      error = await self.get_error_async(response)
+      if error:
+        raise error
+      res_data = (await response.json_async())['data']
+
+    # check for errors creating the zip file
+    if res_data['attributes']['status'] != 'finished':
+      raise APIError(
+          'ServerError',
+          f'Error when creating zip file: {res_data["attributes"]["status"]}')
+
+    # download the zip file
+    await self.__download_async(
+        f'/intelligence/zip_files/{res_data["id"]}/download', zipfile)
+
   def feed(self, feed_type, cursor=None):
     """Returns an iterator for a VirusTotal feed.