Merge pull request #63 from JupiterOne/fix-example-scripts

fix example scripts, improve test suites, add type hints for method inputs, improve input validation

Author: SeaBlooms

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'])
 

examples/02_entity_management.py

@@ -76,13 +76,10 @@ 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'
-            }
-        },
-        timestamp=int(time.time()) * 1000
+            'createdBy': 'terraform',
+            'lastBackup': '2024-01-01T00:00:00Z',
+            'maintenanceWindow': 'sun:03:00-sun:04:00'
+        }
     )
     print(f"Created complex entity: {complex_entity['entity']['_id']}\n")
 
@@ -122,16 +119,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")

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")

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
 

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

examples/bulk_upload.py

@@ -5,9 +5,15 @@
 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 = "e7113c37-1ea8-4d00-9b82-c24952e70916"
+instance_id = "<Integration Instance ID>"
 
 sync_job = j1.start_sync_job(
             instance_id=instance_id,

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
@@ -31,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)
@@ -63,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)

jupiterone/__init__.py

@@ -1,5 +1,13 @@
 from .client import JupiterOneClient
 from .errors import (
     JupiterOneClientError,
-    JupiterOneApiError
+    JupiterOneApiError,
+    JupiterOneApiRetryError
 )
+
+__all__ = [
+    "JupiterOneClient",
+    "JupiterOneClientError", 
+    "JupiterOneApiError",
+    "JupiterOneApiRetryError"
+]