Update JupiterOne fork with auth0 master branch

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.
 

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!

opslevel.yml

@@ -0,0 +1,6 @@
+---
+version: 1
+repository:
+  owner: cloud_security_ops
+  tier:
+  tags:

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=[