cartography 0.102.0rc2__py3-none-any.whl → 0.103.0__py3-none-any.whl

cartography/intel/okta/roles.py

@@ -26,7 +26,7 @@ def _get_user_roles(api_client: ApiClient, user_id: str, okta_org_id: str) -> st
     """
 
     # https://developer.okta.com/docs/reference/api/roles/#list-roles
-    response = api_client.get_path(f'/{user_id}/roles')
+    response = api_client.get_path(f"/{user_id}/roles")
     check_rate_limit(response)
     return response.text
 
@@ -42,7 +42,7 @@ def _get_group_roles(api_client: ApiClient, group_id: str, okta_org_id: str) ->
     """
 
     # https://developer.okta.com/docs/reference/api/roles/#list-roles-assigned-to-group
-    response = api_client.get_path(f'/{group_id}/roles')
+    response = api_client.get_path(f"/{group_id}/roles")
     check_rate_limit(response)
     return response.text
 
@@ -94,7 +94,12 @@ def transform_group_roles_data(data: str, okta_org_id: str) -> List[Dict]:
 
 
 @timeit
-def _load_user_role(neo4j_session: neo4j.Session, user_id: str, roles_data: List[Dict], okta_update_tag: int) -> None:
+def _load_user_role(
+    neo4j_session: neo4j.Session,
+    user_id: str,
+    roles_data: List[Dict],
+    okta_update_tag: int,
+) -> None:
     ingest = """
     MATCH (user:OktaUser{id: $USER_ID})<-[:RESOURCE]-(org:OktaOrganization)
     WITH user,org
@@ -122,7 +127,9 @@ def _load_user_role(neo4j_session: neo4j.Session, user_id: str, roles_data: List
 
 @timeit
 def _load_group_role(
-    neo4j_session: neo4j.Session, group_id: str, roles_data: List[Dict],
+    neo4j_session: neo4j.Session,
+    group_id: str,
+    roles_data: List[Dict],
     okta_update_tag: int,
 ) -> None:
     ingest = """
@@ -152,7 +159,10 @@ def _load_group_role(
 
 @timeit
 def sync_roles(
-    neo4j_session: str, okta_org_id: str, okta_update_tag: int, okta_api_key: str,
+    neo4j_session: str,
+    okta_org_id: str,
+    okta_update_tag: int,
+    okta_api_key: str,
     sync_state: OktaSyncState,
 ) -> None:
     """

cartography/intel/okta/users.py

@@ -12,7 +12,6 @@ from cartography.intel.okta.sync_state import OktaSyncState
 from cartography.intel.okta.utils import check_rate_limit
 from cartography.util import timeit
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -56,7 +55,9 @@ def _get_okta_users(user_client: UsersClient) -> List[Dict]:
 
 
 @timeit
-def transform_okta_user_list(okta_user_list: List[User]) -> Tuple[List[Dict], List[str]]:
+def transform_okta_user_list(
+    okta_user_list: List[User],
+) -> Tuple[List[Dict], List[str]]:
     users: List[Dict] = []
     user_ids: List[str] = []
 
@@ -91,7 +92,9 @@ def transform_okta_user(okta_user: User) -> Dict:
         user_props["activated"] = None
 
     if okta_user.statusChanged:
-        user_props["status_changed"] = okta_user.statusChanged.strftime("%m/%d/%Y, %H:%M:%S")
+        user_props["status_changed"] = okta_user.statusChanged.strftime(
+            "%m/%d/%Y, %H:%M:%S",
+        )
     else:
         user_props["status_changed"] = None
 
@@ -101,12 +104,16 @@ def transform_okta_user(okta_user: User) -> Dict:
         user_props["last_login"] = None
 
     if okta_user.lastUpdated:
-        user_props["okta_last_updated"] = okta_user.lastUpdated.strftime("%m/%d/%Y, %H:%M:%S")
+        user_props["okta_last_updated"] = okta_user.lastUpdated.strftime(
+            "%m/%d/%Y, %H:%M:%S",
+        )
     else:
         user_props["okta_last_updated"] = None
 
     if okta_user.passwordChanged:
-        user_props["password_changed"] = okta_user.passwordChanged.strftime("%m/%d/%Y, %H:%M:%S")
+        user_props["password_changed"] = okta_user.passwordChanged.strftime(
+            "%m/%d/%Y, %H:%M:%S",
+        )
     else:
         user_props["password_changed"] = None
 
@@ -120,7 +127,9 @@ def transform_okta_user(okta_user: User) -> Dict:
 
 @timeit
 def _load_okta_users(
-    neo4j_session: neo4j.Session, okta_org_id: str, user_list: List[Dict],
+    neo4j_session: neo4j.Session,
+    okta_org_id: str,
+    user_list: List[Dict],
     okta_update_tag: int,
 ) -> None:
     """
@@ -175,8 +184,11 @@ def _load_okta_users(
 
 @timeit
 def sync_okta_users(
-    neo4j_session: neo4j.Session, okta_org_id: str, okta_update_tag: int,
-    okta_api_key: str, sync_state: OktaSyncState,
+    neo4j_session: neo4j.Session,
+    okta_org_id: str,
+    okta_update_tag: int,
+    okta_api_key: str,
+    sync_state: OktaSyncState,
 ) -> None:
     """
     Sync okta users

cartography/intel/okta/utils.py

@@ -43,9 +43,9 @@ def check_rate_limit(response: Response) -> None:
     """
     rate_limit_threshold = 0.1
 
-    remaining = response.headers.get('x-rate-limit-remaining')
-    limit = response.headers.get('x-rate-limit-limit')
-    reset_time = response.headers.get('x-rate-limit-reset')
+    remaining = response.headers.get("x-rate-limit-remaining")
+    limit = response.headers.get("x-rate-limit-limit")
+    reset_time = response.headers.get("x-rate-limit-reset")
 
     if remaining and limit and reset_time:
         if (int(remaining) / int(limit)) < rate_limit_threshold:
@@ -58,5 +58,7 @@ def check_rate_limit(response: Response) -> None:
                     f"Okta API limit exceeded. Sleep time of {sleep_time_seconds} would exceed one minute. Crashing "
                     f"Okta sync to avoid blocking.",
                 )
-            logger.warning(f"Okta rate limit threshold reached. Waiting {sleep_time_seconds} seconds.")
+            logger.warning(
+                f"Okta rate limit threshold reached. Waiting {sleep_time_seconds} seconds.",
+            )
             time.sleep(sleep_time_seconds)

cartography/intel/openai/__init__.py

@@ -0,0 +1,86 @@
+import logging
+
+import neo4j
+import requests
+
+import cartography.intel.openai.adminapikeys
+import cartography.intel.openai.apikeys
+import cartography.intel.openai.projects
+import cartography.intel.openai.serviceaccounts
+import cartography.intel.openai.users
+from cartography.config import Config
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def start_openai_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
+    """
+    If this module is configured, perform ingestion of OpenAI data. Otherwise warn and exit
+    :param neo4j_session: Neo4J session for database interface
+    :param config: A cartography.config object
+    :return: None
+    """
+
+    if not config.openai_apikey or not config.openai_org_id:
+        logger.info(
+            "OpenAI import is not configured - skipping this module. "
+            "See docs to configure.",
+        )
+        return
+
+    # Create requests sessions
+    api_session = requests.session()
+    api_session.headers.update(
+        {
+            "Authorization": f"Bearer {config.openai_apikey}",
+            "OpenAI-Organization": config.openai_org_id,
+        }
+    )
+
+    common_job_parameters = {
+        "UPDATE_TAG": config.update_tag,
+        "BASE_URL": "https://api.openai.com/v1",
+        "ORG_ID": config.openai_org_id,
+    }
+
+    # Organization node is created during the users sync
+    cartography.intel.openai.users.sync(
+        neo4j_session,
+        api_session,
+        common_job_parameters,
+        ORG_ID=config.openai_org_id,
+    )
+
+    for project in cartography.intel.openai.projects.sync(
+        neo4j_session,
+        api_session,
+        common_job_parameters,
+        ORG_ID=config.openai_org_id,
+    ):
+        project_job_parameters = {
+            "UPDATE_TAG": config.update_tag,
+            "BASE_URL": "https://api.openai.com/v1",
+            "ORG_ID": config.openai_org_id,
+            "project_id": project["id"],
+        }
+        cartography.intel.openai.serviceaccounts.sync(
+            neo4j_session,
+            api_session,
+            project_job_parameters,
+            project_id=project["id"],
+        )
+        cartography.intel.openai.apikeys.sync(
+            neo4j_session,
+            api_session,
+            project_job_parameters,
+            project_id=project["id"],
+        )
+
+    cartography.intel.openai.adminapikeys.sync(
+        neo4j_session,
+        api_session,
+        common_job_parameters,
+        ORG_ID=config.openai_org_id,
+    )

cartography/intel/openai/adminapikeys.py

@@ -0,0 +1,90 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.openai.util import paginated_get
+from cartography.models.openai.adminapikey import OpenAIAdminApiKeySchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: Dict[str, Any],
+    ORG_ID: str,
+) -> List[Dict]:
+    adminapikeys = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+    )
+    transformed_adminapikeys = transform(adminapikeys)
+    load_adminapikeys(
+        neo4j_session,
+        transformed_adminapikeys,
+        ORG_ID,
+        common_job_parameters["UPDATE_TAG"],
+    )
+    cleanup(neo4j_session, common_job_parameters)
+    return adminapikeys
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+) -> List[Dict[str, Any]]:
+    return list(
+        paginated_get(
+            api_session, f"{base_url}/organization/admin_api_keys", timeout=_TIMEOUT
+        )
+    )
+
+
+def transform(
+    adminapikeys: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    result: List[Dict[str, Any]] = []
+    for adminapikey in adminapikeys:
+        if adminapikey["owner"]["type"] == "user":
+            adminapikey["owner_user_id"] = adminapikey["owner"]["id"]
+        else:
+            adminapikey["owner_sa_id"] = adminapikey["owner"]["id"]
+        result.append(adminapikey)
+    return result
+
+
+@timeit
+def load_adminapikeys(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    ORG_ID: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d OpenAI AdminApiKey into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        OpenAIAdminApiKeySchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=ORG_ID,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(OpenAIAdminApiKeySchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/openai/apikeys.py

@@ -0,0 +1,96 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.openai.util import paginated_get
+from cartography.models.openai.apikey import OpenAIApiKeySchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: Dict[str, Any],
+    project_id: str,
+) -> None:
+    apikeys = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+        project_id,
+    )
+    transformed_apikeys = transform(apikeys)
+    load_apikeys(
+        neo4j_session,
+        transformed_apikeys,
+        project_id,
+        common_job_parameters["UPDATE_TAG"],
+    )
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+    project_id: str,
+) -> List[Dict[str, Any]]:
+    return list(
+        paginated_get(
+            api_session,
+            "{base_url}/organization/projects/{project_id}/api_keys".format(
+                base_url=base_url,
+                project_id=project_id,
+            ),
+            timeout=_TIMEOUT,
+        )
+    )
+
+
+def transform(
+    apikeys: List[Dict[str, Any]],
+) -> List[Dict[str, Any]]:
+    result: List[Dict[str, Any]] = []
+    for apikey in apikeys:
+        if apikey["owner"]["type"] == "user":
+            apikey["owner_user_id"] = apikey["owner"]["user"]["id"]
+        else:
+            apikey["owner_sa_id"] = apikey["owner"]["service_account"]["id"]
+        result.append(apikey)
+    return result
+
+
+@timeit
+def load_apikeys(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    project_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d OpenAI Project APIKey into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        OpenAIApiKeySchema(),
+        data,
+        lastupdated=update_tag,
+        project_id=project_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(OpenAIApiKeySchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/openai/projects.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.openai.util import paginated_get
+from cartography.models.openai.project import OpenAIProjectSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: Dict[str, Any],
+    ORG_ID: str,
+) -> List[Dict]:
+    projects = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+    )
+    for project in projects:
+        project["users"] = []
+        for user in get_project_users(
+            api_session,
+            common_job_parameters["BASE_URL"],
+            project["id"],
+        ):
+            project["users"].append(user["id"])
+    load_projects(neo4j_session, projects, ORG_ID, common_job_parameters["UPDATE_TAG"])
+    cleanup(neo4j_session, common_job_parameters)
+    return projects
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+) -> List[Dict[str, Any]]:
+    return list(
+        paginated_get(
+            api_session, f"{base_url}/organization/projects", timeout=_TIMEOUT
+        )
+    )
+
+
+@timeit
+def get_project_users(
+    api_session: requests.Session,
+    base_url: str,
+    project_id: str,
+) -> List[Dict[str, Any]]:
+    return list(
+        paginated_get(
+            api_session,
+            f"{base_url}/organization/projects/{project_id}/users",
+            timeout=_TIMEOUT,
+        )
+    )
+
+
+@timeit
+def load_projects(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    ORG_ID: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d OpenAIProjectSchema into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        OpenAIProjectSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=ORG_ID,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(OpenAIProjectSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/openai/serviceaccounts.py

@@ -0,0 +1,82 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.openai.util import paginated_get
+from cartography.models.openai.serviceaccount import OpenAIServiceAccountSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: Dict[str, Any],
+    project_id: str,
+) -> None:
+    serviceaccountss = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+        project_id,
+    )
+    load_serviceaccounts(
+        neo4j_session,
+        serviceaccountss,
+        project_id,
+        common_job_parameters["UPDATE_TAG"],
+    )
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+    project_id: str,
+) -> List[Dict[str, Any]]:
+    return list(
+        paginated_get(
+            api_session,
+            "{base_url}/organization/projects/{project_id}/service_accounts".format(
+                base_url=base_url,
+                project_id=project_id,
+            ),
+            timeout=_TIMEOUT,
+        )
+    )
+
+
+@timeit
+def load_serviceaccounts(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    project_id: str,
+    update_tag: int,
+) -> None:
+    logger.info("Loading %d OpenAI ProjectServiceAccount into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        OpenAIServiceAccountSchema(),
+        data,
+        lastupdated=update_tag,
+        project_id=project_id,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(OpenAIServiceAccountSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/openai/users.py

@@ -0,0 +1,78 @@
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+
+import neo4j
+import requests
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.openai.util import paginated_get
+from cartography.models.openai.organization import OpenAIOrganizationSchema
+from cartography.models.openai.user import OpenAIUserSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+# Connect and read timeouts of 60 seconds each; see https://requests.readthedocs.io/en/master/user/advanced/#timeouts
+_TIMEOUT = (60, 60)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    api_session: requests.Session,
+    common_job_parameters: Dict[str, Any],
+    ORG_ID: str,
+) -> None:
+    users = get(
+        api_session,
+        common_job_parameters["BASE_URL"],
+    )
+    load_users(neo4j_session, users, ORG_ID, common_job_parameters["UPDATE_TAG"])
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def get(
+    api_session: requests.Session,
+    base_url: str,
+) -> List[Dict[str, Any]]:
+    return list(
+        paginated_get(api_session, f"{base_url}/organization/users", timeout=_TIMEOUT)
+    )
+
+
+@timeit
+def load_users(
+    neo4j_session: neo4j.Session,
+    data: List[Dict[str, Any]],
+    ORG_ID: str,
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        OpenAIOrganizationSchema(),
+        [{"id": ORG_ID}],
+        lastupdated=update_tag,
+    )
+    logger.info("Loading %d OpenAI User into Neo4j.", len(data))
+    load(
+        neo4j_session,
+        OpenAIUserSchema(),
+        data,
+        lastupdated=update_tag,
+        ORG_ID=ORG_ID,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session, common_job_parameters: Dict[str, Any]
+) -> None:
+    GraphJob.from_node_schema(OpenAIOrganizationSchema(), common_job_parameters).run(
+        neo4j_session
+    )
+    GraphJob.from_node_schema(OpenAIUserSchema(), common_job_parameters).run(
+        neo4j_session
+    )