├── requirements.txt ├── dev-requirements.txt ├── opslevel.yml ├── MANIFEST.in ├── jupiterone ├── __init__.py ├── errors.py ├── constants.py └── client.py ├── .travis.yml ├── tests ├── test_client.py ├── test_update_entity.py ├── test_delete_entity.py ├── test_create_entity.py ├── test_delete_relationship.py ├── test_create_relationship.py └── test_query.py ├── setup.py ├── LICENSE ├── .gitignore └── README.md /requirements.txt: -------------------------------------------------------------------------------- 1 | retrying 2 | requests 3 | -------------------------------------------------------------------------------- /dev-requirements.txt: -------------------------------------------------------------------------------- 1 | -r requirements.txt 2 | pytest 3 | responses -------------------------------------------------------------------------------- /opslevel.yml: -------------------------------------------------------------------------------- 1 | --- 2 | version: 1 3 | repository: 4 | owner: cloud_security_ops 5 | tier: 6 | tags: 7 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include MANIFEST.in 2 | include LICENSE 3 | include README.md 4 | recursive-include *.py docs *.rst -------------------------------------------------------------------------------- /jupiterone/__init__.py: -------------------------------------------------------------------------------- 1 | from .client import JupiterOneClient 2 | from .errors import ( 3 | JupiterOneClientError, 4 | JupiterOneApiError 5 | ) -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - 3.6 4 | - 3.7 5 | 6 | install: 7 | - pip install . 8 | - pip install responses 9 | - pip install tox-travis 10 | 11 | script: 12 | - tox 13 | -------------------------------------------------------------------------------- /jupiterone/errors.py: -------------------------------------------------------------------------------- 1 | 2 | class JupiterOneClientError(Exception): 3 | """ Raised when error creating client """ 4 | 5 | class JupiterOneApiRetryError(Exception): 6 | """ Used to trigger retry on rate limit """ 7 | 8 | class JupiterOneApiError(Exception): 9 | """ Raised when API returns error response """ -------------------------------------------------------------------------------- /tests/test_client.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from jupiterone.client import JupiterOneClient 4 | 5 | 6 | def test_missing_account(): 7 | 8 | with pytest.raises(Exception) as ex: 9 | j1 = JupiterOneClient(token='123') 10 | assert 'account is required' in str(ex.value) 11 | 12 | 13 | def test_missing_token(): 14 | 15 | with pytest.raises(Exception) as ex: 16 | j1 = JupiterOneClient(account='test') 17 | assert 'token is required' in str(ex.value) -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2020-2023 Okta 2 | 3 | from setuptools import setup, find_packages 4 | 5 | install_reqs = [ 6 | 'requests', 7 | 'retrying' 8 | ] 9 | 10 | setup(name='jupiterone', 11 | version='0.2.1', 12 | description='A Python client for the JupiterOne API', 13 | license='MIT License', 14 | author='George Vauter', 15 | author_email='george.vauter@okta.com', 16 | maintainer='Okta', 17 | url='https://github.com/auth0/jupiterone-python-sdk', 18 | install_requires=install_reqs, 19 | classifiers=[ 20 | 'Development Status :: 4 - Beta', 21 | 'Intended Audience :: Developers', 22 | 'Intended Audience :: Information Technology', 23 | 'Intended Audience :: System Administrators', 24 | 'License :: OSI Approved :: MIT License', 25 | 'Natural Language :: English', 26 | 'Operating System :: POSIX :: Linux', 27 | 'Programming Language :: Python', 28 | 'Topic :: Security', 29 | ], 30 | packages=find_packages() 31 | ) 32 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Auth0 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /tests/test_update_entity.py: -------------------------------------------------------------------------------- 1 | import json 2 | import pytest 3 | import responses 4 | 5 | from jupiterone.client import JupiterOneClient 6 | from jupiterone.constants import CREATE_ENTITY 7 | 8 | 9 | @responses.activate 10 | def test_tree_query_v1(): 11 | 12 | def request_callback(request): 13 | headers = { 14 | 'Content-Type': 'application/json' 15 | } 16 | 17 | response = { 18 | 'data': { 19 | 'updateEntity': { 20 | 'entity': { 21 | '_id': '1' 22 | }, 23 | 'vertex': { 24 | 'id': '1' 25 | } 26 | } 27 | } 28 | } 29 | return (200, headers, json.dumps(response)) 30 | 31 | responses.add_callback( 32 | responses.POST, 'https://api.us.jupiterone.io/graphql', 33 | callback=request_callback, 34 | content_type='application/json', 35 | ) 36 | 37 | j1 = JupiterOneClient(account='testAccount', token='testToken') 38 | response = j1.update_entity('1', properties={'testKey': 'testValue'}) 39 | 40 | assert type(response) == dict 41 | assert type(response['entity']) == dict 42 | assert type(response['vertex']) == dict 43 | assert response['entity']['_id'] == '1' 44 | -------------------------------------------------------------------------------- /tests/test_delete_entity.py: -------------------------------------------------------------------------------- 1 | import json 2 | import pytest 3 | import responses 4 | 5 | from jupiterone.client import JupiterOneClient 6 | from jupiterone.constants import CREATE_ENTITY 7 | 8 | 9 | @responses.activate 10 | def test_tree_query_v1(): 11 | 12 | def request_callback(request): 13 | headers = { 14 | 'Content-Type': 'application/json' 15 | } 16 | response = { 17 | 'data': { 18 | 'deleteEntity': { 19 | 'entity': { 20 | '_id': '1' 21 | }, 22 | 'vertex': { 23 | 'id': '1', 24 | 'entity': { 25 | '_id': '1' 26 | }, 27 | 'properties': {} 28 | } 29 | } 30 | } 31 | } 32 | 33 | return (200, headers, json.dumps(response)) 34 | 35 | responses.add_callback( 36 | responses.POST, 'https://api.us.jupiterone.io/graphql', 37 | callback=request_callback, 38 | content_type='application/json', 39 | ) 40 | 41 | j1 = JupiterOneClient(account='testAccount', token='testToken') 42 | response = j1.delete_entity('1') 43 | 44 | assert type(response) == dict 45 | assert type(response['entity']) == dict 46 | assert type(response['vertex']) == dict 47 | assert response['entity']['_id'] == '1' 48 | -------------------------------------------------------------------------------- /tests/test_create_entity.py: -------------------------------------------------------------------------------- 1 | import json 2 | import pytest 3 | import responses 4 | 5 | from jupiterone.client import JupiterOneClient 6 | from jupiterone.constants import CREATE_ENTITY 7 | 8 | 9 | @responses.activate 10 | def test_tree_query_v1(): 11 | 12 | def request_callback(request): 13 | headers = { 14 | 'Content-Type': 'application/json' 15 | } 16 | 17 | response = { 18 | 'data': { 19 | 'createEntity': { 20 | 'entity': { 21 | '_id': '1' 22 | }, 23 | 'vertex': { 24 | 'id': '1', 25 | 'entity': { 26 | '_id': '1' 27 | } 28 | } 29 | } 30 | } 31 | } 32 | return (200, headers, json.dumps(response)) 33 | 34 | responses.add_callback( 35 | responses.POST, 'https://api.us.jupiterone.io/graphql', 36 | callback=request_callback, 37 | content_type='application/json', 38 | ) 39 | 40 | j1 = JupiterOneClient(account='testAccount', token='testToken') 41 | response = j1.create_entity( 42 | entity_key='host1', 43 | entity_type='test_host', 44 | entity_class='Host', 45 | properties={'key1': 'value1'} 46 | ) 47 | 48 | assert type(response) == dict 49 | assert type(response['entity']) == dict 50 | assert type(response['vertex']) == dict 51 | assert response['entity']['_id'] == '1' 52 | -------------------------------------------------------------------------------- /tests/test_delete_relationship.py: -------------------------------------------------------------------------------- 1 | import json 2 | import pytest 3 | import responses 4 | 5 | from jupiterone.client import JupiterOneClient 6 | from jupiterone.constants import CREATE_ENTITY 7 | 8 | 9 | @responses.activate 10 | def test_tree_query_v1(): 11 | 12 | def request_callback(request): 13 | headers = { 14 | 'Content-Type': 'application/json' 15 | } 16 | 17 | response = { 18 | 'data': { 19 | 'deleteRelationship': { 20 | 'relationship': { 21 | '_id': '1' 22 | }, 23 | 'edge': { 24 | 'id': '1', 25 | 'toVertexId': '1', 26 | 'fromVertexId': '2', 27 | 'relationship': { 28 | '_id': '1' 29 | }, 30 | 'properties': {} 31 | } 32 | } 33 | } 34 | } 35 | 36 | return (200, headers, json.dumps(response)) 37 | 38 | responses.add_callback( 39 | responses.POST, 'https://api.us.jupiterone.io/graphql', 40 | callback=request_callback, 41 | content_type='application/json', 42 | ) 43 | 44 | j1 = JupiterOneClient(account='testAccount', token='testToken') 45 | response = j1.delete_relationship('1') 46 | 47 | assert type(response) == dict 48 | assert type(response['relationship']) == dict 49 | assert response['relationship']['_id'] == '1' 50 | assert response['edge']['toVertexId'] == '1' 51 | assert response['edge']['fromVertexId'] == '2' 52 | -------------------------------------------------------------------------------- /tests/test_create_relationship.py: -------------------------------------------------------------------------------- 1 | import json 2 | import pytest 3 | import responses 4 | 5 | from jupiterone.client import JupiterOneClient 6 | from jupiterone.constants import CREATE_ENTITY 7 | 8 | 9 | @responses.activate 10 | def test_tree_query_v1(): 11 | 12 | def request_callback(request): 13 | headers = { 14 | 'Content-Type': 'application/json' 15 | } 16 | 17 | response = { 18 | 'data': { 19 | 'createRelationship': { 20 | 'relationship': { 21 | '_id': '1' 22 | }, 23 | 'edge': { 24 | 'id': '1', 25 | 'toVertexId': '1', 26 | 'fromVertexId': '2', 27 | 'relationship': { 28 | '_id': '1' 29 | }, 30 | 'properties': {} 31 | } 32 | } 33 | } 34 | } 35 | 36 | return (200, headers, json.dumps(response)) 37 | 38 | responses.add_callback( 39 | responses.POST, 'https://api.us.jupiterone.io/graphql', 40 | callback=request_callback, 41 | content_type='application/json', 42 | ) 43 | 44 | j1 = JupiterOneClient(account='testAccount', token='testToken') 45 | response = j1.create_relationship( 46 | relationship_key='relationship1', 47 | relationship_type='test_relationship', 48 | relationship_class='TestRelationship', 49 | from_entity_id='2', 50 | to_entity_id='1' 51 | ) 52 | 53 | assert type(response) == dict 54 | assert type(response['relationship']) == dict 55 | assert response['relationship']['_id'] == '1' 56 | assert response['edge']['toVertexId'] == '1' 57 | assert response['edge']['fromVertexId'] == '2' 58 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | .vscode/ 131 | .DS_Store 132 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # JupiterOne Python SDK 2 | 3 | [![Python 3.6](https://img.shields.io/badge/python-3.6-blue.svg)](https://www.python.org/downloads/release/python-360/) 4 | [![Python 3.7](https://img.shields.io/badge/python-3.7-blue.svg)](https://www.python.org/downloads/release/python-370/) 5 | 6 | 7 | A Python library for the [JupiterOne API](https://support.jupiterone.io/hc/en-us/articles/360022722094-JupiterOne-Platform-API). 8 | 9 | ## Installation 10 | 11 | Requires Python 3.6+ 12 | 13 | `pip install jupiterone` 14 | 15 | 16 | ## Usage 17 | 18 | ##### Create a new client: 19 | 20 | ```python 21 | from jupiterone import JupiterOneClient 22 | 23 | j1 = JupiterOneClient( 24 | account='', 25 | token='' 26 | ) 27 | ``` 28 | 29 | ##### Execute a query: 30 | 31 | ```python 32 | QUERY = 'FIND Host' 33 | query_result = j1.query_v1(QUERY) 34 | 35 | # Using LIMIT and SKIP for pagination 36 | query_result = j1.query_v1(QUERY, limit=5, skip=5) 37 | 38 | # Including deleted entities 39 | query_result = j1.query_v1(QUERY, include_deleted=True) 40 | 41 | # Tree query 42 | QUERY = 'FIND Host RETURN TREE' 43 | query_result = j1.query_v1(QUERY) 44 | ``` 45 | 46 | ##### Create an entity: 47 | 48 | Note that the CreateEntity mutation behaves like an upsert, so an non-existant entity will be created or an existing entity will be updated. 49 | 50 | ```python 51 | properties = { 52 | 'myProperty': 'myValue', 53 | 'tag.myTagProperty': 'value_will_be_a_tag' 54 | } 55 | 56 | entity = j1.create_entity( 57 | entity_key='my-unique-key', 58 | entity_type='my_type', 59 | entity_class='MyClass', 60 | properties=properties, 61 | timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime 62 | ) 63 | print(entity['entity']) 64 | ``` 65 | 66 | 67 | #### Update an existing entity: 68 | Only send in properties you want to add or update, other existing properties will not be modified. 69 | 70 | ```python 71 | properties = { 72 | 'newProperty': 'newPropertyValue' 73 | } 74 | 75 | j1.update_entity( 76 | entity_id='', 77 | properties=properties 78 | ) 79 | ``` 80 | 81 | 82 | #### Delete an entity: 83 | 84 | ```python 85 | j1.delete_entity(entit_id='') 86 | ``` 87 | 88 | ##### Create a relationship 89 | 90 | ```python 91 | j1.create_relationship( 92 | relationship_key='this_entity_relates_to_that_entity', 93 | relationship_type='my_relationship_type', 94 | relationship_class='MYRELATIONSHIP', 95 | from_entity_id='', 96 | to_entity_id='' 97 | ) 98 | ``` 99 | 100 | ##### Delete a relationship 101 | 102 | ```python 103 | j1.delete_relationship(relationship_id='') 104 | ``` 105 | -------------------------------------------------------------------------------- /jupiterone/constants.py: -------------------------------------------------------------------------------- 1 | J1QL_SKIP_COUNT = 250 2 | J1QL_LIMIT_COUNT = 250 3 | 4 | QUERY_V1 = """ 5 | query J1QL($query: String!, $variables: JSON, $dryRun: Boolean, $includeDeleted: Boolean) { 6 | queryV1(query: $query, variables: $variables, dryRun: $dryRun, includeDeleted: $includeDeleted) { 7 | type 8 | data 9 | } 10 | } 11 | """ 12 | 13 | CURSOR_QUERY_V1 = """ 14 | query J1QL_v2($query: String!, $variables: JSON, $flags: QueryV1Flags, $includeDeleted: Boolean, $cursor: String) { 15 | queryV1( 16 | query: $query 17 | variables: $variables 18 | deferredResponse: DISABLED 19 | flags: $flags 20 | includeDeleted: $includeDeleted 21 | cursor: $cursor 22 | ) { 23 | type 24 | data 25 | cursor 26 | __typename 27 | } 28 | } 29 | """ 30 | 31 | CREATE_ENTITY = """ 32 | mutation CreateEntity( 33 | $entityKey: String! 34 | $entityType: String! 35 | $entityClass: [String!]! 36 | $properties: JSON 37 | ) { 38 | createEntity( 39 | entityKey: $entityKey 40 | entityType: $entityType 41 | entityClass: $entityClass 42 | properties: $properties 43 | ) { 44 | entity { 45 | _id 46 | } 47 | vertex { 48 | id 49 | entity { 50 | _id 51 | } 52 | } 53 | } 54 | } 55 | """ 56 | 57 | DELETE_ENTITY = """ 58 | mutation DeleteEntity($entityId: String!, $timestamp: Long) { 59 | deleteEntity(entityId: $entityId, timestamp: $timestamp) { 60 | entity { 61 | _id 62 | } 63 | vertex { 64 | id 65 | entity { 66 | _id 67 | } 68 | properties 69 | } 70 | } 71 | } 72 | """ 73 | 74 | UPDATE_ENTITY = """ 75 | mutation UpdateEntity($entityId: String!, $properties: JSON) { 76 | updateEntity(entityId: $entityId, properties: $properties) { 77 | entity { 78 | _id 79 | } 80 | vertex { 81 | id 82 | } 83 | } 84 | } 85 | """ 86 | 87 | CREATE_RELATIONSHIP = """ 88 | mutation CreateRelationship( 89 | $relationshipKey: String! 90 | $relationshipType: String! 91 | $relationshipClass: String! 92 | $fromEntityId: String! 93 | $toEntityId: String! 94 | $properties: JSON 95 | ) { 96 | createRelationship( 97 | relationshipKey: $relationshipKey 98 | relationshipType: $relationshipType 99 | relationshipClass: $relationshipClass 100 | fromEntityId: $fromEntityId 101 | toEntityId: $toEntityId 102 | properties: $properties 103 | ) { 104 | relationship { 105 | _id 106 | } 107 | edge { 108 | id 109 | toVertexId 110 | fromVertexId 111 | relationship { 112 | _id 113 | } 114 | properties 115 | } 116 | } 117 | } 118 | """ 119 | 120 | DELETE_RELATIONSHIP = """ 121 | mutation DeleteRelationship($relationshipId: String! $timestamp: Long) { 122 | deleteRelationship (relationshipId: $relationshipId, timestamp: $timestamp) { 123 | relationship { 124 | _id 125 | } 126 | edge { 127 | id 128 | toVertexId 129 | fromVertexId 130 | relationship { 131 | _id 132 | } 133 | properties 134 | } 135 | } 136 | } 137 | """ -------------------------------------------------------------------------------- /jupiterone/client.py: -------------------------------------------------------------------------------- 1 | """ Python SDK for JupiterOne GraphQL API """ 2 | # pylint: disable=W0212,no-name-in-module 3 | # see https://github.com/PyCQA/pylint/issues/409 4 | 5 | import json 6 | from typing import Dict, List 7 | 8 | import requests 9 | from retrying import retry 10 | from warnings import warn 11 | 12 | from jupiterone.errors import ( 13 | JupiterOneClientError, 14 | JupiterOneApiRetryError, 15 | JupiterOneApiError 16 | ) 17 | 18 | from jupiterone.constants import ( 19 | J1QL_SKIP_COUNT, 20 | J1QL_LIMIT_COUNT, 21 | QUERY_V1, 22 | CREATE_ENTITY, 23 | DELETE_ENTITY, 24 | UPDATE_ENTITY, 25 | CREATE_RELATIONSHIP, 26 | DELETE_RELATIONSHIP, 27 | CURSOR_QUERY_V1 28 | ) 29 | 30 | def retry_on_429(exc): 31 | """ Used to trigger retry on rate limit """ 32 | return isinstance(exc, JupiterOneApiRetryError) 33 | 34 | 35 | class JupiterOneClient: 36 | """ Python client class for the JupiterOne GraphQL API """ 37 | # pylint: disable=too-many-instance-attributes 38 | 39 | DEFAULT_URL = 'https://api.us.jupiterone.io' 40 | 41 | RETRY_OPTS = { 42 | 'wait_exponential_multiplier': 1000, 43 | 'wait_exponential_max': 10000, 44 | 'stop_max_delay': 300000, 45 | 'retry_on_exception': retry_on_429 46 | } 47 | 48 | def __init__(self, account: str = None, token: str = None, url: str = DEFAULT_URL): 49 | self.account = account 50 | self.token = token 51 | self.url = url 52 | self.query_endpoint = self.url + '/graphql' 53 | self.rules_endpoint = self.url + '/rules/graphql' 54 | self.headers = { 55 | 'Authorization': 'Bearer {}'.format(self.token), 56 | 'LifeOmic-Account': self.account 57 | } 58 | 59 | @property 60 | def account(self): 61 | """ Your JupiterOne account ID """ 62 | return self._account 63 | 64 | @account.setter 65 | def account(self, value: str): 66 | """ Your JupiterOne account ID """ 67 | if not value: 68 | raise JupiterOneClientError('account is required') 69 | self._account = value 70 | 71 | @property 72 | def token(self): 73 | """ Your JupiteOne access token """ 74 | return self._token 75 | 76 | @token.setter 77 | def token(self, value: str): 78 | """ Your JupiteOne access token """ 79 | if not value: 80 | raise JupiterOneClientError('token is required') 81 | self._token = value 82 | 83 | # pylint: disable=R1710 84 | @retry(**RETRY_OPTS) 85 | def _execute_query(self, query: str, variables: Dict = None) -> Dict: 86 | """ Executes query against graphql endpoint """ 87 | 88 | data = { 89 | 'query': query 90 | } 91 | if variables: 92 | data.update(variables=variables) 93 | 94 | response = requests.post(self.query_endpoint, headers=self.headers, json=data) 95 | 96 | # It is still unclear if all responses will have a status 97 | # code of 200 or if 429 will eventually be used to 98 | # indicate rate limitting. J1 devs are aware. 99 | if response.status_code == 200: 100 | if response._content: 101 | content = json.loads(response._content) 102 | if 'errors' in content: 103 | errors = content['errors'] 104 | if len(errors) == 1: 105 | if '429' in errors[0]['message']: 106 | raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded') 107 | raise JupiterOneApiError(content.get('errors')) 108 | return response.json() 109 | 110 | elif response.status_code == 401: 111 | raise JupiterOneApiError('JupiterOne API query is unauthorized, check credentials.') 112 | 113 | elif response.status_code in [429, 503]: 114 | raise JupiterOneApiRetryError('JupiterOne API rate limit exceeded') 115 | 116 | else: 117 | content = response._content 118 | if isinstance(content, (bytes, bytearray)): 119 | content = content.decode("utf-8") 120 | if 'application/json' in response.headers.get('Content-Type', 'text/plain'): 121 | data = json.loads(content) 122 | content = data.get('error', data.get('errors', content)) 123 | raise JupiterOneApiError('{}:{}'.format(response.status_code, content)) 124 | 125 | def _cursor_query(self, query: str, cursor: str = None, include_deleted: bool = False) -> Dict: 126 | """ Performs a V1 graph query using cursor pagination 127 | args: 128 | query (str): Query text 129 | cursor (str): A pagination cursor for the initial query 130 | include_deleted (bool): Include recently deleted entities in query/search 131 | """ 132 | 133 | results: List = [] 134 | while True: 135 | variables = { 136 | 'query': query, 137 | 'includeDeleted': include_deleted 138 | } 139 | 140 | if cursor is not None: 141 | variables['cursor'] = cursor 142 | 143 | response = self._execute_query(query=CURSOR_QUERY_V1, variables=variables) 144 | data = response['data']['queryV1']['data'] 145 | 146 | if 'vertices' in data and 'edges' in data: 147 | return data 148 | 149 | results.extend(data) 150 | 151 | if 'cursor' in response['data']['queryV1'] and response['data']['queryV1']['cursor'] is not None: 152 | cursor = response['data']['queryV1']['cursor'] 153 | else: 154 | break 155 | 156 | return {'data': results} 157 | 158 | def _limit_and_skip_query(self, query: str, skip: int = J1QL_SKIP_COUNT, limit: int = J1QL_LIMIT_COUNT, include_deleted: bool = False) -> Dict: 159 | results: List = [] 160 | page: int = 0 161 | 162 | while True: 163 | variables = { 164 | 'query': f"{query} SKIP {page * skip} LIMIT {limit}", 165 | 'includeDeleted': include_deleted 166 | } 167 | response = self._execute_query( 168 | query=QUERY_V1, 169 | variables=variables 170 | ) 171 | 172 | data = response['data']['queryV1']['data'] 173 | 174 | # If tree query then no pagination 175 | if 'vertices' in data and 'edges' in data: 176 | return data 177 | 178 | if len(data) < J1QL_SKIP_COUNT: 179 | results.extend(data) 180 | break 181 | 182 | results.extend(data) 183 | page += 1 184 | 185 | return {'data': results} 186 | 187 | def query_v1(self, query: str, **kwargs) -> Dict: 188 | """ Performs a V1 graph query 189 | args: 190 | query (str): Query text 191 | skip (int): Skip entity count 192 | limit (int): Limit entity count 193 | cursor (str): A pagination cursor for the initial query 194 | include_deleted (bool): Include recently deleted entities in query/search 195 | """ 196 | uses_limit_and_skip: bool = 'skip' in kwargs.keys() or 'limit' in kwargs.keys() 197 | skip: int = kwargs.pop('skip', J1QL_SKIP_COUNT) 198 | limit: int = kwargs.pop('limit', J1QL_LIMIT_COUNT) 199 | include_deleted: bool = kwargs.pop('include_deleted', False) 200 | cursor: str = kwargs.pop('cursor', None) 201 | 202 | if uses_limit_and_skip: 203 | 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) 204 | return self._limit_and_skip_query( 205 | query=query, 206 | skip=skip, 207 | limit=limit, 208 | include_deleted=include_deleted 209 | ) 210 | else: 211 | return self._cursor_query( 212 | query=query, 213 | cursor=cursor, 214 | include_deleted=include_deleted 215 | ) 216 | 217 | def create_entity(self, **kwargs) -> Dict: 218 | """ Creates an entity in graph. It will also update an existing entity. 219 | 220 | args: 221 | entity_key (str): Unique key for the entity 222 | entity_type (str): Value for _type of entity 223 | entity_class (str): Value for _class of entity 224 | timestamp (int): Specify createdOn timestamp 225 | properties (dict): Dictionary of key/value entity properties 226 | """ 227 | variables = { 228 | 'entityKey': kwargs.pop('entity_key'), 229 | 'entityType': kwargs.pop('entity_type'), 230 | 'entityClass': kwargs.pop('entity_class') 231 | } 232 | 233 | timestamp: int = kwargs.pop('timestamp', None) 234 | properties: Dict = kwargs.pop('properties', None) 235 | 236 | if timestamp: 237 | variables.update(timestamp=timestamp) 238 | if properties: 239 | variables.update(properties=properties) 240 | 241 | response = self._execute_query( 242 | query=CREATE_ENTITY, 243 | variables=variables 244 | ) 245 | return response['data']['createEntity'] 246 | 247 | def delete_entity(self, entity_id: str = None) -> Dict: 248 | """ Deletes an entity from the graph. Note this is a hard delete. 249 | 250 | args: 251 | entity_id (str): Entity ID for entity to delete 252 | """ 253 | variables = { 254 | 'entityId': entity_id 255 | } 256 | response = self._execute_query(DELETE_ENTITY, variables=variables) 257 | return response['data']['deleteEntity'] 258 | 259 | def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: 260 | """ 261 | Update an existing entity. 262 | 263 | args: 264 | entity_id (str): The _id of the entity to udate 265 | properties (dict): Dictionary of key/value entity properties 266 | """ 267 | variables = { 268 | 'entityId': entity_id, 269 | 'properties': properties 270 | } 271 | response = self._execute_query(UPDATE_ENTITY, variables=variables) 272 | return response['data']['updateEntity'] 273 | 274 | def create_relationship(self, **kwargs) -> Dict: 275 | """ 276 | Create a relationship (edge) between two entities (veritces). 277 | 278 | args: 279 | relationship_key (str): Unique key for the relationship 280 | relationship_type (str): Value for _type of relationship 281 | relationship_class (str): Value for _class of relationship 282 | from_entity_id (str): Entity ID of the source vertex 283 | to_entity_id (str): Entity ID of the destination vertex 284 | """ 285 | variables = { 286 | 'relationshipKey': kwargs.pop('relationship_key'), 287 | 'relationshipType': kwargs.pop('relationship_type'), 288 | 'relationshipClass': kwargs.pop('relationship_class'), 289 | 'fromEntityId': kwargs.pop('from_entity_id'), 290 | 'toEntityId': kwargs.pop('to_entity_id') 291 | } 292 | 293 | properties = kwargs.pop('properties', None) 294 | if properties: 295 | variables['properties'] = properties 296 | 297 | response = self._execute_query( 298 | query=CREATE_RELATIONSHIP, 299 | variables=variables 300 | ) 301 | return response['data']['createRelationship'] 302 | 303 | def delete_relationship(self, relationship_id: str = None): 304 | """ Deletes a relationship between two entities. 305 | 306 | args: 307 | relationship_id (str): The ID of the relationship 308 | """ 309 | variables = { 310 | 'relationshipId': relationship_id 311 | } 312 | 313 | response = self._execute_query( 314 | DELETE_RELATIONSHIP, 315 | variables=variables 316 | ) 317 | return response['data']['deleteRelationship'] 318 | -------------------------------------------------------------------------------- /tests/test_query.py: -------------------------------------------------------------------------------- 1 | import json 2 | import pytest 3 | import responses 4 | from collections import Counter 5 | 6 | from jupiterone.client import JupiterOneClient 7 | from jupiterone.constants import QUERY_V1 8 | from jupiterone.errors import JupiterOneApiError 9 | 10 | def build_results(response_code: int = 200, cursor: str = None, max_pages: int = 1): 11 | pages = Counter(requests=0) 12 | 13 | def request_callback(request): 14 | headers = { 15 | 'Content-Type': 'application/json' 16 | } 17 | 18 | response = { 19 | 'data': { 20 | 'queryV1': { 21 | 'type': 'list', 22 | 'data': [ 23 | { 24 | 'id': '1', 25 | 'entity': { 26 | '_rawDataHashes': '1', 27 | '_integrationDefinitionId': '1', 28 | '_integrationName': '1', 29 | '_beginOn': 1580482083079, 30 | 'displayName': 'host1', 31 | '_class': ['Host'], 32 | '_scope': 'aws_instance', 33 | '_version': 1, 34 | '_integrationClass': 'CSP', 35 | '_accountId': 'testAccount', 36 | '_id': '1', 37 | '_key': 'key1', 38 | '_type': ['aws_instance'], 39 | '_deleted': False, 40 | '_integrationInstanceId': '1', 41 | '_integrationType': 'aws', 42 | '_source': 'integration-managed', 43 | '_createdOn': 1578093840019 44 | }, 45 | 'properties': { 46 | 'id': 'host1', 47 | 'active': True 48 | } 49 | } 50 | ] 51 | } 52 | } 53 | } 54 | 55 | if cursor is not None and pages.get('requests') < max_pages: 56 | response['data']['queryV1']['cursor'] = cursor 57 | 58 | pages.update(requests=1) 59 | 60 | return (response_code, headers, json.dumps(response)) 61 | 62 | return request_callback 63 | 64 | 65 | def build_error_results(response_code: int, response_content, response_type: str = 'application/json'): 66 | def request_callback(request): 67 | headers = { 68 | 'Content-Type': response_type 69 | } 70 | return response_code, headers, response_content 71 | 72 | return request_callback 73 | 74 | 75 | @responses.activate 76 | def test_execute_query(): 77 | 78 | responses.add_callback( 79 | responses.POST, 'https://api.us.jupiterone.io/graphql', 80 | callback=build_results(), 81 | content_type='application/json', 82 | ) 83 | 84 | j1 = JupiterOneClient(account='testAccount', token='testToken') 85 | query = "find Host with _id='1'" 86 | variables = { 87 | 'query': query, 88 | 'includeDeleted': False 89 | } 90 | 91 | response = j1._execute_query( 92 | query=QUERY_V1, 93 | variables=variables 94 | ) 95 | assert 'data' in response 96 | assert 'queryV1' in response['data'] 97 | assert len(response['data']['queryV1']['data']) == 1 98 | assert type(response['data']['queryV1']['data']) == list 99 | assert response['data']['queryV1']['data'][0]['entity']['_id'] == '1' 100 | 101 | 102 | @responses.activate 103 | def test_limit_skip_query_v1(): 104 | 105 | responses.add_callback( 106 | responses.POST, 'https://api.us.jupiterone.io/graphql', 107 | callback=build_results(), 108 | content_type='application/json', 109 | ) 110 | 111 | j1 = JupiterOneClient(account='testAccount', token='testToken') 112 | query = "find Host with _id='1'" 113 | response = j1.query_v1( 114 | query=query, 115 | limit=250, 116 | skip=0 117 | ) 118 | 119 | assert type(response) == dict 120 | assert len(response['data']) == 1 121 | assert type(response['data']) == list 122 | assert response['data'][0]['entity']['_id'] == '1' 123 | 124 | @responses.activate 125 | def test_cursor_query_v1(): 126 | 127 | responses.add_callback( 128 | responses.POST, 'https://api.us.jupiterone.io/graphql', 129 | callback=build_results(cursor='cursor_value'), 130 | content_type='application/json', 131 | ) 132 | 133 | responses.add_callback( 134 | responses.POST, 'https://api.us.jupiterone.io/graphql', 135 | callback=build_results(), 136 | content_type='application/json', 137 | ) 138 | 139 | j1 = JupiterOneClient(account='testAccount', token='testToken') 140 | query = "find Host with _id='1'" 141 | 142 | response = j1.query_v1( 143 | query=query, 144 | ) 145 | 146 | assert type(response) == dict 147 | assert len(response['data']) == 2 148 | assert type(response['data']) == list 149 | assert response['data'][0]['entity']['_id'] == '1' 150 | 151 | @responses.activate 152 | def test_limit_skip_tree_query_v1(): 153 | 154 | def request_callback(request): 155 | headers = { 156 | 'Content-Type': 'application/json' 157 | } 158 | 159 | response = { 160 | 'data': { 161 | 'queryV1': { 162 | 'type': 'tree', 163 | 'data': { 164 | 'vertices': [ 165 | { 166 | 'id': '1', 167 | 'entity': {}, 168 | 'properties': {} 169 | } 170 | ], 171 | 'edges': [] 172 | } 173 | } 174 | } 175 | } 176 | 177 | return (200, headers, json.dumps(response)) 178 | 179 | responses.add_callback( 180 | responses.POST, 'https://api.us.jupiterone.io/graphql', 181 | callback=request_callback, 182 | content_type='application/json', 183 | ) 184 | 185 | j1 = JupiterOneClient(account='testAccount', token='testToken') 186 | query = "find Host with _id='1' return tree" 187 | response = j1.query_v1( 188 | query=query, 189 | limit=250, 190 | skip=0 191 | ) 192 | 193 | assert type(response) == dict 194 | assert 'edges' in response 195 | assert 'vertices' in response 196 | assert type(response['edges']) == list 197 | assert type(response['vertices']) == list 198 | assert response['vertices'][0]['id'] == '1' 199 | 200 | @responses.activate 201 | def test_cursor_tree_query_v1(): 202 | 203 | def request_callback(request): 204 | headers = { 205 | 'Content-Type': 'application/json' 206 | } 207 | 208 | response = { 209 | 'data': { 210 | 'queryV1': { 211 | 'type': 'tree', 212 | 'data': { 213 | 'vertices': [ 214 | { 215 | 'id': '1', 216 | 'entity': {}, 217 | 'properties': {} 218 | } 219 | ], 220 | 'edges': [] 221 | } 222 | } 223 | } 224 | } 225 | 226 | return (200, headers, json.dumps(response)) 227 | 228 | responses.add_callback( 229 | responses.POST, 'https://api.us.jupiterone.io/graphql', 230 | callback=request_callback, 231 | content_type='application/json', 232 | ) 233 | 234 | j1 = JupiterOneClient(account='testAccount', token='testToken') 235 | query = "find Host with _id='1' return tree" 236 | response = j1.query_v1(query) 237 | 238 | assert type(response) == dict 239 | assert 'edges' in response 240 | assert 'vertices' in response 241 | assert type(response['edges']) == list 242 | assert type(response['vertices']) == list 243 | assert response['vertices'][0]['id'] == '1' 244 | 245 | @responses.activate 246 | def test_retry_on_limit_skip_query(): 247 | responses.add_callback( 248 | responses.POST, 'https://api.us.jupiterone.io/graphql', 249 | callback=build_results(response_code=429), 250 | content_type='application/json', 251 | ) 252 | 253 | responses.add_callback( 254 | responses.POST, 'https://api.us.jupiterone.io/graphql', 255 | callback=build_results(response_code=503), 256 | content_type='application/json', 257 | ) 258 | 259 | responses.add_callback( 260 | responses.POST, 'https://api.us.jupiterone.io/graphql', 261 | callback=build_results(), 262 | content_type='application/json', 263 | ) 264 | 265 | j1 = JupiterOneClient(account='testAccount', token='testToken') 266 | query = "find Host with _id='1'" 267 | response = j1.query_v1( 268 | query=query, 269 | limit=250, 270 | skip=0 271 | ) 272 | 273 | assert type(response) == dict 274 | assert len(response['data']) == 1 275 | assert type(response['data']) == list 276 | assert response['data'][0]['entity']['_id'] == '1' 277 | 278 | @responses.activate 279 | def test_retry_on_cursor_query(): 280 | responses.add_callback( 281 | responses.POST, 'https://api.us.jupiterone.io/graphql', 282 | callback=build_results(response_code=429), 283 | content_type='application/json', 284 | ) 285 | 286 | responses.add_callback( 287 | responses.POST, 'https://api.us.jupiterone.io/graphql', 288 | callback=build_results(response_code=503), 289 | content_type='application/json', 290 | ) 291 | 292 | responses.add_callback( 293 | responses.POST, 'https://api.us.jupiterone.io/graphql', 294 | callback=build_results(), 295 | content_type='application/json', 296 | ) 297 | 298 | j1 = JupiterOneClient(account='testAccount', token='testToken') 299 | query = "find Host with _id='1'" 300 | response = j1.query_v1( 301 | query=query 302 | ) 303 | 304 | assert type(response) == dict 305 | assert len(response['data']) == 1 306 | assert type(response['data']) == list 307 | assert response['data'][0]['entity']['_id'] == '1' 308 | 309 | @responses.activate 310 | def test_avoid_retry_on_limit_skip_query(): 311 | responses.add_callback( 312 | responses.POST, 'https://api.us.jupiterone.io/graphql', 313 | callback=build_results(response_code=404), 314 | content_type='application/json', 315 | ) 316 | 317 | j1 = JupiterOneClient(account='testAccount', token='testToken') 318 | query = "find Host with _id='1'" 319 | with pytest.raises(JupiterOneApiError): 320 | j1.query_v1( 321 | query=query, 322 | limit=250, 323 | skip=0 324 | ) 325 | 326 | @responses.activate 327 | def test_avoid_retry_on_cursor_query(): 328 | responses.add_callback( 329 | responses.POST, 'https://api.us.jupiterone.io/graphql', 330 | callback=build_results(response_code=404), 331 | content_type='application/json', 332 | ) 333 | 334 | j1 = JupiterOneClient(account='testAccount', token='testToken') 335 | query = "find Host with _id='1'" 336 | with pytest.raises(JupiterOneApiError): 337 | j1.query_v1( 338 | query=query 339 | ) 340 | 341 | @responses.activate 342 | def test_warn_limit_and_skip_deprecated(): 343 | responses.add_callback( 344 | responses.POST, 'https://api.us.jupiterone.io/graphql', 345 | callback=build_results(), 346 | content_type='application/json', 347 | ) 348 | 349 | j1 = JupiterOneClient(account='testAccount', token='testToken') 350 | query = "find Host with _id='1'" 351 | 352 | with pytest.warns(DeprecationWarning): 353 | j1.query_v1( 354 | query=query, 355 | limit=250, 356 | skip=0 357 | ) 358 | 359 | 360 | @responses.activate 361 | def test_unauthorized_query_v1(): 362 | responses.add_callback( 363 | responses.POST, 'https://api.us.jupiterone.io/graphql', 364 | callback=build_error_results(401, b'Unauthorized', 'text/plain'), 365 | content_type='application/json', 366 | ) 367 | 368 | j1 = JupiterOneClient(account='testAccount', token='bogusToken') 369 | query = "find Host with _id='1' return tree" 370 | 371 | with pytest.raises(JupiterOneApiError) as exc_info: 372 | j1.query_v1(query) 373 | 374 | assert exc_info.value.args[0] == 'JupiterOne API query is unauthorized, check credentials.' 375 | 376 | 377 | @responses.activate 378 | def test_unexpected_string_error_query_v1(): 379 | responses.add_callback( 380 | responses.POST, 'https://api.us.jupiterone.io/graphql', 381 | callback=build_error_results(500, 'String exception on server', 'text/plain'), 382 | content_type='application/json', 383 | ) 384 | 385 | j1 = JupiterOneClient(account='testAccount', token='bogusToken') 386 | query = "find Host with _id='1' return tree" 387 | 388 | with pytest.raises(JupiterOneApiError) as exc_info: 389 | j1.query_v1(query) 390 | 391 | assert exc_info.value.args[0] == '500:String exception on server' 392 | 393 | 394 | @responses.activate 395 | def test_unexpected_json_error_query_v1(): 396 | error_json = { 397 | 'error': 'Bad Gateway' 398 | } 399 | 400 | responses.add_callback( 401 | responses.POST, 'https://api.us.jupiterone.io/graphql', 402 | callback=build_error_results(502, json.dumps(error_json), ), 403 | content_type='application/json', 404 | ) 405 | 406 | j1 = JupiterOneClient(account='testAccount', token='bogusToken') 407 | query = "find Host with _id='1' return tree" 408 | 409 | with pytest.raises(JupiterOneApiError) as exc_info: 410 | j1.query_v1(query) 411 | 412 | assert exc_info.value.args[0] == '502:Bad Gateway' 413 | 414 | 415 | @responses.activate 416 | def test_unexpected_json_errors_query_v1(): 417 | errors_json = { 418 | 'errors': ['First error', 'Second error', 'Third error'] 419 | } 420 | 421 | responses.add_callback( 422 | responses.POST, 'https://api.us.jupiterone.io/graphql', 423 | callback=build_error_results(500, json.dumps(errors_json), ), 424 | content_type='application/json', 425 | ) 426 | 427 | j1 = JupiterOneClient(account='testAccount', token='bogusToken') 428 | query = "find Host with _id='1' return tree" 429 | 430 | with pytest.raises(JupiterOneApiError) as exc_info: 431 | j1.query_v1(query) 432 | 433 | assert exc_info.value.args[0] == "500:['First error', 'Second error', 'Third error']" 434 | --------------------------------------------------------------------------------