diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index 19bcdce..0000000 --- a/.travis.yml +++ /dev/null @@ -1,12 +0,0 @@ -language: python -python: - - 3.6 - - 3.7 - -install: - - pip install . - - pip install responses - - pip install tox-travis - -script: - - tox diff --git a/LICENSE b/LICENSE index d2618d4..d9da973 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2020 Auth0 +Copyright (c) 2024 JupiterOne Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index cca35b9..8075f04 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![Python 3.7](https://img.shields.io/badge/python-3.7-blue.svg)](https://www.python.org/downloads/release/python-370/) -A Python library for the [JupiterOne API](https://support.jupiterone.io/hc/en-us/articles/360022722094-JupiterOne-Platform-API). +A Python library for the [JupiterOne API](https://docs.jupiterone.io/reference). ## Installation @@ -41,6 +41,10 @@ query_result = j1.query_v1(QUERY, include_deleted=True) # Tree query QUERY = 'FIND Host RETURN TREE' query_result = j1.query_v1(QUERY) + +# Using cursor graphQL variable to return full set of paginated results +QUERY = "FIND (Device | Person)" +cursor_query_r = j1._cursor_query(QUERY) ``` ##### Create an entity: diff --git a/examples/examples.py b/examples/examples.py new file mode 100644 index 0000000..402bdef --- /dev/null +++ b/examples/examples.py @@ -0,0 +1,91 @@ +from jupiterone.client import JupiterOneClient +import random +import time +import os + +account = os.environ.get("JUPITERONE_ACCOUNT") +token = os.environ.get("JUPITERONE_TOKEN") + +j1 = JupiterOneClient(account=account, token=token) + +# query_v1 +q = "FIND jupiterone_user" +query_r = j1.query_v1(q) +print(query_r) + +# create_entity +num1 = random.randrange(1, 999, 1) + +# create_entity +properties = { + 'displayName': 'test{}'.format(num1), + 'customProperty': 'customVal', + 'tag.Production': 'false', + 'owner': 'user.name@jupiterone.com' + } + +create_r = j1.create_entity( + entity_key='jupiterone-api-client-python:{}'.format(num1), + entity_type='python_client_create_entity', + entity_class='Record', + properties=properties, + timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime +) +print(create_r) + +properties = { + 'customProperty': 'customValUpdated' + } + +# update_entity +update_r = j1.update_entity( + entity_id='{}'.format(create_r['entity']['_id']), + properties=properties +) +print(update_r) + +# create_entity_2 +num2 = random.randrange(1, 999, 1) + +properties = { + 'displayName': 'test{}'.format(num2), + 'customProperty': 'customVal', + 'tag.Production': 'false', + 'owner': 'user.name@jupiterone.com' + } + +create_r_2 = j1.create_entity( + entity_key='jupiterone-api-client-python:{}'.format(num2), + entity_type='python_client_create_entity', + entity_class='Record', + properties=properties, + timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime +) +print(create_r_2) + +# create_relationship +create_relationship_r = j1.create_relationship( + relationship_key='{}:{}'.format(create_r['entity']['_id'], create_r_2['entity']['_id']), + relationship_type='jupiterone-api-client-python:create_relationship', + relationship_class='HAS', + from_entity_id=create_r['entity']['_id'], + to_entity_id=create_r_2['entity']['_id'] +) +print(create_relationship_r) + +# delete_relationship +delete_relationship_r = j1.delete_relationship(relationship_id=create_relationship_r['relationship']['_id']) +print(delete_relationship_r) + +# delete_entity +delete_entity_r1 = j1.delete_entity(entity_id=create_r['entity']['_id']) +print(delete_entity_r1) + +delete_entity_r2 = j1.delete_entity(entity_id=create_r_2['entity']['_id']) +print(delete_entity_r2) + + +q = "FIND (Device | Person)" +cursor_query_r = j1._cursor_query(q) +print(cursor_query_r) +print(len(cursor_query_r['data'])) diff --git a/jupiterone/__init__.py b/jupiterone/__init__.py index 50b1190..0d227bb 100644 --- a/jupiterone/__init__.py +++ b/jupiterone/__init__.py @@ -2,4 +2,4 @@ from .errors import ( JupiterOneClientError, JupiterOneApiError -) \ No newline at end of file +) diff --git a/jupiterone/client.py b/jupiterone/client.py index 422d24e..542e626 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,9 +23,11 @@ DELETE_ENTITY, UPDATE_ENTITY, CREATE_RELATIONSHIP, - DELETE_RELATIONSHIP + DELETE_RELATIONSHIP, + CURSOR_QUERY_V1 ) + def retry_on_429(exc): """ Used to trigger retry on rate limit """ return isinstance(exc, JupiterOneApiRetryError) @@ -68,12 +71,12 @@ def account(self, value: str): @property def token(self): - """ Your JupiteOne access token """ + """ Your JupiterOne access token """ return self._token @token.setter def token(self, value: str): - """ Your JupiteOne access token """ + """ Your JupiterOne access token """ if not value: raise JupiterOneClientError('token is required') self._token = value @@ -89,11 +92,11 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: if variables: data.update(variables=variables) - response = requests.post(self.query_endpoint, headers=self.headers, json=data) + response = requests.post(self.query_endpoint, headers=self.headers, json=data, timeout=60) # It is still unclear if all responses will have a status # code of 200 or if 429 will eventually be used to - # indicate rate limitting. J1 devs are aware. + # indicate rate limits being hit. J1 devs are aware. if response.status_code == 200: if response._content: content = json.loads(response._content) @@ -108,29 +111,59 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: elif response.status_code == 401: raise JupiterOneApiError('401: Unauthorized. Please supply a valid account id and API token.') - elif response.status_code in [429, 500]: + elif response.status_code in [429, 503]: raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded') - else: - try: - content = json.loads(response._content) - raise JupiterOneApiError('{}: {}'.format(response.status_code, content.get('error') or 'Unknown Error')) - except ValueError as e: - raise JupiterOneApiError('{}: {}'.format(response.status_code, 'Unknown Error')); - + elif response.status_code in [500]: + raise JupiterOneApiError('JupiterOne API internal server error.') - def query_v1(self, query: str, **kwargs) -> Dict: - """ Performs a V1 graph query + else: + 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 @@ -159,6 +192,39 @@ 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://docs.jupiterone.io/features/admin/parameters#query-parameterlist', + 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. @@ -206,7 +272,7 @@ def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: Update an existing entity. args: - entity_id (str): The _id of the entity to udate + entity_id (str): The _id of the entity to update properties (dict): Dictionary of key/value entity properties """ variables = { @@ -218,7 +284,7 @@ def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: def create_relationship(self, **kwargs) -> Dict: """ - Create a relationship (edge) between two entities (veritces). + Create a relationship (edge) between two entities (vertices). args: relationship_key (str): Unique key for the relationship 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/jupiterone/errors.py b/jupiterone/errors.py index aaf1b3e..eccf017 100644 --- a/jupiterone/errors.py +++ b/jupiterone/errors.py @@ -2,8 +2,10 @@ class JupiterOneClientError(Exception): """ Raised when error creating client """ + class JupiterOneApiRetryError(Exception): """ Used to trigger retry on rate limit """ + class JupiterOneApiError(Exception): - """ Raised when API returns error response """ \ No newline at end of file + """ Raised when API returns error response """ diff --git a/requirements.txt b/requirements.txt index d5f6187..3c9c733 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ retrying requests +warnings diff --git a/setup.py b/setup.py index eb477c7..9cb3f2a 100644 --- a/setup.py +++ b/setup.py @@ -1,20 +1,20 @@ -# Copyright (c) 2020 Auth0 - +# Copyright (c) 2020-2025 JupiterOne from setuptools import setup, find_packages install_reqs = [ 'requests', - 'retrying' + 'retrying', + 'warnings' ] setup(name='jupiterone', - version='0.1.0', + version='1.0.0', description='A Python client for the JupiterOne API', license='MIT License', - author='George Vauter', - author_email='george.vauter@auth0.com', - maintainer='Auth0', - url='https://github.com/auth0/jupiterone-python-sdk', + author='JupiterOne', + author_email='solutions@jupiterone.com', + maintainer='JupiterOne', + url='https://github.com/JupiterOne/jupiterone-api-client-python', install_requires=install_reqs, classifiers=[ 'Development Status :: 4 - Beta', @@ -28,4 +28,4 @@ 'Topic :: Security', ], packages=find_packages() -) + ) diff --git a/tests/test_query.py b/tests/test_query.py index 708d573..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,57 +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_401(): - with pytest.raises(Exception) as ex: +def test_limit_skip_tree_query_v1(): - def response_401_callback(r): - headers = { - 'Content-Type': 'application/json' - } + def request_callback(request): + headers = { + 'Content-Type': 'application/json' + } - response = { - 'test': ['Unauthorized'] + response = { + 'data': { + 'queryV1': { + 'type': 'tree', + 'data': { + 'vertices': [ + { + 'id': '1', + 'entity': {}, + 'properties': {} + } + ], + 'edges': [] + } + } } - return (401, headers, json.dumps(response)) + } - responses.add_callback( - responses.POST, 'https://api.us.jupiterone.io/graphql', - callback=response_401_callback, - content_type='application/text', - ) + return (200, headers, json.dumps(response)) - j1 = JupiterOneClient(account='testAccount', token='testToken') - query = "find Host with _id='1'" - j1.query_v1(query) - - assert '401: Unauthorized. Please supply a valid account id and API token.' in str(ex.value) + 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_tree_query_v1(): +def test_cursor_tree_query_v1(): def request_callback(request): headers = { @@ -160,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']"