cartography 0.93.0rc1__py3-none-any.whl → 0.123.0__py3-none-any.whl

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 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,97 @@
+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"] = []
+        project["admins"] = []
+        for user in get_project_users(
+            api_session,
+            common_job_parameters["BASE_URL"],
+            project["id"],
+        ):
+            project["users"].append(user["id"])
+            if user["role"] == "owner":
+                project["admins"].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 OpenAI Projects 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 ServiceAccount 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,75 @@
+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(OpenAIUserSchema(), common_job_parameters).run(
+        neo4j_session
+    )

cartography/intel/openai/util.py

@@ -0,0 +1,45 @@
+from typing import Any
+from typing import Generator
+
+import requests
+
+
+def paginated_get(
+    api_session: requests.Session,
+    url: str,
+    timeout: tuple[int, int],
+    after: str | None = None,
+) -> Generator[dict[str, Any], None, None]:
+    """Helper function to get paginated data from the OpenAI API.
+
+    This function handles the pagination of the API requests and returns
+    the results as a generator. It will continue to make requests until
+    all pages of data have been retrieved. The results are returned as a
+    list of dictionaries, where each dictionary represents a single
+    entity.
+
+    Args:
+        api_session (requests.Session): The requests session to use for making API calls.
+        url (str): The URL to make the API call to.
+        timeout (tuple[int, int]): The timeout for the API call.
+        after (str | None): The ID of the last item retrieved in the previous request.
+            If None, the first page of results will be retrieved.
+    Returns:
+        Generator[dict[str, Any], None, None]: A generator yielding dictionaries representing the results.
+    """
+    params = {"after": after} if after else {}
+    req = api_session.get(
+        url,
+        params=params,
+        timeout=timeout,
+    )
+    req.raise_for_status()
+    result = req.json()
+    yield from result.get("data", [])
+    if result.get("has_more"):
+        yield from paginated_get(
+            api_session,
+            url,
+            timeout=timeout,
+            after=result.get("last_id"),
+        )

cartography/intel/pagerduty/__init__.py

@@ -4,9 +4,7 @@ import neo4j
 from pdpyras import APISession
 
 from cartography.config import Config
-from cartography.intel.pagerduty.escalation_policies import (
-    sync_escalation_policies,
-)
+from cartography.intel.pagerduty.escalation_policies import sync_escalation_policies
 from cartography.intel.pagerduty.schedules import sync_schedules
 from cartography.intel.pagerduty.services import sync_services
 from cartography.intel.pagerduty.teams import sync_teams
@@ -23,7 +21,8 @@ stat_handler = get_stats_client(__name__)
 
 @timeit
 def start_pagerduty_ingestion(
-    neo4j_session: neo4j.Session, config: Config,
+    neo4j_session: neo4j.Session,
+    config: Config,
 ) -> None:
     """
     Perform ingestion of pagerduty data.
@@ -35,7 +34,9 @@ def start_pagerduty_ingestion(
         "UPDATE_TAG": config.update_tag,
     }
     if not config.pagerduty_api_key:
-        logger.info('PagerDuty import is not configured - skipping this module. See docs to configure.')
+        logger.info(
+            "PagerDuty import is not configured - skipping this module. See docs to configure.",
+        )
         return
     session = APISession(config.pagerduty_api_key)
     if config.pagerduty_request_timeout is not None:
@@ -54,8 +55,8 @@ def start_pagerduty_ingestion(
 
     merge_module_sync_metadata(
         neo4j_session,
-        group_type='pagerduty',
-        group_id='module',
+        group_type="pagerduty",
+        group_id="module",
         synced_type="pagerduty",
         update_tag=config.update_tag,
         stat_handler=stat_handler,

cartography/intel/pagerduty/escalation_policies.py

@@ -6,6 +6,7 @@ from typing import List
 import neo4j
 from pdpyras import APISession
 
+from cartography.client.core.tx import run_write_query
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
@@ -31,7 +32,9 @@ def get_escalation_policies(pd_session: APISession) -> List[Dict[str, Any]]:
 
 
 def load_escalation_policy_data(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Transform and load escalation_policy information
@@ -70,7 +73,8 @@ def load_escalation_policy_data(
             for team in policy["teams"]:
                 teams.append({"escalation_policy": policy["id"], "team": team["id"]})
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         EscalationPolicies=data,
         update_tag=update_tag,
@@ -82,7 +86,9 @@ def load_escalation_policy_data(
 
 
 def _attach_rules(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add escalation policy rules, and attach them to targets.
@@ -111,7 +117,8 @@ def _attach_rules(
                 elif target["type"] == "schedule":
                     schedules.append({"rule": rule["id"], "schedule": target["id"]})
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Rules=data,
         update_tag=update_tag,
@@ -122,7 +129,9 @@ def _attach_rules(
 
 
 def _attach_user_targets(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add relationship between escalation policy and services.
@@ -134,7 +143,8 @@ def _attach_user_targets(
         MERGE (p)-[r:ASSOCIATED_WITH]->(u)
         ON CREATE SET r.firstseen = timestamp()
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Relations=data,
         update_tag=update_tag,
@@ -142,7 +152,9 @@ def _attach_user_targets(
 
 
 def _attach_schedule_targets(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add relationship between escalation policy and services.
@@ -154,7 +166,8 @@ def _attach_schedule_targets(
         MERGE (p)-[r:ASSOCIATED_WITH]->(s)
         ON CREATE SET r.firstseen = timestamp()
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Relations=data,
         update_tag=update_tag,
@@ -162,7 +175,9 @@ def _attach_schedule_targets(
 
 
 def _attach_services(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add relationship between escalation policy and services.
@@ -174,7 +189,8 @@ def _attach_services(
         MERGE (s)-[r:ASSOCIATED_WITH]->(p)
         ON CREATE SET r.firstseen = timestamp()
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Relations=data,
         update_tag=update_tag,
@@ -182,7 +198,9 @@ def _attach_services(
 
 
 def _attach_teams(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add relationship between escalation policy and teams.
@@ -194,7 +212,8 @@ def _attach_teams(
         MERGE (t)-[r:ASSOCIATED_WITH]->(p)
         ON CREATE SET r.firstseen = timestamp()
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Relations=data,
         update_tag=update_tag,

cartography/intel/pagerduty/schedules.py

@@ -7,6 +7,7 @@ import dateutil.parser
 import neo4j
 from pdpyras import APISession
 
+from cartography.client.core.tx import run_write_query
 from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
@@ -32,7 +33,9 @@ def get_schedules(pd_session: APISession) -> List[Dict[str, Any]]:
 
 
 def load_schedule_data(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Transform and load schedule information
@@ -61,7 +64,8 @@ def load_schedule_data(
                 layer["_schedule_id"] = schedule["id"]
                 layers.append(layer)
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Schedules=data,
         update_tag=update_tag,
@@ -72,7 +76,9 @@ def load_schedule_data(
 
 
 def _attach_users(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add relationship between schedule and users.
@@ -83,7 +89,8 @@ def _attach_users(
         MERGE (u)-[r:MEMBER_OF]->(s)
         ON CREATE SET r.firstseen = timestamp()
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Relations=data,
         update_tag=update_tag,
@@ -91,7 +98,9 @@ def _attach_users(
 
 
 def _attach_layers(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Create layers for a schedule and attach them together
@@ -123,7 +132,8 @@ def _attach_layers(
                 users.append(
                     {"layer_id": layer["_layer_id"], "user": user["user"]["id"]},
                 )
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Layers=data,
         update_tag=update_tag,
@@ -133,7 +143,9 @@ def _attach_layers(
 
 
 def _attach_layer_users(
-    neo4j_session: neo4j.Session, data: List[Dict], update_tag: int,
+    neo4j_session: neo4j.Session,
+    data: List[Dict],
+    update_tag: int,
 ) -> None:
     """
     Add relationship between schedule layers and users.
@@ -144,7 +156,8 @@ def _attach_layer_users(
         MERGE (u)-[r:MEMBER_OF]->(l)
         ON CREATE SET r.firstseen = timestamp()
     """
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingestion_cypher_query,
         Relations=data,
         update_tag=update_tag,