From e86cea0cba235db6d2e37b3787da7bc1a20240c0 Mon Sep 17 00:00:00 2001 From: davidt99 Date: Tue, 7 May 2024 16:30:33 +0300 Subject: [PATCH] feat: add download scan, stream seek to start after download --- CHANGES | 9 ++++++++ intezer_sdk/__init__.py | 2 +- intezer_sdk/_api.py | 31 ++++++++++++++++++++++++++- intezer_sdk/_endpoint_analysis_api.py | 3 +-- intezer_sdk/endpoint_analysis.py | 16 ++++++++++++++ 5 files changed, 57 insertions(+), 4 deletions(-) diff --git a/CHANGES b/CHANGES index 2e40803..f15bccd 100644 --- a/CHANGES +++ b/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 diff --git a/intezer_sdk/__init__.py b/intezer_sdk/__init__.py index 0f2616a..6b50f52 100644 --- a/intezer_sdk/__init__.py +++ b/intezer_sdk/__init__.py @@ -1 +1 @@ -__version__ = '1.20' +__version__ = '1.21' diff --git a/intezer_sdk/_api.py b/intezer_sdk/_api.py index 84f4357..c44d5c8 100644 --- a/intezer_sdk/_api.py +++ b/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. diff --git a/intezer_sdk/_endpoint_analysis_api.py b/intezer_sdk/_endpoint_analysis_api.py index 765d753..2c739fe 100644 --- a/intezer_sdk/_endpoint_analysis_api.py +++ b/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, diff --git a/intezer_sdk/endpoint_analysis.py b/intezer_sdk/endpoint_analysis.py index 04dfdbb..498e1b4 100644 --- a/intezer_sdk/endpoint_analysis.py +++ b/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)