From 86eef4adf7971d99c1b9951b0b1aee416ab9455f Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 19:56:28 -0600 Subject: [PATCH 01/11] fix example scripts and bump version --- examples/02_entity_management.py | 24 +++++++++--------------- examples/03_relationship_management.py | 16 ++++++---------- setup.py | 2 +- 3 files changed, 16 insertions(+), 26 deletions(-) diff --git a/examples/02_entity_management.py b/examples/02_entity_management.py index a618f0a..07b0b6f 100644 --- a/examples/02_entity_management.py +++ b/examples/02_entity_management.py @@ -76,11 +76,9 @@ def create_entity_examples(j1): 'backupRetentionPeriod': 7, 'tag.Environment': 'production', 'tag.Team': 'data', - 'metadata': { - 'createdBy': 'terraform', - 'lastBackup': '2024-01-01T00:00:00Z', - 'maintenanceWindow': 'sun:03:00-sun:04:00' - } + 'createdBy': 'terraform', + 'lastBackup': '2024-01-01T00:00:00Z', + 'maintenanceWindow': 'sun:03:00-sun:04:00' }, timestamp=int(time.time()) * 1000 ) @@ -122,16 +120,12 @@ def update_entity_examples(j1, entity_id): entity_id=entity_id, properties={ 'isActive': False, - 'maintenanceWindow': { - 'start': '2024-01-01T00:00:00Z', - 'end': '2024-01-01T04:00:00Z', - 'reason': 'scheduled_maintenance' - }, - 'metadata': { - 'maintenancePerformedBy': 'admin@company.com', - 'maintenanceType': 'security_patches', - 'estimatedDuration': '4 hours' - } + 'maintenanceWindowStart': '2024-01-01T00:00:00Z', + 'maintenanceWindowEnd': '2024-01-01T04:00:00Z', + 'maintenanceReason': 'scheduled_maintenance', + 'maintenancePerformedBy': 'admin@company.com', + 'maintenanceType': 'security_patches', + 'estimatedDuration': '4 hours' } ) print(f"Updated with complex properties\n") diff --git a/examples/03_relationship_management.py b/examples/03_relationship_management.py index 73037cb..04e7997 100644 --- a/examples/03_relationship_management.py +++ b/examples/03_relationship_management.py @@ -100,11 +100,9 @@ def create_relationship_examples(j1, from_entity_id, to_entity_id): 'version': '2.1.0', 'installPath': '/usr/local/bin/software', 'permissions': ['read', 'execute'], - 'metadata': { - 'installer': 'package_manager', - 'verified': True, - 'checksum': 'sha256:abc123...' - }, + 'installer': 'package_manager', + 'verified': True, + 'checksum': 'sha256:abc123...', 'tag.InstallationType': 'automated', 'tag.Verified': 'true' } @@ -142,11 +140,9 @@ def update_relationship_examples(j1, relationship_id, from_entity_id, to_entity_ 'lastModified': int(time.time()) * 1000, 'modifiedBy': 'security_team', 'expiresOn': int(time.time() + 86400) * 1000, # 24 hours from now - 'auditLog': { - 'previousLevel': 'write', - 'reason': 'promotion_requested', - 'approvedBy': 'security_manager' - } + 'previousLevel': 'write', + 'promotionReason': 'promotion_requested', + 'approvedBy': 'security_manager' } ) print(f"Updated with complex properties\n") diff --git a/setup.py b/setup.py index 57b1d2d..318fe06 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="jupiterone", - version="2.0.0", + version="2.0.1", description="A Python client for the JupiterOne API", license="MIT License", author="JupiterOne", From 09286277d8dea579cf3aa7a15a6d95d596978123 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 20:05:14 -0600 Subject: [PATCH 02/11] Update bulk_upload.py --- examples/bulk_upload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/bulk_upload.py b/examples/bulk_upload.py index a8ac5ea..2f63942 100644 --- a/examples/bulk_upload.py +++ b/examples/bulk_upload.py @@ -7,7 +7,7 @@ j1 = JupiterOneClient(account=account, token=token, url=url) -instance_id = "e7113c37-1ea8-4d00-9b82-c24952e70916" +instance_id = "" sync_job = j1.start_sync_job( instance_id=instance_id, From 7633738c953addd1c242c1209ac2af8bccec1445 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 20:10:24 -0600 Subject: [PATCH 03/11] Update README.md --- README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index a7cf46a..2e87557 100644 --- a/README.md +++ b/README.md @@ -128,8 +128,7 @@ entity = j1.create_entity( entity_key='my-unique-key', entity_type='my_type', entity_class='MyClass', - properties=properties, - timestamp=int(time.time()) * 1000 # Optional, defaults to current datetime + properties=properties ) print(entity['entity']) From e89b28c7273d2fb3138065cfcefccb191f25d489 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 20:23:13 -0600 Subject: [PATCH 04/11] Update test_alert_rule_methods.py --- tests/test_alert_rule_methods.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/tests/test_alert_rule_methods.py b/tests/test_alert_rule_methods.py index 0634a40..f373afa 100644 --- a/tests/test_alert_rule_methods.py +++ b/tests/test_alert_rule_methods.py @@ -216,6 +216,10 @@ def test_update_alert_rule_basic(self, mock_execute_query, mock_get_details): "pollingInterval": "ONE_DAY", "tags": ["old-tag"], "labels": [], + "triggerActionsOnNewEntitiesOnly": True, + "ignorePreviousResults": False, + "notifyOnFailure": True, + "templates": {}, "operations": [{ "__typename": "Operation", "when": {"type": "FILTER", "condition": ["AND", ["queries.query0.total", ">", 0]]}, @@ -300,24 +304,25 @@ def test_fetch_evaluation_result_download_url(self, mock_execute_query): assert result == mock_response mock_execute_query.assert_called_once() - @patch('jupiterone.client.requests.get') - def test_fetch_downloaded_evaluation_results_success(self, mock_get): + def test_fetch_downloaded_evaluation_results_success(self): """Test fetch_downloaded_evaluation_results method - success""" mock_response = Mock() mock_response.json.return_value = {"data": [{"id": "result-1"}]} - mock_get.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.get.return_value = mock_response - result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") + result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") - assert result == {"data": [{"id": "result-1"}]} - mock_get.assert_called_once() + assert result == {"data": [{"id": "result-1"}]} + mock_session.get.assert_called_once_with("https://example.com/download", timeout=60) - @patch('jupiterone.client.requests.get') - def test_fetch_downloaded_evaluation_results_exception(self, mock_get): + def test_fetch_downloaded_evaluation_results_exception(self): """Test fetch_downloaded_evaluation_results method - exception""" - mock_get.side_effect = Exception("Network error") + with patch.object(self.client, 'session') as mock_session: + mock_session.get.side_effect = Exception("Network error") - result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") + result = self.client.fetch_downloaded_evaluation_results(download_url="https://example.com/download") - assert isinstance(result, Exception) - assert str(result) == "Network error" \ No newline at end of file + assert isinstance(result, Exception) + assert str(result) == "Network error" \ No newline at end of file From ca4433097e471f675511f82e26d34a80c1c67194 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 20:32:03 -0600 Subject: [PATCH 05/11] Update test_client_init.py --- tests/test_client_init.py | 132 ++++++++++++++++++++------------------ 1 file changed, 71 insertions(+), 61 deletions(-) diff --git a/tests/test_client_init.py b/tests/test_client_init.py index e12d61b..dad689d 100644 --- a/tests/test_client_init.py +++ b/tests/test_client_init.py @@ -127,127 +127,137 @@ def setup_method(self): """Set up test fixtures""" self.client = JupiterOneClient(account="test-account", token="test-token") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_401_error(self, mock_session): + def test_execute_query_401_error(self): """Test _execute_query method with 401 error""" mock_response = Mock() mock_response.status_code = 401 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiError, match="401: Unauthorized"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiError, match="401: Unauthorized"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_429_error(self, mock_session): + def test_execute_query_429_error(self): """Test _execute_query method with 429 error""" mock_response = Mock() mock_response.status_code = 429 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_503_error(self, mock_session): + def test_execute_query_503_error(self): """Test _execute_query method with 503 error""" mock_response = Mock() mock_response.status_code = 503 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_504_error(self, mock_session): + def test_execute_query_504_error(self): """Test _execute_query method with 504 error""" mock_response = Mock() mock_response.status_code = 504 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="Gateway Timeout"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="Gateway Timeout"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_500_error(self, mock_session): + def test_execute_query_500_error(self): """Test _execute_query method with 500 error""" mock_response = Mock() mock_response.status_code = 500 - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiError, match="internal server error"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiError, match="internal server error"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_200_with_errors(self, mock_session): + def test_execute_query_200_with_errors(self): """Test _execute_query method with 200 status but GraphQL errors""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "errors": [{"message": "GraphQL error"}] } - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiError): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiError): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_200_with_429_in_errors(self, mock_session): + def test_execute_query_200_with_429_in_errors(self): """Test _execute_query method with 200 status but 429 in GraphQL errors""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "errors": [{"message": "429 rate limit exceeded"}] } - mock_session.post.return_value = mock_response + + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response - with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): - self.client._execute_query("test query") + with pytest.raises(JupiterOneApiRetryError, match="rate limit exceeded"): + self.client._execute_query("test query") - @patch.object(JupiterOneClient, 'session') - def test_execute_query_200_success(self, mock_session): + def test_execute_query_200_success(self): """Test _execute_query method with successful 200 response""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "data": {"result": "success"} } - mock_session.post.return_value = mock_response - - result = self.client._execute_query("test query") - assert result == {"data": {"result": "success"}} + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response + + result = self.client._execute_query("test query") + + assert result == {"data": {"result": "success"}} - @patch.object(JupiterOneClient, 'session') - def test_execute_query_with_variables(self, mock_session): + def test_execute_query_with_variables(self): """Test _execute_query method with variables""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "data": {"result": "success"} } - mock_session.post.return_value = mock_response - - variables = {"key": "value"} - self.client._execute_query("test query", variables=variables) - # Verify that variables were included in the request - call_args = mock_session.post.call_args - assert "variables" in call_args[1]["json"] - assert call_args[1]["json"]["variables"] == variables - - @patch.object(JupiterOneClient, 'session') - def test_execute_query_with_flags(self, mock_session): + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response + + variables = {"key": "value"} + self.client._execute_query("test query", variables=variables) + + # Verify that variables were included in the request + call_args = mock_session.post.call_args + assert "variables" in call_args[1]["json"] + assert call_args[1]["json"]["variables"] == variables + + def test_execute_query_with_flags(self): """Test _execute_query method includes flags""" mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = { "data": {"result": "success"} } - mock_session.post.return_value = mock_response - - self.client._execute_query("test query") - # Verify that flags were included in the request - call_args = mock_session.post.call_args - assert "flags" in call_args[1]["json"] - assert call_args[1]["json"]["flags"] == {"variableResultSize": True} \ No newline at end of file + with patch.object(self.client, 'session') as mock_session: + mock_session.post.return_value = mock_response + + self.client._execute_query("test query") + + # Verify that flags were included in the request + call_args = mock_session.post.call_args + assert "flags" in call_args[1]["json"] + assert call_args[1]["json"]["flags"] == {"variableResultSize": True} \ No newline at end of file From 1b7c588c6e0ee638678b4a6e356ffb68fa6289e3 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 20:37:16 -0600 Subject: [PATCH 06/11] example script usability improvements --- examples/09_custom_file_transfer_example.py | 6 +++++- examples/J1QLdeferredResponse.py | 4 ++-- examples/bulk_upload.py | 6 ++++++ examples/examples.py | 6 ++++++ 4 files changed, 19 insertions(+), 3 deletions(-) diff --git a/examples/09_custom_file_transfer_example.py b/examples/09_custom_file_transfer_example.py index 67c75cd..5bbfacf 100644 --- a/examples/09_custom_file_transfer_example.py +++ b/examples/09_custom_file_transfer_example.py @@ -20,7 +20,11 @@ import sys # Add the parent directory to the path so we can import the jupiterone client -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +try: + sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +except NameError: + # Handle case when __file__ is not available (e.g., when exec'd) + sys.path.append('..') from jupiterone.client import JupiterOneClient diff --git a/examples/J1QLdeferredResponse.py b/examples/J1QLdeferredResponse.py index 244cf3a..8c69709 100644 --- a/examples/J1QLdeferredResponse.py +++ b/examples/J1QLdeferredResponse.py @@ -13,8 +13,8 @@ # JupiterOne GraphQL API headers j1_graphql_headers = { 'Content-Type': 'application/json', - 'Authorization': 'Bearer ' + token, - 'Jupiterone-Account': acct + 'Authorization': 'Bearer ' + (token or ''), + 'Jupiterone-Account': acct or '' } gql_query = """ diff --git a/examples/bulk_upload.py b/examples/bulk_upload.py index 2f63942..b64f68d 100644 --- a/examples/bulk_upload.py +++ b/examples/bulk_upload.py @@ -5,6 +5,12 @@ token = os.environ.get("JUPITERONE_TOKEN") url = "https://graphql.us.jupiterone.io" +# Check if credentials are available +if not account or not token: + print("Error: JUPITERONE_ACCOUNT and JUPITERONE_TOKEN environment variables must be set") + print("This example script requires valid JupiterOne credentials to run") + exit(1) + j1 = JupiterOneClient(account=account, token=token, url=url) instance_id = "" diff --git a/examples/examples.py b/examples/examples.py index 061e509..e692540 100644 --- a/examples/examples.py +++ b/examples/examples.py @@ -8,6 +8,12 @@ token = os.environ.get("JUPITERONE_TOKEN") url = "https://graphql.us.jupiterone.io" +# Check if credentials are available +if not account or not token: + print("Error: JUPITERONE_ACCOUNT and JUPITERONE_TOKEN environment variables must be set") + print("This example script requires valid JupiterOne credentials to run") + exit(1) + j1 = JupiterOneClient(account=account, token=token, url=url) # query_v1 From c9f01bf4e8e49a5dabfb47f8f7a76e7919175263 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 20:56:10 -0600 Subject: [PATCH 07/11] fixed full test suite - updated skip limit query handling --- jupiterone/client.py | 2 +- tests/test_cursor_query_edge_cases.py | 6 +-- tests/test_limit_skip_query_edge_cases.py | 62 ++++++++++++++--------- 3 files changed, 43 insertions(+), 27 deletions(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index ec4fa52..2d74955 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -304,7 +304,7 @@ def _limit_and_skip_query( if "vertices" in data and "edges" in data: return data - if len(data) < J1QL_SKIP_COUNT: + if len(data) < skip: results.extend(data) break diff --git a/tests/test_cursor_query_edge_cases.py b/tests/test_cursor_query_edge_cases.py index 8c8cb20..6aafe30 100644 --- a/tests/test_cursor_query_edge_cases.py +++ b/tests/test_cursor_query_edge_cases.py @@ -211,7 +211,7 @@ def test_cursor_query_with_include_deleted_true(self, mock_execute_query): # Verify the query was called with includeDeleted=True mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is True # Verify the result @@ -241,7 +241,7 @@ def test_cursor_query_with_include_deleted_false(self, mock_execute_query): # Verify the query was called with includeDeleted=False mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is False # Verify the result @@ -271,7 +271,7 @@ def test_cursor_query_with_initial_cursor(self, mock_execute_query): # Verify the query was called with the initial cursor mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["cursor"] == "initial_cursor" # Verify the result diff --git a/tests/test_limit_skip_query_edge_cases.py b/tests/test_limit_skip_query_edge_cases.py index 7a83250..926cf6e 100644 --- a/tests/test_limit_skip_query_edge_cases.py +++ b/tests/test_limit_skip_query_edge_cases.py @@ -79,7 +79,9 @@ def test_limit_and_skip_query_multiple_pages_with_break(self, mock_execute_query "queryV1": { "data": [ {"id": "1", "name": "entity1"}, - {"id": "2", "name": "entity2"} + {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"} ] } } @@ -98,7 +100,7 @@ def test_limit_and_skip_query_multiple_pages_with_break(self, mock_execute_query mock_execute_query.side_effect = [mock_response1, mock_response2] - result = self.client._limit_and_skip_query("FIND * LIMIT 10") + result = self.client._limit_and_skip_query("FIND * LIMIT 10", skip=3) # Should call twice, but break on second page assert mock_execute_query.call_count == 2 @@ -107,6 +109,8 @@ def test_limit_and_skip_query_multiple_pages_with_break(self, mock_execute_query assert result == {"data": [ {"id": "1", "name": "entity1"}, {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"}, {"id": "3", "name": "entity3"} ]} @@ -134,8 +138,8 @@ def test_limit_and_skip_query_with_custom_skip_limit(self, mock_execute_query): # Verify the query was called with custom skip/limit mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - query = call_args[0][0] - assert "SKIP 0 LIMIT 25" in query # First page starts at 0 + variables = call_args[1]["variables"] + assert "SKIP 0 LIMIT 25" in variables["query"] # First page starts at 0 # Verify the result assert result == {"data": [ @@ -151,7 +155,9 @@ def test_limit_and_skip_query_multiple_pages_with_custom_values(self, mock_execu "queryV1": { "data": [ {"id": "1", "name": "entity1"}, - {"id": "2", "name": "entity2"} + {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"} ] } } @@ -172,7 +178,7 @@ def test_limit_and_skip_query_multiple_pages_with_custom_values(self, mock_execu result = self.client._limit_and_skip_query( "FIND * LIMIT 10", - skip=10, + skip=3, limit=10 ) @@ -181,13 +187,15 @@ def test_limit_and_skip_query_multiple_pages_with_custom_values(self, mock_execu # Verify the queries were called with correct skip values call_args_list = mock_execute_query.call_args_list - assert "SKIP 0 LIMIT 10" in call_args_list[0][0][0] # First page - assert "SKIP 10 LIMIT 10" in call_args_list[1][0][0] # Second page + assert "SKIP 0 LIMIT 10" in call_args_list[0][1]["variables"]["query"] # First page + assert "SKIP 3 LIMIT 10" in call_args_list[1][1]["variables"]["query"] # Second page # Verify the result assert result == {"data": [ {"id": "1", "name": "entity1"}, {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"}, {"id": "3", "name": "entity3"} ]} @@ -215,7 +223,7 @@ def test_limit_and_skip_query_with_include_deleted_true(self, mock_execute_query # Verify the query was called with includeDeleted=True mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is True # Verify the result @@ -248,7 +256,7 @@ def test_limit_and_skip_query_with_include_deleted_false(self, mock_execute_quer # Verify the query was called with includeDeleted=False mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - variables = call_args[0][1] + variables = call_args[1]["variables"] assert variables["includeDeleted"] is False # Verify the result @@ -305,11 +313,11 @@ def test_limit_and_skip_query_complex_query(self, mock_execute_query): # Verify the query was called with the complex query mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - query = call_args[0][0] - assert "FIND aws_instance" in query - assert "THAT RELATES TO aws_vpc" in query - assert "WITH tag.Environment = 'production'" in query - assert "SKIP 0 LIMIT" in query # Should have skip/limit added + variables = call_args[1]["variables"] + assert "FIND aws_instance" in variables["query"] + assert "THAT RELATES TO aws_vpc" in variables["query"] + assert "WITH tag.Environment = 'production'" in variables["query"] + assert "SKIP 0 LIMIT" in variables["query"] # Should have skip/limit added # Verify the result assert result == {"data": [ @@ -325,7 +333,9 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): "queryV1": { "data": [ {"id": "1", "name": "entity1"}, - {"id": "2", "name": "entity2"} + {"id": "2", "name": "entity2"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"} ] } } @@ -337,7 +347,9 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): "queryV1": { "data": [ {"id": "3", "name": "entity3"}, - {"id": "4", "name": "entity4"} + {"id": "4", "name": "entity4"}, + {"id": "5", "name": "entity5"}, + {"id": "6", "name": "entity6"} ] } } @@ -356,16 +368,16 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): mock_execute_query.side_effect = [mock_response1, mock_response2, mock_response3] - result = self.client._limit_and_skip_query("FIND * LIMIT 10") + result = self.client._limit_and_skip_query("FIND * LIMIT 10", skip=3) # Should call three times assert mock_execute_query.call_count == 3 # Verify the pagination math call_args_list = mock_execute_query.call_args_list - assert "SKIP 0 LIMIT" in call_args_list[0][0][0] # Page 0: SKIP 0 - assert "SKIP 100 LIMIT" in call_args_list[1][0][0] # Page 1: SKIP 100 - assert "SKIP 200 LIMIT" in call_args_list[2][0][0] # Page 2: SKIP 200 + assert "SKIP 0 LIMIT" in call_args_list[0][1]["variables"]["query"] # Page 0: SKIP 0 + assert "SKIP 3 LIMIT" in call_args_list[1][1]["variables"]["query"] # Page 1: SKIP 3 + assert "SKIP 6 LIMIT" in call_args_list[2][1]["variables"]["query"] # Page 2: SKIP 6 # Verify the result assert result == {"data": [ @@ -373,6 +385,10 @@ def test_limit_and_skip_query_pagination_math(self, mock_execute_query): {"id": "2", "name": "entity2"}, {"id": "3", "name": "entity3"}, {"id": "4", "name": "entity4"}, + {"id": "3", "name": "entity3"}, + {"id": "4", "name": "entity4"}, + {"id": "5", "name": "entity5"}, + {"id": "6", "name": "entity6"}, {"id": "5", "name": "entity5"} ]} @@ -396,8 +412,8 @@ def test_limit_and_skip_query_default_constants(self, mock_execute_query): # Verify the query was called with default constants mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - query = call_args[0][0] - assert f"SKIP 0 LIMIT {J1QL_LIMIT_COUNT}" in query + variables = call_args[1]["variables"] + assert f"SKIP 0 LIMIT {J1QL_LIMIT_COUNT}" in variables["query"] # Verify the result assert result == {"data": [ From f6af5643c62dfb1af6acb2a56ea8db4d8acd8677 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 21:18:21 -0600 Subject: [PATCH 08/11] add type hints for all methods - helps with documentation of inputs and IDE auto-complete --- jupiterone/__init__.py | 10 +- jupiterone/client.py | 256 +++++++++++++++++++++++++++++------------ jupiterone/errors.py | 5 +- 3 files changed, 197 insertions(+), 74 deletions(-) diff --git a/jupiterone/__init__.py b/jupiterone/__init__.py index 0d227bb..4199fb5 100644 --- a/jupiterone/__init__.py +++ b/jupiterone/__init__.py @@ -1,5 +1,13 @@ from .client import JupiterOneClient from .errors import ( JupiterOneClientError, - JupiterOneApiError + JupiterOneApiError, + JupiterOneApiRetryError ) + +__all__ = [ + "JupiterOneClient", + "JupiterOneClientError", + "JupiterOneApiError", + "JupiterOneApiRetryError" +] diff --git a/jupiterone/client.py b/jupiterone/client.py index 2d74955..ceb0bbf 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -2,11 +2,12 @@ import json import os from warnings import warn -from typing import Dict, List, Union, Optional +from typing import Dict, List, Union, Optional, Any, Callable, Iterator, Tuple, Set from datetime import datetime import time import re import requests +import urllib.parse from requests.adapters import HTTPAdapter, Retry import concurrent.futures @@ -67,28 +68,30 @@ class JupiterOneClient: # pylint: disable=too-many-instance-attributes - DEFAULT_URL = "https://graphql.us.jupiterone.io" - SYNC_API_URL = "https://api.us.jupiterone.io" + DEFAULT_URL: str = "https://graphql.us.jupiterone.io" + SYNC_API_URL: str = "https://api.us.jupiterone.io" def __init__( self, - account: str = None, - token: str = None, + account: Optional[str] = None, + token: Optional[str] = None, url: str = DEFAULT_URL, sync_url: str = SYNC_API_URL, - ): - self.account = account - self.token = token - self.graphql_url = url - self.sync_url = sync_url - self.headers = { - "Authorization": "Bearer {}".format(self.token), - "JupiterOne-Account": self.account, + ) -> None: + # Validate inputs + self._validate_constructor_inputs(account, token, url, sync_url) + self.account: Optional[str] = account + self.token: Optional[str] = token + self.graphql_url: str = url + self.sync_url: str = sync_url + self.headers: Dict[str, str] = { + "Authorization": "Bearer {}".format(self.token or ""), + "JupiterOne-Account": self.account or "", "Content-Type": "application/json", } # Initialize session with retry logic - self.session = requests.Session() + self.session: requests.Session = requests.Session() retries = Retry( total=5, backoff_factor=1, @@ -97,40 +100,129 @@ def __init__( ) self.session.mount("https://", HTTPAdapter(max_retries=retries)) + def _validate_constructor_inputs( + self, + account: Optional[str], + token: Optional[str], + url: str, + sync_url: str + ) -> None: + """Validate constructor inputs""" + # Validate account + if account is not None: + if not isinstance(account, str): + raise JupiterOneClientError("Account must be a string") + if not account.strip(): + raise JupiterOneClientError("Account cannot be empty") + if len(account) < 3: + raise JupiterOneClientError("Account ID appears to be too short") + + # Validate token + if token is not None: + if not isinstance(token, str): + raise JupiterOneClientError("Token must be a string") + if not token.strip(): + raise JupiterOneClientError("Token cannot be empty") + if len(token) < 10: + raise JupiterOneClientError("Token appears to be too short") + + # Validate URLs + self._validate_url(url, "GraphQL URL") + self._validate_url(sync_url, "Sync API URL") + + def _validate_url(self, url: str, url_name: str) -> None: + """Validate URL format""" + if not isinstance(url, str): + raise JupiterOneClientError(f"{url_name} must be a string") + + if not url.strip(): + raise JupiterOneClientError(f"{url_name} cannot be empty") + + try: + parsed = urllib.parse.urlparse(url) + if not parsed.scheme or not parsed.netloc: + raise ValueError("Invalid URL format") + if parsed.scheme not in ['http', 'https']: + raise ValueError("URL must use http or https protocol") + except Exception as e: + raise JupiterOneClientError(f"Invalid {url_name}: {str(e)}") + + def _validate_entity_id(self, entity_id: str, param_name: str = "entity_id") -> None: + """Validate entity ID format""" + if not isinstance(entity_id, str): + raise JupiterOneClientError(f"{param_name} must be a string") + + if not entity_id.strip(): + raise JupiterOneClientError(f"{param_name} cannot be empty") + + if len(entity_id) < 10: + raise JupiterOneClientError(f"{param_name} appears to be too short") + + def _validate_query_string(self, query: str, param_name: str = "query") -> None: + """Validate J1QL query string""" + if not isinstance(query, str): + raise JupiterOneClientError(f"{param_name} must be a string") + + if not query.strip(): + raise JupiterOneClientError(f"{param_name} cannot be empty") + + # Basic J1QL validation + query_upper = query.upper().strip() + if not query_upper.startswith(('FIND', 'MATCH', 'RETURN')): + raise JupiterOneClientError(f"{param_name} must be a valid J1QL query starting with FIND, MATCH, or RETURN") + + def _validate_properties(self, properties: Dict[str, Any], param_name: str = "properties") -> None: + """Validate entity/relationship properties""" + if not isinstance(properties, dict): + raise JupiterOneClientError(f"{param_name} must be a dictionary") + + # Check for nested objects (not supported by JupiterOne API) + for key, value in properties.items(): + if isinstance(value, dict): + raise JupiterOneClientError( + f"Nested objects in {param_name} are not supported by JupiterOne API. " + f"Key '{key}' contains a nested dictionary. Please flatten the structure." + ) + if isinstance(value, list) and any(isinstance(item, dict) for item in value): + raise JupiterOneClientError( + f"Lists containing dictionaries in {param_name} are not supported by JupiterOne API. " + f"Key '{key}' contains a list with dictionaries. Please flatten the structure." + ) + @property - def account(self): + def account(self) -> Optional[str]: """Your JupiterOne account ID""" return self._account @account.setter - def account(self, value: str): + def account(self, value: Optional[str]) -> None: """Your JupiterOne account ID""" if not value: raise JupiterOneClientError("account is required") self._account = value @property - def token(self): + def token(self) -> Optional[str]: """Your JupiterOne access token""" return self._token @token.setter - def token(self, value: str): + def token(self, value: Optional[str]) -> None: """Your JupiterOne access token""" if not value: raise JupiterOneClientError("token is required") self._token = value # pylint: disable=R1710 - def _execute_query(self, query: str, variables: Dict = None) -> Dict: + def _execute_query(self, query: str, variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Executes query against graphql endpoint""" - data = {"query": query} + data: Dict[str, Any] = {"query": query} if variables: - data.update(variables=variables) + data["variables"] = variables # Always ask for variableResultSize - data.update(flags={"variableResultSize": True}) + data["flags"] = {"variableResultSize": True} response = self.session.post( self.graphql_url, @@ -179,10 +271,10 @@ def _execute_query(self, query: str, variables: Dict = None) -> Dict: def _cursor_query( self, query: str, - cursor: str = None, + cursor: Optional[str] = None, include_deleted: bool = False, max_workers: Optional[int] = None - ) -> Dict: + ) -> Dict[str, Any]: """Performs a V1 graph query using cursor pagination with optional parallel processing args: @@ -201,9 +293,9 @@ def _cursor_query( else: result_limit = False - results: List = [] + results: List[Dict[str, Any]] = [] - def fetch_page(cursor: Optional[str] = None) -> Dict: + def fetch_page(cursor: Optional[str] = None) -> Dict[str, Any]: variables = {"query": query, "includeDeleted": include_deleted} if cursor is not None: variables["cursor"] = cursor @@ -287,8 +379,8 @@ def _limit_and_skip_query( skip: int = J1QL_SKIP_COUNT, limit: int = J1QL_LIMIT_COUNT, include_deleted: bool = False, - ) -> Dict: - results: List = [] + ) -> Dict[str, Any]: + results: List[Dict[str, Any]] = [] page: int = 0 while True: @@ -313,7 +405,7 @@ def _limit_and_skip_query( return {"data": results} - def query_with_deferred_response(self, query, cursor=None): + def query_with_deferred_response(self, query: str, cursor: Optional[str] = None) -> List[Dict[str, Any]]: """ Execute a J1QL query that returns a deferred response for handling large result sets. @@ -389,7 +481,7 @@ def query_with_deferred_response(self, query, cursor=None): return all_query_results - def _execute_syncapi_request(self, endpoint: str, payload: Dict = None) -> Dict: + def _execute_syncapi_request(self, endpoint: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Executes POST request to SyncAPI endpoints""" # initiate requests session and implement retry logic of 5 request retries with 1 second between retries @@ -412,6 +504,7 @@ def _execute_syncapi_request(self, endpoint: str, payload: Dict = None) -> Dict: ) raise JupiterOneApiError(content.get("errors")) return response.json() + return {} elif response.status_code == 401: raise JupiterOneApiError( @@ -436,7 +529,7 @@ def _execute_syncapi_request(self, endpoint: str, payload: Dict = None) -> Dict: content = data.get("error", data.get("errors", content)) raise JupiterOneApiError("{}:{}".format(response.status_code, content)) - def query_v1(self, query: str, **kwargs) -> Dict: + def query_v1(self, query: str, **kwargs: Any) -> Dict[str, Any]: """Performs a V1 graph query args: query (str): Query text @@ -445,6 +538,25 @@ def query_v1(self, query: str, **kwargs) -> Dict: cursor (str): A pagination cursor for the initial query include_deleted (bool): Include recently deleted entities in query/search """ + # Validate inputs + self._validate_query_string(query) + + # Validate kwargs + if 'skip' in kwargs and kwargs['skip'] is not None: + if not isinstance(kwargs['skip'], int) or kwargs['skip'] < 0: + raise JupiterOneClientError("skip must be a non-negative integer") + + if 'limit' in kwargs and kwargs['limit'] is not None: + if not isinstance(kwargs['limit'], int) or kwargs['limit'] <= 0: + raise JupiterOneClientError("limit must be a positive integer") + + if 'cursor' in kwargs and kwargs['cursor'] is not None: + if not isinstance(kwargs['cursor'], str) or not kwargs['cursor'].strip(): + raise JupiterOneClientError("cursor must be a non-empty string") + + if 'include_deleted' in kwargs and kwargs['include_deleted'] is not None: + if not isinstance(kwargs['include_deleted'], bool): + raise JupiterOneClientError("include_deleted must be a boolean") 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) @@ -467,7 +579,7 @@ def query_v1(self, query: str, **kwargs) -> Dict: query=query, cursor=cursor, include_deleted=include_deleted ) - def create_entity(self, **kwargs) -> Dict: + def create_entity(self, **kwargs: Any) -> Dict[str, Any]: """Creates an entity in graph. It will also update an existing entity. args: @@ -494,7 +606,7 @@ def create_entity(self, **kwargs) -> Dict: response = self._execute_query(query=CREATE_ENTITY, variables=variables) return response["data"]["createEntity"] - def delete_entity(self, entity_id: str = None, timestamp: int = None, hard_delete: bool = True) -> Dict: + def delete_entity(self, entity_id: Optional[str] = None, timestamp: Optional[int] = None, hard_delete: bool = True) -> Dict[str, Any]: """Deletes an entity from the graph. args: @@ -502,13 +614,13 @@ def delete_entity(self, entity_id: str = None, timestamp: int = None, hard_delet timestamp (int, optional): Timestamp for the deletion. Defaults to None. hard_delete (bool): Whether to perform a hard delete. Defaults to True. """ - variables = {"entityId": entity_id, "hardDelete": hard_delete} + variables: Dict[str, Any] = {"entityId": entity_id, "hardDelete": hard_delete} if timestamp: variables["timestamp"] = timestamp response = self._execute_query(DELETE_ENTITY, variables=variables) return response["data"]["deleteEntityV2"] - def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: + def update_entity(self, entity_id: Optional[str] = None, properties: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """ Update an existing entity. @@ -520,7 +632,7 @@ def update_entity(self, entity_id: str = None, properties: Dict = None) -> Dict: response = self._execute_query(UPDATE_ENTITY, variables=variables) return response["data"]["updateEntity"] - def create_relationship(self, **kwargs) -> Dict: + def create_relationship(self, **kwargs: Any) -> Dict[str, Any]: """ Create a relationship (edge) between two entities (vertices). @@ -546,7 +658,7 @@ def create_relationship(self, **kwargs) -> Dict: response = self._execute_query(query=CREATE_RELATIONSHIP, variables=variables) return response["data"]["createRelationship"] - def update_relationship(self, **kwargs) -> Dict: + def update_relationship(self, **kwargs: Any) -> Dict[str, Any]: """ Update a relationship (edge) between two entities (vertices). @@ -568,7 +680,7 @@ def update_relationship(self, **kwargs) -> Dict: response = self._execute_query(query=UPDATE_RELATIONSHIP, variables=variables) return response["data"]["updateRelationship"] - def delete_relationship(self, relationship_id: str = None): + def delete_relationship(self, relationship_id: Optional[str] = None) -> Dict[str, Any]: """Deletes a relationship between two entities. args: @@ -581,11 +693,11 @@ def delete_relationship(self, relationship_id: str = None): def create_integration_instance( self, - instance_name: str = None, - instance_description: str = None, + instance_name: Optional[str] = None, + instance_description: Optional[str] = None, integration_definition_id: str = "8013680b-311a-4c2e-b53b-c8735fd97a5c", - resource_group_id: str = None, - ): + resource_group_id: Optional[str] = None, + ) -> Dict[str, Any]: """Creates a new Custom Integration Instance. args: @@ -614,7 +726,7 @@ def create_integration_instance( response = self._execute_query(CREATE_INSTANCE, variables=variables) return response["data"]["createIntegrationInstance"] - def fetch_all_entity_properties(self): + def fetch_all_entity_properties(self) -> List[Dict[str, Any]]: """Fetch list of aggregated property keys from all entities in the graph.""" response = self._execute_query(query=ALL_PROPERTIES) @@ -629,7 +741,7 @@ def fetch_all_entity_properties(self): return return_list - def fetch_all_entity_tags(self): + def fetch_all_entity_tags(self) -> List[Dict[str, Any]]: """Fetch list of aggregated property keys from all entities in the graph.""" response = self._execute_query(query=ALL_PROPERTIES) @@ -644,7 +756,7 @@ def fetch_all_entity_tags(self): return return_list - def fetch_entity_raw_data(self, entity_id: str = None): + def fetch_entity_raw_data(self, entity_id: Optional[str] = None) -> Dict[str, Any]: """Fetch the contents of raw data for a given entity in a J1 Account.""" variables = {"entityId": entity_id, "source": "integration-managed"} @@ -654,10 +766,10 @@ def fetch_entity_raw_data(self, entity_id: str = None): def start_sync_job( self, - instance_id: str = None, - sync_mode: str = None, - source: str = None, - ): + instance_id: Optional[str] = None, + sync_mode: Optional[str] = None, + source: Optional[str] = None, + ) -> Dict[str, Any]: """Start a synchronization job. args: @@ -754,7 +866,7 @@ def bulk_delete_entities( return response - def finalize_sync_job(self, instance_job_id: str = None): + def finalize_sync_job(self, instance_job_id: Optional[str] = None) -> Dict[str, Any]: """Start a synchronization job. args: @@ -768,7 +880,7 @@ def finalize_sync_job(self, instance_job_id: str = None): return response - def fetch_integration_jobs(self, instance_id: str = None): + def fetch_integration_jobs(self, instance_id: Optional[str] = None) -> List[Dict[str, Any]]: """Fetch Integration Job details from defined integration instance. args: @@ -801,21 +913,21 @@ def fetch_integration_job_events( return response["data"]["integrationEvents"] - def get_integration_definition_details(self, integration_type: str = None): + def get_integration_definition_details(self, integration_type: Optional[str] = None) -> Dict[str, Any]: """Fetch the Integration Definition Details for a given integration type.""" variables = {"integrationType": integration_type, "includeConfig": True} response = self._execute_query(FIND_INTEGRATION_DEFINITION, variables=variables) return response - def fetch_integration_instances(self, definition_id: str = None): + def fetch_integration_instances(self, definition_id: Optional[str] = None) -> List[Dict[str, Any]]: """Fetch all configured Instances for a given integration type.""" variables = {"definitionId": definition_id, "limit": 100} response = self._execute_query(INTEGRATION_INSTANCES, variables=variables) return response - def get_integration_instance_details(self, instance_id: str = None): + def get_integration_instance_details(self, instance_id: Optional[str] = None) -> Dict[str, Any]: """Fetch configuration details for a single configured Integration Instance.""" variables = {"integrationInstanceId": instance_id} @@ -964,7 +1076,7 @@ def generate_j1ql(self, natural_language_prompt: str = None): return response["data"]["j1qlFromNaturalLanguage"] - def list_alert_rules(self): + def list_alert_rules(self) -> List[Dict[str, Any]]: """List all defined Alert Rules configured in J1 account""" results = [] @@ -993,7 +1105,7 @@ def list_alert_rules(self): return results - def get_alert_rule_details(self, rule_id: str = None): + def get_alert_rule_details(self, rule_id: Optional[str] = None) -> Dict[str, Any]: """Get details of a single defined Alert Rule configured in J1 account""" results = [] @@ -1100,7 +1212,7 @@ def create_alert_rule( return response["data"]["createInlineQuestionRuleInstance"] - def delete_alert_rule(self, rule_id: str = None): + def delete_alert_rule(self, rule_id: Optional[str] = None) -> Dict[str, Any]: """Delete a single Alert Rule configured in J1 account""" variables = { "id": rule_id @@ -1320,7 +1432,7 @@ def fetch_downloaded_evaluation_results(self, download_url: str = None): return e - def list_questions(self, search_query: str = None, tags: List[str] = None): + def list_questions(self, search_query: Optional[str] = None, tags: Optional[List[str]] = None) -> Dict[str, Any]: """List all defined Questions configured in J1 Account Questions Library Args: @@ -1392,7 +1504,7 @@ def list_questions(self, search_query: str = None, tags: List[str] = None): return results - def get_question_details(self, question_id: str = None): + def get_question_details(self, question_id: Optional[str] = None) -> Dict[str, Any]: """Get details of a specific question by ID Args: @@ -1422,9 +1534,9 @@ def get_question_details(self, question_id: str = None): def create_question( self, title: str, - queries: List[Dict], - **kwargs - ): + queries: List[Dict[str, Any]], + **kwargs: Any + ) -> Dict[str, Any]: """Creates a new Question in the J1 account. Args: @@ -1510,12 +1622,12 @@ def create_question( def update_question( self, question_id: str, - title: str = None, - description: str = None, - queries: List[Dict] = None, - tags: List[str] = None, - **kwargs - ) -> Dict: + title: Optional[str] = None, + description: Optional[str] = None, + queries: Optional[List[Dict[str, Any]]] = None, + tags: Optional[List[str]] = None, + **kwargs: Any + ) -> Dict[str, Any]: """ Update an existing question in the J1 account. @@ -1629,7 +1741,7 @@ def update_question( response = self._execute_query(UPDATE_QUESTION, variables) return response["data"]["updateQuestion"] - def delete_question(self, question_id: str) -> Dict: + def delete_question(self, question_id: str) -> Dict[str, Any]: """ Delete an existing question from the J1 account. @@ -1729,7 +1841,7 @@ def create_update_parameter( response = self._execute_query(UPSERT_PARAMETER, variables=variables) return response - def update_entity_v2(self, entity_id: str = None, properties: Dict = None) -> Dict: + def update_entity_v2(self, entity_id: Optional[str] = None, properties: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """ Update an existing entity by adding new or updating existing properties. @@ -1746,7 +1858,7 @@ def update_entity_v2(self, entity_id: str = None, properties: Dict = None) -> Di response = self._execute_query(UPDATE_ENTITYV2, variables=variables) return response["data"]["updateEntityV2"] - def get_cft_upload_url(self, integration_instance_id: str, filename: str, dataset_id: str) -> Dict: + def get_cft_upload_url(self, integration_instance_id: str, filename: str, dataset_id: str) -> Dict[str, Any]: """ Get an upload URL for Custom File Transfer integration. @@ -1792,7 +1904,7 @@ def get_cft_upload_url(self, integration_instance_id: str, filename: str, datase response = self._execute_query(query, variables) return response["data"]["integrationFileTransferUploadUrl"] - def upload_cft_file(self, upload_url: str, file_path: str) -> Dict: + def upload_cft_file(self, upload_url: str, file_path: str) -> Dict[str, Any]: """ Upload a CSV file to the Custom File Transfer integration using the provided upload URL. diff --git a/jupiterone/errors.py b/jupiterone/errors.py index eccf017..018a5ef 100644 --- a/jupiterone/errors.py +++ b/jupiterone/errors.py @@ -1,11 +1,14 @@ class JupiterOneClientError(Exception): - """ Raised when error creating client """ + """ Raised when error creating client """ + pass class JupiterOneApiRetryError(Exception): """ Used to trigger retry on rate limit """ + pass class JupiterOneApiError(Exception): """ Raised when API returns error response """ + pass From 6ee5741ab30a617cd72c616b362c5f9ea52bc2c6 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 21:43:54 -0600 Subject: [PATCH 09/11] implement improved input validation --- jupiterone/client.py | 96 ++++++++++++++++++++++++++++++- tests/test_client_init.py | 4 +- tests/test_create_entity.py | 2 +- tests/test_create_relationship.py | 6 +- tests/test_delete_entity.py | 8 +-- tests/test_delete_relationship.py | 2 +- tests/test_query.py | 20 +++---- tests/test_update_entity.py | 2 +- 8 files changed, 116 insertions(+), 24 deletions(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index ceb0bbf..4727ca8 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -168,8 +168,8 @@ def _validate_query_string(self, query: str, param_name: str = "query") -> None: # Basic J1QL validation query_upper = query.upper().strip() - if not query_upper.startswith(('FIND', 'MATCH', 'RETURN')): - raise JupiterOneClientError(f"{param_name} must be a valid J1QL query starting with FIND, MATCH, or RETURN") + if not query_upper.startswith('FIND'): + raise JupiterOneClientError(f"{param_name} must be a valid J1QL query starting with FIND (case-insensitive)") def _validate_properties(self, properties: Dict[str, Any], param_name: str = "properties") -> None: """Validate entity/relationship properties""" @@ -216,6 +216,11 @@ def token(self, value: Optional[str]) -> None: # pylint: disable=R1710 def _execute_query(self, query: str, variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Executes query against graphql endpoint""" + # Validate credentials before making API calls + if not self.account: + raise JupiterOneClientError("Account is required. Please set the account property.") + if not self.token: + raise JupiterOneClientError("Token is required. Please set the token property.") data: Dict[str, Any] = {"query": query} if variables: @@ -483,6 +488,11 @@ def query_with_deferred_response(self, query: str, cursor: Optional[str] = None) def _execute_syncapi_request(self, endpoint: str, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Executes POST request to SyncAPI endpoints""" + # Validate credentials before making API calls + if not self.account: + raise JupiterOneClientError("Account is required. Please set the account property.") + if not self.token: + raise JupiterOneClientError("Token is required. Please set the token property.") # initiate requests session and implement retry logic of 5 request retries with 1 second between retries response = self.session.post( @@ -589,6 +599,34 @@ def create_entity(self, **kwargs: Any) -> Dict[str, Any]: timestamp (int): Specify createdOn timestamp properties (dict): Dictionary of key/value entity properties """ + # Validate required parameters + entity_key = kwargs.get("entity_key") + entity_type = kwargs.get("entity_type") + entity_class = kwargs.get("entity_class") + + if not entity_key: + raise JupiterOneClientError("entity_key is required") + if not isinstance(entity_key, str) or not entity_key.strip(): + raise JupiterOneClientError("entity_key must be a non-empty string") + + if not entity_type: + raise JupiterOneClientError("entity_type is required") + if not isinstance(entity_type, str) or not entity_type.strip(): + raise JupiterOneClientError("entity_type must be a non-empty string") + + if not entity_class: + raise JupiterOneClientError("entity_class is required") + if not isinstance(entity_class, str) or not entity_class.strip(): + raise JupiterOneClientError("entity_class must be a non-empty string") + + # Validate properties if provided + if "properties" in kwargs and kwargs["properties"] is not None: + self._validate_properties(kwargs["properties"]) + + # Validate timestamp if provided + if "timestamp" in kwargs and kwargs["timestamp"] is not None: + if not isinstance(kwargs["timestamp"], int) or kwargs["timestamp"] <= 0: + raise JupiterOneClientError("timestamp must be a positive integer") variables = { "entityKey": kwargs.pop("entity_key"), "entityType": kwargs.pop("entity_type"), @@ -614,6 +652,19 @@ def delete_entity(self, entity_id: Optional[str] = None, timestamp: Optional[int timestamp (int, optional): Timestamp for the deletion. Defaults to None. hard_delete (bool): Whether to perform a hard delete. Defaults to True. """ + # Validate required parameters + if not entity_id: + raise JupiterOneClientError("entity_id is required") + self._validate_entity_id(entity_id) + + # Validate timestamp if provided + if timestamp is not None: + if not isinstance(timestamp, int) or timestamp <= 0: + raise JupiterOneClientError("timestamp must be a positive integer") + + # Validate hard_delete + if not isinstance(hard_delete, bool): + raise JupiterOneClientError("hard_delete must be a boolean") variables: Dict[str, Any] = {"entityId": entity_id, "hardDelete": hard_delete} if timestamp: variables["timestamp"] = timestamp @@ -628,6 +679,14 @@ def update_entity(self, entity_id: Optional[str] = None, properties: Optional[Di entity_id (str): The _id of the entity to update properties (dict): Dictionary of key/value entity properties """ + # Validate required parameters + if not entity_id: + raise JupiterOneClientError("entity_id is required") + self._validate_entity_id(entity_id) + + if not properties: + raise JupiterOneClientError("properties is required") + self._validate_properties(properties) variables = {"entityId": entity_id, "properties": properties} response = self._execute_query(UPDATE_ENTITY, variables=variables) return response["data"]["updateEntity"] @@ -643,6 +702,39 @@ def create_relationship(self, **kwargs: Any) -> Dict[str, Any]: from_entity_id (str): Entity ID of the source vertex to_entity_id (str): Entity ID of the destination vertex """ + # Validate required parameters + relationship_key = kwargs.get("relationship_key") + relationship_type = kwargs.get("relationship_type") + relationship_class = kwargs.get("relationship_class") + from_entity_id = kwargs.get("from_entity_id") + to_entity_id = kwargs.get("to_entity_id") + + if not relationship_key: + raise JupiterOneClientError("relationship_key is required") + if not isinstance(relationship_key, str) or not relationship_key.strip(): + raise JupiterOneClientError("relationship_key must be a non-empty string") + + if not relationship_type: + raise JupiterOneClientError("relationship_type is required") + if not isinstance(relationship_type, str) or not relationship_type.strip(): + raise JupiterOneClientError("relationship_type must be a non-empty string") + + if not relationship_class: + raise JupiterOneClientError("relationship_class is required") + if not isinstance(relationship_class, str) or not relationship_class.strip(): + raise JupiterOneClientError("relationship_class must be a non-empty string") + + if not from_entity_id: + raise JupiterOneClientError("from_entity_id is required") + self._validate_entity_id(from_entity_id, "from_entity_id") + + if not to_entity_id: + raise JupiterOneClientError("to_entity_id is required") + self._validate_entity_id(to_entity_id, "to_entity_id") + + # Validate properties if provided + if "properties" in kwargs and kwargs["properties"] is not None: + self._validate_properties(kwargs["properties"]) variables = { "relationshipKey": kwargs.pop("relationship_key"), "relationshipType": kwargs.pop("relationship_type"), diff --git a/tests/test_client_init.py b/tests/test_client_init.py index dad689d..39811e2 100644 --- a/tests/test_client_init.py +++ b/tests/test_client_init.py @@ -45,12 +45,12 @@ def test_client_init_missing_token(self): def test_client_init_empty_account(self): """Test client initialization with empty account""" - with pytest.raises(JupiterOneClientError, match="account is required"): + with pytest.raises(JupiterOneClientError, match="Account cannot be empty"): JupiterOneClient(account="", token="test-token") def test_client_init_empty_token(self): """Test client initialization with empty token""" - with pytest.raises(JupiterOneClientError, match="token is required"): + with pytest.raises(JupiterOneClientError, match="Token cannot be empty"): JupiterOneClient(account="test-account", token="") def test_client_init_none_account(self): diff --git a/tests/test_create_entity.py b/tests/test_create_entity.py index b31b316..cc18fdf 100644 --- a/tests/test_create_entity.py +++ b/tests/test_create_entity.py @@ -37,7 +37,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.create_entity( entity_key='host1', entity_type='test_host', diff --git a/tests/test_create_relationship.py b/tests/test_create_relationship.py index f63d4b5..c53bb65 100644 --- a/tests/test_create_relationship.py +++ b/tests/test_create_relationship.py @@ -41,13 +41,13 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.create_relationship( relationship_key='relationship1', relationship_type='test_relationship', relationship_class='TestRelationship', - from_entity_id='2', - to_entity_id='1' + from_entity_id='entity-id-12345', + to_entity_id='entity-id-67890' ) assert type(response) == dict diff --git a/tests/test_delete_entity.py b/tests/test_delete_entity.py index 16c123d..f232ddf 100644 --- a/tests/test_delete_entity.py +++ b/tests/test_delete_entity.py @@ -32,7 +32,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('1') assert type(response) == dict @@ -68,7 +68,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('2', timestamp=1640995200000) assert type(response) == dict @@ -103,7 +103,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('3', hard_delete=False) assert type(response) == dict @@ -138,7 +138,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_entity('4', timestamp=1640995200000, hard_delete=True) assert type(response) == dict diff --git a/tests/test_delete_relationship.py b/tests/test_delete_relationship.py index 42493b3..dba6d5d 100644 --- a/tests/test_delete_relationship.py +++ b/tests/test_delete_relationship.py @@ -41,7 +41,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.delete_relationship('1') assert type(response) == dict diff --git a/tests/test_query.py b/tests/test_query.py index 4ae28c4..890a37f 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -83,7 +83,7 @@ def test_execute_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" variables = { 'query': query, @@ -110,7 +110,7 @@ def test_limit_skip_query_v1(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( query=query, @@ -139,7 +139,7 @@ def test_cursor_query_v1(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( @@ -186,7 +186,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1' return tree" response = j1.query_v1( query=query, @@ -236,7 +236,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1' return tree" response = j1.query_v1(query) @@ -268,7 +268,7 @@ def test_retry_on_limit_skip_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( query=query, @@ -302,7 +302,7 @@ def test_retry_on_cursor_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" response = j1.query_v1( query=query @@ -322,7 +322,7 @@ def test_avoid_retry_on_limit_skip_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" with pytest.raises(JupiterOneApiError): j1.query_v1( @@ -340,7 +340,7 @@ def test_avoid_retry_on_cursor_query(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" with pytest.raises(JupiterOneApiError): j1.query_v1( @@ -356,7 +356,7 @@ def test_warn_limit_and_skip_deprecated(): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') query = "find Host with _id='1'" with pytest.warns(DeprecationWarning): diff --git a/tests/test_update_entity.py b/tests/test_update_entity.py index cbd24c3..37b0514 100644 --- a/tests/test_update_entity.py +++ b/tests/test_update_entity.py @@ -34,7 +34,7 @@ def request_callback(request): content_type='application/json', ) - j1 = JupiterOneClient(account='testAccount', token='testToken') + j1 = JupiterOneClient(account='testAccount', token='testToken1234567890') response = j1.update_entity('1', properties={'testKey': 'testValue'}) assert type(response) == dict From c86de28184e62813c680d09441ad1a944c281ce5 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 21:48:51 -0600 Subject: [PATCH 10/11] remove timestamp input from create_entity --- examples/02_entity_management.py | 3 +-- examples/examples.py | 6 ++---- jupiterone/client.py | 8 -------- 3 files changed, 3 insertions(+), 14 deletions(-) diff --git a/examples/02_entity_management.py b/examples/02_entity_management.py index 07b0b6f..c415934 100644 --- a/examples/02_entity_management.py +++ b/examples/02_entity_management.py @@ -79,8 +79,7 @@ def create_entity_examples(j1): 'createdBy': 'terraform', 'lastBackup': '2024-01-01T00:00:00Z', 'maintenanceWindow': 'sun:03:00-sun:04:00' - }, - timestamp=int(time.time()) * 1000 + } ) print(f"Created complex entity: {complex_entity['entity']['_id']}\n") diff --git a/examples/examples.py b/examples/examples.py index e692540..6a710f0 100644 --- a/examples/examples.py +++ b/examples/examples.py @@ -37,8 +37,7 @@ 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 + properties=properties ) print("create_entity()") print(create_r) @@ -69,8 +68,7 @@ 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 + properties=properties ) print("create_entity()") print(create_r_2) diff --git a/jupiterone/client.py b/jupiterone/client.py index 4727ca8..e5beee4 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -596,7 +596,6 @@ def create_entity(self, **kwargs: Any) -> Dict[str, Any]: entity_key (str): Unique key for the entity entity_type (str): Value for _type of entity entity_class (str): Value for _class of entity - timestamp (int): Specify createdOn timestamp properties (dict): Dictionary of key/value entity properties """ # Validate required parameters @@ -623,21 +622,14 @@ def create_entity(self, **kwargs: Any) -> Dict[str, Any]: if "properties" in kwargs and kwargs["properties"] is not None: self._validate_properties(kwargs["properties"]) - # Validate timestamp if provided - if "timestamp" in kwargs and kwargs["timestamp"] is not None: - if not isinstance(kwargs["timestamp"], int) or kwargs["timestamp"] <= 0: - raise JupiterOneClientError("timestamp must be a positive integer") variables = { "entityKey": kwargs.pop("entity_key"), "entityType": kwargs.pop("entity_type"), "entityClass": kwargs.pop("entity_class"), } - timestamp: int = kwargs.pop("timestamp", None) properties: Dict = kwargs.pop("properties", None) - if timestamp: - variables.update(timestamp=timestamp) if properties: variables.update(properties=properties) From b4d7f41d8d63e9e0fa99c9485df7a2a501e4ffd3 Mon Sep 17 00:00:00 2001 From: SeaBlooms Date: Thu, 16 Oct 2025 21:59:02 -0600 Subject: [PATCH 11/11] remove unused imports --- jupiterone/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jupiterone/client.py b/jupiterone/client.py index e5beee4..e5f9ae9 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -2,7 +2,7 @@ import json import os from warnings import warn -from typing import Dict, List, Union, Optional, Any, Callable, Iterator, Tuple, Set +from typing import Dict, List, Union, Optional, Any from datetime import datetime import time import re