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

cartography/intel/okta/users.py

@@ -8,11 +8,11 @@ import neo4j
 from okta import UsersClient
 from okta.models.user import User
 
+from cartography.client.core.tx import run_write_query
 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 +56,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 +93,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 +105,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 +128,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:
     """
@@ -150,7 +160,14 @@ def _load_okta_users(
     new_user.okta_last_updated = user_data.okta_last_updated,
     new_user.password_changed = user_data.password_changed,
     new_user.transition_to_status = user_data.transition_to_status,
-    new_user.lastupdated = $okta_update_tag
+    new_user.lastupdated = $okta_update_tag,
+    new_user:UserAccount,
+    new_user._module_name = "cartography:okta",
+    new_user._ont_email = user_data.email,
+    new_user._ont_firstname = user_data.first_name,
+    new_user._ont_lastname = user_data.last_name,
+    new_user._ont_lastactivity = user_data.last_login,
+    new_user._ont_source = "okta"
     WITH new_user, org
     MERGE (org)-[org_r:RESOURCE]->(new_user)
     ON CREATE SET org_r.firstseen = timestamp()
@@ -164,7 +181,8 @@ def _load_okta_users(
     SET h.lastupdated = $okta_update_tag
     """
 
-    neo4j_session.run(
+    run_write_query(
+        neo4j_session,
         ingest_statement,
         ORG_ID=okta_org_id,
         USER_LIST=user_list,
@@ -174,8 +192,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/ontology/__init__.py

@@ -0,0 +1,44 @@
+import logging
+
+import neo4j
+
+import cartography.intel.ontology.devices
+import cartography.intel.ontology.users
+from cartography.config import Config
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def run(neo4j_session: neo4j.Session, config: Config) -> None:
+    common_job_parameters = {
+        "UPDATE_TAG": config.update_tag,
+    }
+
+    # Get source of truth from config
+    if config.ontology_users_source:
+        users_source_of_truth = [
+            source.strip() for source in config.ontology_users_source.split(",")
+        ]
+    else:
+        users_source_of_truth = []
+    if config.ontology_devices_source:
+        computers_source_of_truth = [
+            source.strip() for source in config.ontology_devices_source.split(",")
+        ]
+    else:
+        computers_source_of_truth = []
+
+    cartography.intel.ontology.users.sync(
+        neo4j_session,
+        users_source_of_truth,
+        config.update_tag,
+        common_job_parameters,
+    )
+    cartography.intel.ontology.devices.sync(
+        neo4j_session,
+        computers_source_of_truth,
+        config.update_tag,
+        common_job_parameters,
+    )

cartography/intel/ontology/devices.py

@@ -0,0 +1,54 @@
+import logging
+from typing import Any
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.ontology.utils import get_source_nodes_from_graph
+from cartography.intel.ontology.utils import link_ontology_nodes
+from cartography.models.ontology.device import DeviceSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    source_of_truth: list[str],
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    data = get_source_nodes_from_graph(neo4j_session, source_of_truth, "devices")
+    load_devices(
+        neo4j_session,
+        data,
+        update_tag,
+    )
+    link_ontology_nodes(neo4j_session, "devices", update_tag)
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def load_devices(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        DeviceSchema(),
+        data,
+        lastupdated=update_tag,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    GraphJob.from_node_schema(DeviceSchema(), common_job_parameters).run(
+        neo4j_session,
+    )

cartography/intel/ontology/users.py

@@ -0,0 +1,54 @@
+import logging
+from typing import Any
+
+import neo4j
+
+from cartography.client.core.tx import load
+from cartography.graph.job import GraphJob
+from cartography.intel.ontology.utils import get_source_nodes_from_graph
+from cartography.intel.ontology.utils import link_ontology_nodes
+from cartography.models.ontology.user import UserSchema
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def sync(
+    neo4j_session: neo4j.Session,
+    source_of_truth: list[str],
+    update_tag: int,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    data = get_source_nodes_from_graph(neo4j_session, source_of_truth, "users")
+    load_users(
+        neo4j_session,
+        data,
+        update_tag,
+    )
+    link_ontology_nodes(neo4j_session, "users", update_tag)
+    cleanup(neo4j_session, common_job_parameters)
+
+
+@timeit
+def load_users(
+    neo4j_session: neo4j.Session,
+    data: list[dict[str, Any]],
+    update_tag: int,
+) -> None:
+    load(
+        neo4j_session,
+        UserSchema(),
+        data,
+        lastupdated=update_tag,
+    )
+
+
+@timeit
+def cleanup(
+    neo4j_session: neo4j.Session,
+    common_job_parameters: dict[str, Any],
+) -> None:
+    GraphJob.from_node_schema(UserSchema(), common_job_parameters).run(
+        neo4j_session,
+    )

cartography/intel/ontology/utils.py

@@ -0,0 +1,176 @@
+import logging
+from dataclasses import asdict
+from typing import Any
+
+import neo4j
+
+from cartography.client.core.tx import read_list_of_dicts_tx
+from cartography.graph.job import GraphJob
+from cartography.models.ontology.mapping import ONTOLOGY_MODELS
+from cartography.models.ontology.mapping import ONTOLOGY_NODES_MAPPING
+from cartography.models.ontology.mapping.specs import OntologyNodeMapping
+from cartography.util import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@timeit
+def _run_source_node_single_query(
+    module_name: str,
+    node: OntologyNodeMapping,
+    neo4j_session: neo4j.Session,
+    query: str,
+    results: dict[str, dict[str, Any]],
+    **kwargs: Any,
+) -> dict[str, dict[str, Any]]:
+    # DOC
+    for row in neo4j_session.execute_read(read_list_of_dicts_tx, query, **kwargs):
+        node_data = row["n"]
+        result: dict[str, Any] = {}
+        skip_node: bool = False
+
+        # Extract only the fields defined in the ontology mapping
+        for field in node.fields:
+            value = node_data.get(field.node_field)
+            # Skip nodes missing required fields
+            if field.required and not value:
+                logger.debug(
+                    "Skipping node with label '%s' due to missing required field '%s'.",
+                    node.node_label,
+                    field.node_field,
+                )
+                skip_node = True
+                break
+            result[field.ontology_field] = value
+        if skip_node:
+            continue
+
+        # Merge results based on the node's id field to avoid duplicates
+        ontology_model = ONTOLOGY_MODELS[module_name]
+        if ontology_model is None:
+            # Should not happen as we skip non-eligible nodes above
+            logger.warning(
+                "No ontology model found for module '%s'. Skipping node label '%s'.",
+                module_name,
+                node.node_label,
+            )
+            continue
+        id_field = ontology_model().properties.id.name
+        existing = results.get(result[id_field])
+        if existing:
+            logger.debug("Merging node: %s to %s", result[id_field], existing[id_field])
+            # Merge existing data with new data, prioritizing non-None values
+            for key, value in result.items():
+                if existing.get(key) is None and value is not None:
+                    existing[key] = value
+        else:
+            logger.debug("Adding new node: %s", result[id_field])
+            results[result[id_field]] = result
+    return results
+
+
+@timeit
+def get_source_nodes_from_graph(
+    neo4j_session: neo4j.Session,
+    source_of_truth: list[str],
+    module_name: str,
+) -> list[dict[str, Any]]:
+    """Retrieve source nodes from the Neo4j graph database based on the ontology mapping.
+
+    This function queries the Neo4j database for nodes that match the labels
+    defined in the ontology mapping for the specified module and source of truth.
+    It returns a list of dictionaries containing the relevant fields for each node.
+
+    If no source of truth is provided, default to all sources defined in the mapping.
+
+    Args:
+        neo4j_session (neo4j.Session): The Neo4j session to use for querying the database.
+        source_of_truth (list[str]): A list of source of truth identifiers to filter the modules.
+        module_name (str): The name of the ontology module to use for the mapping (eg. users, devices, etc.).
+
+    Returns:
+        list[dict[str, Any]]: A list of dictionaries, each containing a node details formatted according to the ontology mapping.
+    """
+    results: dict[str, dict[str, Any]] = {}
+    modules_mapping = ONTOLOGY_NODES_MAPPING[module_name]
+    if len(source_of_truth) == 0:
+        source_of_truth = list(modules_mapping.keys())
+    # Check if ontology nodes are used in mapping
+    _has_ontology = False
+    if modules_mapping.get("ontology") is not None:
+        _has_ontology = True
+        for node in modules_mapping["ontology"].nodes:
+            if not node.eligible_for_source:
+                logger.debug(
+                    "Skipping ontology node with label '%s' as it is not eligible for source of truth.",
+                    node.node_label,
+                )
+                continue
+            # Run the query for every source
+            for source in source_of_truth:
+                # Use parameterized query to prevent Cypher injection attacks
+                query = f"MATCH (n:{node.node_label} {{_ont_source: $source}}) RETURN n"
+                results = _run_source_node_single_query(
+                    module_name, node, neo4j_session, query, results, source=source
+                )
+
+    # Run queries for each source of truth
+    for source in source_of_truth:
+        if source not in modules_mapping:
+            if not _has_ontology:
+                logger.warning(
+                    "Source of truth '%s' is not supported for '%s'.",
+                    source,
+                    module_name,
+                )
+            continue
+        for node in modules_mapping[source].nodes:
+            if not node.eligible_for_source:
+                logger.debug(
+                    "Skipping node with label '%s' as it is not eligible for source of truth '%s'.",
+                    node.node_label,
+                    source,
+                )
+                continue
+            query = f"MATCH (n:{node.node_label}) RETURN n"
+            results = _run_source_node_single_query(
+                module_name, node, neo4j_session, query, results
+            )
+
+    return list(results.values())
+
+
+@timeit
+def link_ontology_nodes(
+    neo4j_session: neo4j.Session,
+    module_name: str,
+    update_tag: int,
+) -> None:
+    """Link ontology nodes in the Neo4j graph database based on the ontology mapping.
+
+    This function retrieves the ontology mapping for the specified module and
+    executes the relationship statements defined in the mapping to link nodes
+    in the Neo4j graph database.
+
+    Args:
+        neo4j_session (neo4j.Session): The Neo4j session to use for executing the relationship statements.
+        module_name (str): The name of the ontology module for which to link nodes (eg. users, devices, etc.).
+        update_tag (int): The update tag of the current run, used to tag the changes in the graph.
+    """
+    modules_mapping = ONTOLOGY_NODES_MAPPING.get(module_name)
+    if modules_mapping is None:
+        logger.warning("No ontology mapping found for module '%s'.", module_name)
+        return
+    for source, mapping in modules_mapping.items():
+        if len(mapping.rels) == 0:
+            continue
+        formated_json = {
+            "name": f"Linking ontology nodes for {module_name} for source {source}",
+            "statements": [asdict(rel) for rel in mapping.rels],
+        }
+        GraphJob.run_from_json(
+            neo4j_session,
+            formated_json,
+            {"UPDATE_TAG": update_tag},
+            short_name=f"ontology.{module_name}.{source}.linking",
+        )

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,89 @@
+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,
+) -> None:
+    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)
+
+
+@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
+    )