Merge pull request #140 from intezer/feature/download-scan

feat: add download scan, stream seek to start after download

CHANGES

@@ -1,3 +1,8 @@
+1.21
+_______
+- Add download endpoint scan
+- Fix: downloaded file to stream seek to the start of the stream
+
 1.20
 _______
 - Support getting alerts history with filters from server.
@@ -6,6 +11,10 @@ _______
 _______
 - Add sandbox_machine_type to FileAnalysis
 
+1.19.16
+_______
+- Add timeout in seconds to IntezerApi
+
 1.19.15
 _______
 - Remove pip-system-certs from dependencies

intezer_sdk/__init__.py

@@ -1 +1 @@
-__version__ = '1.20'
+__version__ = '1.21'

intezer_sdk/_api.py

@@ -583,7 +583,7 @@ def download_file_by_sha256(self,
         raise_for_status(response)
 
         if password_protected:
-            output_location = output_stream if output_stream else os.path.join(path,'{sha256}.zip')
+            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):
                     output.write(chunk)
@@ -592,6 +592,7 @@ def download_file_by_sha256(self,
         else:
             if output_stream:
                 output_stream.write(response.content)
+                output_stream.seek(0, 0)
             else:
                 if should_extract_name_from_request:
                     try:
@@ -605,6 +606,34 @@ def download_file_by_sha256(self,
                     for chunk in response.iter_content(chunk_size=8192):
                         file.write(chunk)
 
+    def download_endpoint_scanner(self, platform: str = None, path: str = None, output_stream: IO = None):
+        if not path and not output_stream:
+            raise ValueError('You must provide either path or output_stream')
+        elif path and output_stream:
+            raise ValueError('You must provide either path or output_stream, not both')
+
+        json_data = {'platform': platform} if platform else None
+        response = self.api.request_with_refresh_expired_access_token(path=f'/endpoint-scanner/download',
+                                                                      method='GET',
+                                                                      data=json_data,
+                                                                      stream=True)
+        raise_for_status(response)
+        if output_stream:
+            output_stream.write(response.content)
+            output_stream.seek(0, 0)
+        else:
+            if os.path.isdir(path):
+                try:
+                    filename = response.headers['content-disposition'].split('filename=')[1]
+                except Exception:
+                    filename = 'IntezerScanner.exe' if platform == 'windows' else 'IntezerScanner'
+
+                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:
         """
         Index a file by its sha256 hash.

intezer_sdk/_endpoint_analysis_api.py

@@ -15,7 +15,7 @@ def __init__(self, scan_id: str, api: IntezerApiClient, max_upload_retries: int
         if not scan_id:
             raise ValueError('scan_id must be provided')
         self.scan_id = scan_id
-        self.base_url = f"{api.base_url.replace('/api/','')}/scans/scans/{scan_id}"
+        self.base_url = f"{api.base_url.replace('/api/', '')}/scans/scans/{scan_id}"
         self.max_upload_retries = max_upload_retries
 
     def request_with_refresh_expired_access_token(self, *args, **kwargs):
@@ -106,7 +106,6 @@ def upload_collected_binary(self, file_path: str, collected_from: str):
             except Exception:
                 raise
 
-
     def end_scan(self, scan_summary: dict):
         response = self.request_with_refresh_expired_access_token(path='/end',
                                                                   data=scan_summary,

intezer_sdk/endpoint_analysis.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import pathlib
+from typing import IO
 from typing import List
 from typing import Optional
 
@@ -281,3 +282,18 @@ def _send_memory_module_dump_info_and_upload_required(self):
                         logger.warning(f'Endpoint analysis: {self.analysis_id}, file {file_to_upload}.sample does not exist')
                 for future in concurrent.futures.as_completed(futures):
                     future.result()
+
+
+def download_endpoint_scanner(platform: str = None,
+                              path: str = None,
+                              output_stream: IO = None,
+                              api: IntezerApiClient = None):
+    """
+    Download the endpoint scanner to a file or stream.
+    :param platform: The platform to download the scanner for.
+    :param path: The path to save the scanner to.
+    :param output_stream: The stream to write the scanner to.
+    :param api: The API connection to Intezer.
+    """
+    api = api or get_global_api()
+    return IntezerApi(api).download_endpoint_scanner(platform, path, output_stream)