From 7f656cc2abae13dedd45fa677bb0043e8c0a3d3e Mon Sep 17 00:00:00 2001 From: Tyler Walch Date: Thu, 18 Nov 2021 14:07:45 -0500 Subject: [PATCH 1/9] calling the J1 sync results API --- jupiterone/client.py | 73 ++++++- jupiterone/constants.py | 18 ++ tests/test_query.py | 411 +++++++++++++++++++++++++++++++++++----- 3 files changed, 444 insertions(+), 58 deletions(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index a9be725..d814acf 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -7,6 +7,7 @@ import requests from retrying import retry +from warnings import warn from jupiterone.errors import ( JupiterOneClientError, @@ -22,7 +23,8 @@ DELETE_ENTITY, UPDATE_ENTITY, CREATE_RELATIONSHIP, - DELETE_RELATIONSHIP + DELETE_RELATIONSHIP, + CURSOR_QUERY_V1 ) def retry_on_429(exc): @@ -105,25 +107,48 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: raise JupiterOneApiError(content.get('errors')) return response.json() - elif response.status_code in [429, 500]: + elif response.status_code in [429, 503]: raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded') else: content = json.loads(response._content) raise JupiterOneApiError('{}:{}'.format(response.status_code, content.get('error'))) - def query_v1(self, query: str, **kwargs) -> Dict: - """ Performs a V1 graph query + def _cursor_query(self, query: str, cursor: str = None, include_deleted: bool = False) -> Dict: + """ Performs a V1 graph query using cursor pagination args: query (str): Query text - skip (int): Skip entity count - limit (int): Limit entity count + cursor (str): A pagination cursor for the initial query include_deleted (bool): Include recently deleted entities in query/search """ - skip: int = kwargs.pop('skip', J1QL_SKIP_COUNT) - limit: int = kwargs.pop('limit', J1QL_LIMIT_COUNT) - include_deleted: bool = kwargs.pop('include_deleted', False) + results: List = [] + while True: + variables = { + 'query': query, + 'includeDeleted': include_deleted + } + + if cursor is not None: + variables['cursor'] = cursor + + response = self._execute_query(query=CURSOR_QUERY_V1, variables=variables) + data = response['data']['queryV1'] + + if 'vertices' in data and 'edges' in data: + return data + + results.extend(data) + + if 'cursor' in response['data']['queryV1']: + print('cursor', response['data']['queryV1']['cursor'], len(data), len(results)) + cursor = response['data']['queryV1']['cursor'] + else: + break + + return {'data': results} + + def _limit_and_skip_query(self, query: str, skip: int = J1QL_SKIP_COUNT, limit: int = J1QL_LIMIT_COUNT, include_deleted: bool = False) -> Dict: results: List = [] page: int = 0 @@ -152,6 +177,36 @@ def query_v1(self, query: str, **kwargs) -> Dict: return {'data': results} + def query_v1(self, query: str, **kwargs) -> Dict: + """ Performs a V1 graph query + args: + query (str): Query text + skip (int): Skip entity count + limit (int): Limit entity count + cursor (str): A pagination cursor for the initial query + include_deleted (bool): Include recently deleted entities in query/search + """ + uses_limit_and_skip: bool = 'skip' in kwargs.keys() or 'limit' in kwargs.keys() + skip: int = kwargs.pop('skip', J1QL_SKIP_COUNT) + limit: int = kwargs.pop('limit', J1QL_LIMIT_COUNT) + include_deleted: bool = kwargs.pop('include_deleted', False) + cursor: str = kwargs.pop('cursor', None) + + if uses_limit_and_skip: + warn('limit and skip pagination is no longer a recommended method for pagination. To read more about using cursors checkout the JupiterOne documentation: https://support.jupiterone.io/hc/en-us/articles/360022722094#entityandrelationshipqueries', DeprecationWarning, stacklevel=2) + return self._limit_and_skip_query( + query=query, + skip=skip, + limit=limit, + include_deleted=include_deleted + ) + else: + return self._cursor_query( + query=query, + cursor=cursor, + include_deleted=include_deleted + ) + def create_entity(self, **kwargs) -> Dict: """ Creates an entity in graph. It will also update an existing entity. diff --git a/jupiterone/constants.py b/jupiterone/constants.py index 7b58dad..d8d350c 100644 --- a/jupiterone/constants.py +++ b/jupiterone/constants.py @@ -6,6 +6,24 @@ queryV1(query: $query, variables: $variables, dryRun: $dryRun, includeDeleted: $includeDeleted) { type data + url + } + } +""" + +CURSOR_QUERY_V1 = """ + query J1QL_v2($query: String!, $variables: JSON, $flags: QueryV1Flags, $includeDeleted: Boolean, $cursor: String) { + queryV1( + query: $query + variables: $variables + flags: $flags + includeDeleted: $includeDeleted + cursor: $cursor + ) { + type + data + cursor + __typename } } """ diff --git a/tests/test_query.py b/tests/test_query.py index 9228d7f..f9f5860 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -1,60 +1,146 @@ import json import pytest import responses +from collections import Counter from jupiterone.client import JupiterOneClient -from jupiterone.constants import QUERY_V1 +from jupiterone.constants import QUERY_V1, DEFERRED_RESULTS_COMPLETED +from jupiterone.errors import JupiterOneApiError -def request_callback(request): - headers = { - 'Content-Type': 'application/json' - } +API_ENDPOINT = 'https://api.us.jupiterone.io/graphql' +STATE_FILE_ENDPOINT = 'https://api.us.jupiterone.io/state' +RESULTS_ENDPOINT = 'https://api.us.jupiterone.io/results' - response = { - 'data': { - 'queryV1': { - 'type': 'list', - 'data': [ - { - 'id': '1', - 'entity': { - '_rawDataHashes': '1', - '_integrationDefinitionId': '1', - '_integrationName': '1', - '_beginOn': 1580482083079, - 'displayName': 'host1', - '_class': ['Host'], - '_scope': 'aws_instance', - '_version': 1, - '_integrationClass': 'CSP', - '_accountId': 'testAccount', - '_id': '1', - '_key': 'key1', - '_type': ['aws_instance'], - '_deleted': False, - '_integrationInstanceId': '1', - '_integrationType': 'aws', - '_source': 'integration-managed', - '_createdOn': 1578093840019 - }, - 'properties': { - 'id': 'host1', - 'active': True - } - } - ] +def build_deferred_query_response(response_code: int = 200, url: str = STATE_FILE_ENDPOINT): + def request_callback(request): + headers = { + 'Content-Type': 'application/json' + } + + response = { + 'data': { + 'queryV1': { + 'url': url + } } } - } - return (200, headers, json.dumps(response)) + return (response_code, headers, json.dumps(response)) + return request_callback +def build_state_response(response_code: int = 200, status: str = DEFERRED_RESULTS_COMPLETED, url: str = RESULTS_ENDPOINT): + def request_callback(request): + headers = { + 'Content-Type': 'application/json' + } + + response = { + 'status': status, + 'url': url + } + return (response_code, headers, json.dumps(response)) + return request_callback + +def build_deferred_query_results(response_code: int = 200, cursor: str = None, max_pages: int = 1): + pages = Counter(requests=0) + + def request_callback(request): + headers = { + 'Content-Type': 'application/json' + } + + response = { + 'data': [ + { + 'id': '1', + 'entity': { + '_rawDataHashes': '1', + '_integrationDefinitionId': '1', + '_integrationName': '1', + '_beginOn': 1580482083079, + 'displayName': 'host1', + '_class': ['Host'], + '_scope': 'aws_instance', + '_version': 1, + '_integrationClass': 'CSP', + '_accountId': 'testAccount', + '_id': '1', + '_key': 'key1', + '_type': ['aws_instance'], + '_deleted': False, + '_integrationInstanceId': '1', + '_integrationType': 'aws', + '_source': 'integration-managed', + '_createdOn': 1578093840019 + }, + 'properties': { + 'id': 'host1', + 'active': True + } + } + ] + } + + if cursor is not None and pages.get('requests') < max_pages: + response['cursor'] = cursor + + pages.update(requests=1) + + return (response_code, headers, json.dumps(response)) + + return request_callback + +def build_results(response_code: int = 200): + def request_callback(request): + headers = { + 'Content-Type': 'application/json' + } + + response = { + 'data': { + 'queryV1': { + 'type': 'list', + 'data': [ + { + 'id': '1', + 'entity': { + '_rawDataHashes': '1', + '_integrationDefinitionId': '1', + '_integrationName': '1', + '_beginOn': 1580482083079, + 'displayName': 'host1', + '_class': ['Host'], + '_scope': 'aws_instance', + '_version': 1, + '_integrationClass': 'CSP', + '_accountId': 'testAccount', + '_id': '1', + '_key': 'key1', + '_type': ['aws_instance'], + '_deleted': False, + '_integrationInstanceId': '1', + '_integrationType': 'aws', + '_source': 'integration-managed', + '_createdOn': 1578093840019 + }, + 'properties': { + 'id': 'host1', + 'active': True + } + } + ] + } + } + } + return (response_code, headers, json.dumps(response)) + return request_callback + @responses.activate def test_execute_query(): responses.add_callback( responses.POST, 'https://api.us.jupiterone.io/graphql', - callback=request_callback, + callback=build_results(), content_type='application/json', ) @@ -73,30 +159,66 @@ def test_execute_query(): assert 'queryV1' in response['data'] assert len(response['data']['queryV1']['data']) == 1 assert type(response['data']['queryV1']['data']) == list - assert response['data']['queryV1']['data'][0]['entity']['_id'] == '1' + assert response['data']['queryV1']['data'][0]['entity']['_id'] == '1' @responses.activate -def test_query_v1(): +def test_limit_skip_query_v1(): responses.add_callback( responses.POST, 'https://api.us.jupiterone.io/graphql', - callback=request_callback, + callback=build_results(), content_type='application/json', ) j1 = JupiterOneClient(account='testAccount', token='testToken') query = "find Host with _id='1'" - response = j1.query_v1(query) + response = j1.query_v1( + query=query, + limit=250, + skip=0 + ) assert type(response) == dict assert len(response['data']) == 1 assert type(response['data']) == list assert response['data'][0]['entity']['_id'] == '1' +@responses.activate +def test_cursor_query_v1(): + + responses.add_callback( + responses.POST, API_ENDPOINT, + callback=build_deferred_query_response(url=STATE_FILE_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, STATE_FILE_ENDPOINT, + callback=build_state_response(url=RESULTS_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, RESULTS_ENDPOINT, + callback=build_deferred_query_results(cursor='cursor_value'), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1'" + + response = j1.query_v1( + query=query, + ) + + assert type(response) == dict + assert len(response['data']) == 2 + assert type(response['data']) == list + assert response['data'][0]['entity']['_id'] == '1' @responses.activate -def test_tree_query_v1(): +def test_limit_skip_tree_query_v1(): def request_callback(request): headers = { @@ -131,12 +253,203 @@ def request_callback(request): j1 = JupiterOneClient(account='testAccount', token='testToken') query = "find Host with _id='1' return tree" - response = j1.query_v1(query) + response = j1.query_v1( + query=query, + limit=250, + skip=0 + ) - assert type(response) == dict + assert type(response) == dict assert 'edges' in response assert 'vertices' in response assert type(response['edges']) == list assert type(response['vertices']) == list assert response['vertices'][0]['id'] == '1' - \ No newline at end of file + +@responses.activate +def test_cursor_tree_query_v1(): + + def request_callback(request): + headers = { + 'Content-Type': 'application/json' + } + + response = { + 'data': { + 'vertices': [ + { + 'id': '1', + 'entity': {}, + 'properties': {} + } + ], + 'edges': [] + } + } + + return (200, headers, json.dumps(response)) + + responses.add_callback( + responses.POST, API_ENDPOINT, + callback=build_deferred_query_response(url=STATE_FILE_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, STATE_FILE_ENDPOINT, + callback=build_state_response(url=RESULTS_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, RESULTS_ENDPOINT, + callback=request_callback, + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1' return tree" + response = j1.query_v1( + query=query + ) + + assert type(response) == dict + assert 'edges' in response + assert 'vertices' in response + assert type(response['edges']) == list + assert type(response['vertices']) == list + assert response['vertices'][0]['id'] == '1' + +@responses.activate +def test_retry_on_limit_skip_query(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(response_code=429), + content_type='application/json', + ) + + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(response_code=503), + content_type='application/json', + ) + + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1'" + response = j1.query_v1( + query=query, + limit=250, + skip=0 + ) + + assert type(response) == dict + assert len(response['data']) == 1 + assert type(response['data']) == list + assert response['data'][0]['entity']['_id'] == '1' + +@responses.activate +def test_retry_on_cursor_query(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(response_code=429), + content_type='application/json', + ) + + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(response_code=503), + content_type='application/json', + ) + + responses.add_callback( + responses.POST, API_ENDPOINT, + callback=build_deferred_query_response(url=STATE_FILE_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, STATE_FILE_ENDPOINT, + callback=build_state_response(response_code=503, url=RESULTS_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, STATE_FILE_ENDPOINT, + callback=build_state_response(url=RESULTS_ENDPOINT), + content_type='application/json', + ) + + responses.add_callback( + responses.GET, RESULTS_ENDPOINT, + callback=build_deferred_query_results(), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1'" + response = j1.query_v1( + query=query + ) + + assert type(response) == dict + assert len(response['data']) == 1 + assert type(response['data']) == list + assert response['data'][0]['entity']['_id'] == '1' + +@responses.activate +def test_avoid_retry_on_limit_skip_query(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(response_code=404), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1'" + with pytest.raises(JupiterOneApiError): + j1.query_v1( + query=query, + limit=250, + skip=0 + ) + +@responses.activate +def test_avoid_retry_on_cursor_query(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(response_code=404), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1'" + with pytest.raises(JupiterOneApiError): + j1.query_v1( + query=query, + limit=250, + skip=0 + ) + +@responses.activate +def test_warn_limit_and_skip_deprecated(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='testToken') + query = "find Host with _id='1'" + + with pytest.warns(DeprecationWarning): + j1.query_v1( + query=query, + limit=250, + skip=0 + ) From af23ad9e5a614beb1a6c24f71e2728837e3cfd6f Mon Sep 17 00:00:00 2001 From: Tyler Walch Date: Thu, 18 Nov 2021 15:56:50 -0500 Subject: [PATCH 2/9] benchmarked to be faster: "FIND DataStore" 0.62s user 0.14s system 1% cpu 41.754 total 0.56s user 0.10s system 2% cpu 32.367 total "FIND (Service|DataStore)" 0.67s user 0.13s system 1% cpu 55.249 total 0.67s user 0.12s system 1% cpu 39.487 total --- jupiterone/client.py | 5 ++--- jupiterone/constants.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index d814acf..2c362bc 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -133,15 +133,14 @@ def _cursor_query(self, query: str, cursor: str = None, include_deleted: bool = variables['cursor'] = cursor response = self._execute_query(query=CURSOR_QUERY_V1, variables=variables) - data = response['data']['queryV1'] + data = response['data']['queryV1']['data'] if 'vertices' in data and 'edges' in data: return data results.extend(data) - if 'cursor' in response['data']['queryV1']: - print('cursor', response['data']['queryV1']['cursor'], len(data), len(results)) + if 'cursor' in response['data']['queryV1'] and response['data']['queryV1']['cursor'] is not None: cursor = response['data']['queryV1']['cursor'] else: break diff --git a/jupiterone/constants.py b/jupiterone/constants.py index d8d350c..1383e97 100644 --- a/jupiterone/constants.py +++ b/jupiterone/constants.py @@ -6,7 +6,6 @@ queryV1(query: $query, variables: $variables, dryRun: $dryRun, includeDeleted: $includeDeleted) { type data - url } } """ @@ -16,6 +15,7 @@ queryV1( query: $query variables: $variables + deferredResponse: DISABLED flags: $flags includeDeleted: $includeDeleted cursor: $cursor From b5e715e91a6cfc38c14badd3bb80a6898c040517 Mon Sep 17 00:00:00 2001 From: Tyler Walch Date: Thu, 18 Nov 2021 16:04:24 -0500 Subject: [PATCH 3/9] Fixing tests --- tests/test_query.py | 165 ++++++++------------------------------------ 1 file changed, 29 insertions(+), 136 deletions(-) diff --git a/tests/test_query.py b/tests/test_query.py index f9f5860..a524f98 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -4,93 +4,12 @@ from collections import Counter from jupiterone.client import JupiterOneClient -from jupiterone.constants import QUERY_V1, DEFERRED_RESULTS_COMPLETED +from jupiterone.constants import QUERY_V1 from jupiterone.errors import JupiterOneApiError -API_ENDPOINT = 'https://api.us.jupiterone.io/graphql' -STATE_FILE_ENDPOINT = 'https://api.us.jupiterone.io/state' -RESULTS_ENDPOINT = 'https://api.us.jupiterone.io/results' - -def build_deferred_query_response(response_code: int = 200, url: str = STATE_FILE_ENDPOINT): - def request_callback(request): - headers = { - 'Content-Type': 'application/json' - } - - response = { - 'data': { - 'queryV1': { - 'url': url - } - } - } - return (response_code, headers, json.dumps(response)) - return request_callback - - -def build_state_response(response_code: int = 200, status: str = DEFERRED_RESULTS_COMPLETED, url: str = RESULTS_ENDPOINT): - def request_callback(request): - headers = { - 'Content-Type': 'application/json' - } - - response = { - 'status': status, - 'url': url - } - return (response_code, headers, json.dumps(response)) - return request_callback - -def build_deferred_query_results(response_code: int = 200, cursor: str = None, max_pages: int = 1): +def build_results(response_code: int = 200, cursor: str = None, max_pages: int = 1): pages = Counter(requests=0) - def request_callback(request): - headers = { - 'Content-Type': 'application/json' - } - - response = { - 'data': [ - { - 'id': '1', - 'entity': { - '_rawDataHashes': '1', - '_integrationDefinitionId': '1', - '_integrationName': '1', - '_beginOn': 1580482083079, - 'displayName': 'host1', - '_class': ['Host'], - '_scope': 'aws_instance', - '_version': 1, - '_integrationClass': 'CSP', - '_accountId': 'testAccount', - '_id': '1', - '_key': 'key1', - '_type': ['aws_instance'], - '_deleted': False, - '_integrationInstanceId': '1', - '_integrationType': 'aws', - '_source': 'integration-managed', - '_createdOn': 1578093840019 - }, - 'properties': { - 'id': 'host1', - 'active': True - } - } - ] - } - - if cursor is not None and pages.get('requests') < max_pages: - response['cursor'] = cursor - - pages.update(requests=1) - - return (response_code, headers, json.dumps(response)) - - return request_callback - -def build_results(response_code: int = 200): def request_callback(request): headers = { 'Content-Type': 'application/json' @@ -132,7 +51,14 @@ def request_callback(request): } } } + + if cursor is not None and pages.get('requests') < max_pages: + response['data']['queryV1']['cursor'] = cursor + + pages.update(requests=1) + return (response_code, headers, json.dumps(response)) + return request_callback @responses.activate @@ -188,20 +114,14 @@ def test_limit_skip_query_v1(): def test_cursor_query_v1(): responses.add_callback( - responses.POST, API_ENDPOINT, - callback=build_deferred_query_response(url=STATE_FILE_ENDPOINT), - content_type='application/json', - ) - - responses.add_callback( - responses.GET, STATE_FILE_ENDPOINT, - callback=build_state_response(url=RESULTS_ENDPOINT), + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(cursor='cursor_value'), content_type='application/json', ) responses.add_callback( - responses.GET, RESULTS_ENDPOINT, - callback=build_deferred_query_results(cursor='cursor_value'), + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(), content_type='application/json', ) @@ -276,33 +196,26 @@ def request_callback(request): response = { 'data': { - 'vertices': [ - { - 'id': '1', - 'entity': {}, - 'properties': {} + 'queryV1': { + 'type': 'tree', + 'data': { + 'vertices': [ + { + 'id': '1', + 'entity': {}, + 'properties': {} + } + ], + 'edges': [] } - ], - 'edges': [] + } } } return (200, headers, json.dumps(response)) responses.add_callback( - responses.POST, API_ENDPOINT, - callback=build_deferred_query_response(url=STATE_FILE_ENDPOINT), - content_type='application/json', - ) - - responses.add_callback( - responses.GET, STATE_FILE_ENDPOINT, - callback=build_state_response(url=RESULTS_ENDPOINT), - content_type='application/json', - ) - - responses.add_callback( - responses.GET, RESULTS_ENDPOINT, + responses.POST, 'https://api.us.jupiterone.io/graphql', callback=request_callback, content_type='application/json', ) @@ -368,26 +281,8 @@ def test_retry_on_cursor_query(): ) responses.add_callback( - responses.POST, API_ENDPOINT, - callback=build_deferred_query_response(url=STATE_FILE_ENDPOINT), - content_type='application/json', - ) - - responses.add_callback( - responses.GET, STATE_FILE_ENDPOINT, - callback=build_state_response(response_code=503, url=RESULTS_ENDPOINT), - content_type='application/json', - ) - - responses.add_callback( - responses.GET, STATE_FILE_ENDPOINT, - callback=build_state_response(url=RESULTS_ENDPOINT), - content_type='application/json', - ) - - responses.add_callback( - responses.GET, RESULTS_ENDPOINT, - callback=build_deferred_query_results(), + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_results(), content_type='application/json', ) @@ -431,9 +326,7 @@ def test_avoid_retry_on_cursor_query(): query = "find Host with _id='1'" with pytest.raises(JupiterOneApiError): j1.query_v1( - query=query, - limit=250, - skip=0 + query=query ) @responses.activate From f2e1dcea379578b969948aeb11dff8ee82b67a1c Mon Sep 17 00:00:00 2001 From: Tyler Walch Date: Thu, 18 Nov 2021 16:11:33 -0500 Subject: [PATCH 4/9] Clean up --- tests/test_query.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_query.py b/tests/test_query.py index a524f98..9c7859e 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -222,9 +222,7 @@ def request_callback(request): j1 = JupiterOneClient(account='testAccount', token='testToken') query = "find Host with _id='1' return tree" - response = j1.query_v1( - query=query - ) + response = j1.query_v1(query) assert type(response) == dict assert 'edges' in response From b2b178f1d578b53472cb5b815aa10abdae4f2023 Mon Sep 17 00:00:00 2001 From: Tyler Walch Date: Thu, 18 Nov 2021 16:16:08 -0500 Subject: [PATCH 5/9] Bumping version --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index eb477c7..17ba317 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ ] setup(name='jupiterone', - version='0.1.0', + version='0.2.0', description='A Python client for the JupiterOne API', license='MIT License', author='George Vauter', From b5165965f21ce8969232a59d5f0c58cbe92da3ab Mon Sep 17 00:00:00 2001 From: lmarkscp Date: Fri, 15 Jul 2022 14:21:36 -0400 Subject: [PATCH 6/9] Added handling of 401 and other responses that aren't JSON. Added tests for the new error handling. --- jupiterone/client.py | 11 +++++-- tests/test_query.py | 69 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index 2c362bc..c3a0786 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -107,12 +107,19 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: raise JupiterOneApiError(content.get('errors')) return response.json() + elif response.status_code == 401: + raise JupiterOneApiError('JupiterOne API query is unauthorized, check credentials.') + elif response.status_code in [429, 503]: raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded') else: - content = json.loads(response._content) - raise JupiterOneApiError('{}:{}'.format(response.status_code, content.get('error'))) + content = response._content + if isinstance(content, (bytes, bytearray)): + content = content.decode("utf-8") + if 'application/json' in response.headers.get('Content-Type', 'application/json') and '"error":' in content: + content = json.loads(content).get('error') + raise JupiterOneApiError('{}:{}'.format(response.status_code, content)) def _cursor_query(self, query: str, cursor: str = None, include_deleted: bool = False) -> Dict: """ Performs a V1 graph query using cursor pagination diff --git a/tests/test_query.py b/tests/test_query.py index 9c7859e..5172c17 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -61,6 +61,17 @@ def request_callback(request): return request_callback + +def build_error_results(response_code: int, response_content, response_type: str = 'application/json'): + def request_callback(request): + headers = { + 'Content-Type': response_type + } + return response_code, headers, response_content + + return request_callback + + @responses.activate def test_execute_query(): @@ -344,3 +355,61 @@ def test_warn_limit_and_skip_deprecated(): limit=250, skip=0 ) + + +@responses.activate +def test_unauthorized_query_v1(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_error_results(401, b'Unauthorized', 'text/plain'), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='bogusToken') + query = "find Host with _id='1' return tree" + + with pytest.raises(JupiterOneApiError) as exc_info: + j1.query_v1(query) + + assert type(exc_info.value) == JupiterOneApiError + assert exc_info.value.args[0] == 'JupiterOne API query is unauthorized, check credentials.' + + +@responses.activate +def test_unexpected_string_error_query_v1(): + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_error_results(500, 'String exception on server', 'text/plain'), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='bogusToken') + query = "find Host with _id='1' return tree" + + with pytest.raises(JupiterOneApiError) as exc_info: + j1.query_v1(query) + + assert type(exc_info.value) == JupiterOneApiError + assert exc_info.value.args[0] == '500:String exception on server' + + +@responses.activate +def test_unexpected_json_error_query_v1(): + error_json = { + 'error': 'Bad Gateway' + } + + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_error_results(502, json.dumps(error_json), ), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='bogusToken') + query = "find Host with _id='1' return tree" + + with pytest.raises(JupiterOneApiError) as exc_info: + j1.query_v1(query) + + assert type(exc_info.value) == JupiterOneApiError + assert exc_info.value.args[0] == '502:Bad Gateway' From 0d3a5aed0a8fccb1b1f6fe4490041872db04fcc9 Mon Sep 17 00:00:00 2001 From: lmarkscp Date: Mon, 18 Jul 2022 13:42:06 -0400 Subject: [PATCH 7/9] Cleaned up handling of error(s) in json --- jupiterone/client.py | 5 +++-- tests/test_query.py | 24 +++++++++++++++++++++--- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index c3a0786..f6287d1 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -117,8 +117,9 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: content = response._content if isinstance(content, (bytes, bytearray)): content = content.decode("utf-8") - if 'application/json' in response.headers.get('Content-Type', 'application/json') and '"error":' in content: - content = json.loads(content).get('error') + if 'application/json' in response.headers.get('Content-Type', 'text/plain'): + data = json.loads(content) + content = data.get('error', data.get('errors', content)) raise JupiterOneApiError('{}:{}'.format(response.status_code, content)) def _cursor_query(self, query: str, cursor: str = None, include_deleted: bool = False) -> Dict: diff --git a/tests/test_query.py b/tests/test_query.py index 5172c17..28dbc58 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -371,7 +371,6 @@ def test_unauthorized_query_v1(): with pytest.raises(JupiterOneApiError) as exc_info: j1.query_v1(query) - assert type(exc_info.value) == JupiterOneApiError assert exc_info.value.args[0] == 'JupiterOne API query is unauthorized, check credentials.' @@ -389,7 +388,6 @@ def test_unexpected_string_error_query_v1(): with pytest.raises(JupiterOneApiError) as exc_info: j1.query_v1(query) - assert type(exc_info.value) == JupiterOneApiError assert exc_info.value.args[0] == '500:String exception on server' @@ -411,5 +409,25 @@ def test_unexpected_json_error_query_v1(): with pytest.raises(JupiterOneApiError) as exc_info: j1.query_v1(query) - assert type(exc_info.value) == JupiterOneApiError assert exc_info.value.args[0] == '502:Bad Gateway' + + +@responses.activate +def test_unexpected_json_errors_query_v1(): + errors_json = { + 'errors': ['First error', 'Second error', 'Third error'] + } + + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + callback=build_error_results(500, json.dumps(errors_json), ), + content_type='application/json', + ) + + j1 = JupiterOneClient(account='testAccount', token='bogusToken') + query = "find Host with _id='1' return tree" + + with pytest.raises(JupiterOneApiError) as exc_info: + j1.query_v1(query) + + assert exc_info.value.args[0] == "500:['First error', 'Second error', 'Third error']" From f6a9d56fb402470f93fedcefc1dcac7a87f59f65 Mon Sep 17 00:00:00 2001 From: "sre-57-opslevel[bot]" <113727212+sre-57-opslevel[bot]@users.noreply.github.com> Date: Mon, 17 Oct 2022 13:09:37 +0000 Subject: [PATCH 8/9] Upload OpsLevel YAML --- opslevel.yml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 opslevel.yml diff --git a/opslevel.yml b/opslevel.yml new file mode 100644 index 0000000..3a5e183 --- /dev/null +++ b/opslevel.yml @@ -0,0 +1,6 @@ +--- +version: 1 +repository: + owner: cloud_security_ops + tier: + tags: From 7b76a4609c9d89f93b3d8a6f5b6acfe66276a321 Mon Sep 17 00:00:00 2001 From: "sara.perez" Date: Tue, 22 Aug 2023 15:37:10 +0200 Subject: [PATCH 9/9] bumping version --- setup.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 17ba317..46c9dfb 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020 Auth0 +# Copyright (c) 2020-2023 Okta from setuptools import setup, find_packages @@ -8,12 +8,12 @@ ] setup(name='jupiterone', - version='0.2.0', + version='0.2.1', description='A Python client for the JupiterOne API', license='MIT License', author='George Vauter', - author_email='george.vauter@auth0.com', - maintainer='Auth0', + author_email='george.vauter@okta.com', + maintainer='Okta', url='https://github.com/auth0/jupiterone-python-sdk', install_requires=install_reqs, classifiers=[