diff --git a/jupiterone/client.py b/jupiterone/client.py index a9be725..f6287d1 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,55 @@ 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 == 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'))) - - def query_v1(self, query: str, **kwargs) -> Dict: - """ Performs a V1 graph query + content = response._content + if isinstance(content, (bytes, bytearray)): + content = content.decode("utf-8") + 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: + """ 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']['data'] + + if 'vertices' in data and 'edges' in data: + return data + + results.extend(data) + + if 'cursor' in response['data']['queryV1'] and response['data']['queryV1']['cursor'] is not None: + 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 +184,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..1383e97 100644 --- a/jupiterone/constants.py +++ b/jupiterone/constants.py @@ -10,6 +10,24 @@ } """ +CURSOR_QUERY_V1 = """ + query J1QL_v2($query: String!, $variables: JSON, $flags: QueryV1Flags, $includeDeleted: Boolean, $cursor: String) { + queryV1( + query: $query + variables: $variables + deferredResponse: DISABLED + flags: $flags + includeDeleted: $includeDeleted + cursor: $cursor + ) { + type + data + cursor + __typename + } + } +""" + CREATE_ENTITY = """ mutation CreateEntity( $entityKey: String! 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: diff --git a/setup.py b/setup.py index eb477c7..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.1.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=[ diff --git a/tests/test_query.py b/tests/test_query.py index 9228d7f..28dbc58 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -1,52 +1,75 @@ import json import pytest import responses +from collections import Counter from jupiterone.client import JupiterOneClient from jupiterone.constants import QUERY_V1 +from jupiterone.errors import JupiterOneApiError -def request_callback(request): - headers = { - 'Content-Type': 'application/json' - } +def build_results(response_code: int = 200, cursor: str = None, max_pages: int = 1): + pages = Counter(requests=0) - 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 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 (200, headers, json.dumps(response)) + + 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 + + +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 @@ -54,7 +77,7 @@ 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 +96,109 @@ 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, 'https://api.us.jupiterone.io/graphql', + callback=build_results(cursor='cursor_value'), + 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, + ) + + 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 = { + 'Content-Type': 'application/json' + } + + response = { + 'data': { + 'queryV1': { + 'type': 'tree', + 'data': { + 'vertices': [ + { + 'id': '1', + 'entity': {}, + 'properties': {} + } + ], + 'edges': [] + } + } + } + } + + return (200, headers, json.dumps(response)) + + responses.add_callback( + responses.POST, 'https://api.us.jupiterone.io/graphql', + 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, + limit=250, + skip=0 + ) + + 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_cursor_tree_query_v1(): def request_callback(request): headers = { @@ -133,10 +235,199 @@ def request_callback(request): query = "find Host with _id='1' return tree" response = j1.query_v1(query) - 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_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, '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 + ) + + 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 + ) + +@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 + ) + + +@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 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 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 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']"