Merge pull request #124 from intezer/support-download-sha256-sample-w…

…ith-zip-password

Support download sha256 sample with zip password protected

CHANGES

@@ -1,3 +1,7 @@
+1.19.8
+------
+- Add download sha256 sample with zip password protected
+
 1.19.7
 ------
 - Fix system URL for EndpointAnalysis

intezer_sdk/__init__.py

@@ -1 +1 @@
-__version__ = '1.19.7'
+__version__ = '1.19.8'

intezer_sdk/_api.py

@@ -1,4 +1,5 @@
 import io
+import json
 import os
 from http import HTTPStatus
 from typing import Any
@@ -539,13 +540,18 @@ def get_url_result(self, url: str) -> dict:
 
         return result
 
-    def download_file_by_sha256(self, sha256: str, path: str = None, output_stream: IO = None) -> None:
+    def download_file_by_sha256(self,
+                                sha256: str,
+                                path: str = None,
+                                output_stream: IO = None,
+                                password_protected: str = None) -> None:
         """
         Download a file by its sha256.
 
         :param sha256: The sha256 of the file to download.
         :param path: The path to download the file to.
         :param output_stream: The output stream to write the file to.
+        :param password_protected: protect with password and download as zip.
         """
         if not path and not output_stream:
             raise ValueError('You must provide either path or output_stream')
@@ -559,25 +565,37 @@ def download_file_by_sha256(self, sha256: str, path: str = None, output_stream:
             elif os.path.isfile(path):
                 raise FileExistsError()
 
+        json_data = {'password_protected': password_protected} if password_protected else None
+
         response = self.api.request_with_refresh_expired_access_token('GET',
                                                                       f'/files/{sha256}/download',
-                                                                      stream=bool(path))
+                                                                      stream=bool(path),
+                                                                      data=json_data)
 
         raise_for_status(response)
-        if output_stream:
-            output_stream.write(response.content)
-        else:
-            if should_extract_name_from_request:
-                try:
-                    filename = response.headers['content-disposition'].split('filename=')[1]
-                except Exception:
-                    filename = f'{sha256}.sample'
 
-                path = os.path.join(path, filename)
-
-            with open(path, 'wb') as file:
+        if password_protected:
+            output_location = output_stream if output_stream else os.path.join(path,'{sha256}.zip')
+            with open(output_location, 'wb') as output:
                 for chunk in response.iter_content(chunk_size=8192):
-                    file.write(chunk)
+                    output.write(chunk)
+            if output_stream:
+                output_stream.seek(0)
+        else:
+            if output_stream:
+                output_stream.write(response.content)
+            else:
+                if should_extract_name_from_request:
+                    try:
+                        filename = response.headers['content-disposition'].split('filename=')[1]
+                    except Exception:
+                        filename = f'{sha256}.sample'
+
+                    path = os.path.join(path, filename)
+
+                with open(path, 'wb') as file:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        file.write(chunk)
 
     def index_by_sha256(self, sha256: str, index_as: IndexType, family_name: str = None) -> Response:
         """

intezer_sdk/analysis.py

@@ -227,15 +227,16 @@ def _init_sub_analyses(self):
             else:
                 self._sub_analyses.append(sub_analysis_object)
 
-    def download_file(self, path: str = None, output_stream: IO = None):
+    def download_file(self, path: str = None, output_stream: IO = None, password_protection: str = None):
         """
         Downloads the analysis's file.
 
         ``path`` or ``output_stream`` must be provided.
         :param path: A path to where to save the file, it can be either a directory or non-existing file path.
         :param output_stream: A file-like object to write the file's content to.
+        :param password_protection: set password protection to download file as zip with password.
         """
-        self._api.download_file_by_sha256(self.result()['sha256'], path, output_stream)
+        self._api.download_file_by_sha256(self.result()['sha256'], path, output_stream, password_protection)
 
     @property
     def iocs(self) -> dict:

intezer_sdk/sub_analysis.py

@@ -186,11 +186,12 @@ def get_string_related_samples(self,
                                           wait,
                                           wait_timeout)
 
-    def download_file(self, path: str = None, output_stream: IO = None):
+    def download_file(self, path: str = None, output_stream: IO = None, password_protection: str = None):
         """
         Downloads the analysis's file.
         `path` or `output_stream` must be provided.
         :param path: A path to where to save the file, it can be either a directory or non-existing file path.
         :param output_stream: A file-like object to write the file's content to.
+        :param password_protection: set password protection to download file as zip with password.
         """
-        self._api.download_file_by_sha256(self.sha256, path, output_stream)
+        self._api.download_file_by_sha256(self.sha256, path, output_stream, password_protection)