From 49c07b6176730c7ee454986e694b9419ca51de71 Mon Sep 17 00:00:00 2001 From: "Meroujan.Antonyan" Date: Fri, 18 Oct 2024 11:46:13 +0000 Subject: [PATCH 1/4] Add max retry parameter to the execution --- msticpy/data/drivers/cybereason_driver.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/msticpy/data/drivers/cybereason_driver.py b/msticpy/data/drivers/cybereason_driver.py index 3901f7f1f..2750b010d 100644 --- a/msticpy/data/drivers/cybereason_driver.py +++ b/msticpy/data/drivers/cybereason_driver.py @@ -401,6 +401,7 @@ def __execute_query( page: int = 0, page_size: int = 2000, pagination_token: str = None, + max_retry: int = 3, ) -> Dict[str, Any]: """ Run query with pagination enabled. @@ -436,7 +437,8 @@ def __execute_query( headers = {} params = {"page": page, "itemsPerPage": page_size} status = None - while status != "SUCCESS": + cur_try = 0 + while status != "SUCCESS" and cur_try < max_retry: response = self.client.post( self.search_endpoint, json={**body, **pagination}, @@ -446,6 +448,7 @@ def __execute_query( response.raise_for_status() json_result = response.json() status = json_result["status"] + cur_try += 1 return json_result async def __run_threaded_queries( From 8c7278381e53e7ad0afac0bb8509be419ffbcc6b Mon Sep 17 00:00:00 2001 From: "Meroujan.Antonyan" Date: Tue, 22 Oct 2024 08:32:46 +0000 Subject: [PATCH 2/4] Add tests and raise exception after max retry --- msticpy/data/drivers/cybereason_driver.py | 8 +++ tests/data/drivers/test_cybereason_driver.py | 66 ++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/msticpy/data/drivers/cybereason_driver.py b/msticpy/data/drivers/cybereason_driver.py index 2750b010d..05edc2f15 100644 --- a/msticpy/data/drivers/cybereason_driver.py +++ b/msticpy/data/drivers/cybereason_driver.py @@ -449,6 +449,14 @@ def __execute_query( json_result = response.json() status = json_result["status"] cur_try += 1 + + if cur_try >= max_retry: + raise httpx.HTTPStatusError( + f"{status}: {json_result['message']}", + request=response.request, + response=response + ) + return json_result async def __run_threaded_queries( diff --git a/tests/data/drivers/test_cybereason_driver.py b/tests/data/drivers/test_cybereason_driver.py index 0dc8f8d66..317b91aad 100644 --- a/tests/data/drivers/test_cybereason_driver.py +++ b/tests/data/drivers/test_cybereason_driver.py @@ -12,6 +12,7 @@ import pytest import pytest_check as check import respx +import httpx from msticpy.data.core.query_defns import Formatters from msticpy.data.drivers.cybereason_driver import CybereasonDriver @@ -140,6 +141,46 @@ }, ] +_CR_PARTIAL_SUCCESS_RESULT = { + "data": { + "resultIdToElementDataMap": {}, + "suspicionsMap": {}, + "evidenceMap": {}, + "totalResults": 0, + "totalPossibleResults": 0, + "guessedPossibleResults": 0, + "queryLimits": { + "totalResultLimit": 1000, + "perGroupLimit": 100, + "perFeatureLimit": 100, + "groupingFeature": { + "elementInstanceType": "Process", + "featureName": "imageFileHash", + }, + "sortInGroupFeature": None, + }, + "queryTerminated": False, + "pathResultCounts": None, + "guids": [], + "paginationToken": None, + "executionUUID": None, + "quapiMeasurementData": { + "timeToGetGuids": [], + "timeToGetData": [], + "timeToGetAdditionalData": [], + "totalQuapiQueryTime": [], + "startTime": [], + "endTime": [], + }, + }, + "status": "PARTIAL_SUCCESS", + "hidePartialSuccess": False, + "message": "Received Non-OK status code HTTP/1.1 500 Internal Server Error", + "expectedResults": 0, + "failures": 0, + "failedServersInfo": None, +} + _CR_QUERY = { "query": """ { @@ -241,6 +282,31 @@ def test_query(driver): check.is_true(query.called) check.is_instance(data, pd.DataFrame) +@respx.mock +def test_partial_success_query(driver): + """Test query calling returns data in expected format.""" + connect = respx.post( + re.compile(r"^https://[a-zA-Z0-9\-]+\.cybereason\.net/login\.html") + ).respond(200) + query = respx.post( + re.compile( + r"^https://[a-zA-Z0-9\-]+\.cybereason\.net/rest/visualsearch/query/simple" + ) + ) + query.side_effect = [ + httpx.Response(200, json=_CR_PARTIAL_SUCCESS_RESULT), + httpx.Response(200, json=_CR_PARTIAL_SUCCESS_RESULT), + httpx.Response(200, json=_CR_PARTIAL_SUCCESS_RESULT), + httpx.Response(200, json=_CR_PARTIAL_SUCCESS_RESULT), + ] + with custom_mp_config(MP_PATH): + with pytest.raises(httpx.HTTPStatusError, match=r"PARTIAL_SUCCESS:.*"): + driver.connect() + driver.query('{"test": "test"}') + + check.is_true(connect.called or driver.connected) + check.is_true(query.called) + check.equal(query.call_count, 3) @respx.mock def test_paginated_query(driver): From 6e453dc4d707f3e94366e1b72397cab95eb094b9 Mon Sep 17 00:00:00 2001 From: "Meroujan.Antonyan" Date: Tue, 22 Oct 2024 08:57:07 +0000 Subject: [PATCH 3/4] doc and keyword params --- msticpy/data/drivers/cybereason_driver.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/msticpy/data/drivers/cybereason_driver.py b/msticpy/data/drivers/cybereason_driver.py index 05edc2f15..a61b92b9c 100644 --- a/msticpy/data/drivers/cybereason_driver.py +++ b/msticpy/data/drivers/cybereason_driver.py @@ -398,6 +398,7 @@ def _create_paginated_query_tasks( def __execute_query( self, body: Dict[str, Any], + *, page: int = 0, page_size: int = 2000, pagination_token: str = None, @@ -406,6 +407,8 @@ def __execute_query( """ Run query with pagination enabled. + :raises httpx.HTTPStatusError: if max_retry reached + Parameters ---------- body: Dict[str, Any] @@ -416,6 +419,8 @@ def __execute_query( Page number to query pagination_token: str Token of the current search + max_retry: int + Maximum retries in case of API no cuccess response Returns ------- From f21438dd5d5d4b27b28fbc1a3878682f90f152ce Mon Sep 17 00:00:00 2001 From: "Meroujan.Antonyan" Date: Tue, 22 Oct 2024 15:39:57 +0000 Subject: [PATCH 4/4] return cybereason instance name in dataframe --- msticpy/data/drivers/cybereason_driver.py | 2 +- tests/data/drivers/test_cybereason_driver.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/msticpy/data/drivers/cybereason_driver.py b/msticpy/data/drivers/cybereason_driver.py index a61b92b9c..bd9a1e87d 100644 --- a/msticpy/data/drivers/cybereason_driver.py +++ b/msticpy/data/drivers/cybereason_driver.py @@ -159,7 +159,7 @@ def query( ) else: df_result = self._format_result_to_dataframe(result=response) - + df_result["instance"] = self.instance return df_result def _exec_paginated_queries( diff --git a/tests/data/drivers/test_cybereason_driver.py b/tests/data/drivers/test_cybereason_driver.py index 317b91aad..10547e314 100644 --- a/tests/data/drivers/test_cybereason_driver.py +++ b/tests/data/drivers/test_cybereason_driver.py @@ -281,6 +281,7 @@ def test_query(driver): check.is_true(connect.called or driver.connected) check.is_true(query.called) check.is_instance(data, pd.DataFrame) + check.is_true("instance" in data.columns) @respx.mock def test_partial_success_query(driver):