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

cartography/graph/cleanupbuilder.py

@@ -3,6 +3,7 @@ from string import Template
 from typing import Dict
 from typing import List
 
+from cartography.graph.querybuilder import _asdict_with_validate_relprops
 from cartography.graph.querybuilder import _build_match_clause
 from cartography.graph.querybuilder import rel_present_on_node_schema
 from cartography.models.core.common import PropertyRef
@@ -15,35 +16,138 @@ from cartography.models.core.relationships import TargetNodeMatcher
 def build_cleanup_queries(node_schema: CartographyNodeSchema) -> List[str]:
     """
     Generates queries to clean up stale nodes and relationships from the given CartographyNodeSchema.
+    Properly handles cases where a node schema has a scoped cleanup or not.
     Note that auto-cleanups for a node with no relationships is not currently supported.
-    Algorithm:
-    1. First delete all stale nodes attached to the node_schema's sub resource
-    2. Delete all stale node to sub resource relationships
-        - We don't expect this to be very common (never for AWS resources, at least), but in case it is possible for an
-          asset to change sub resources, we want to handle it properly.
-    3. For all relationships defined on the node schema, delete all stale ones.
     :param node_schema: The given CartographyNodeSchema
     :return: A list of Neo4j queries to clean up nodes and relationships.
     """
-    if not node_schema.sub_resource_relationship:
+    # If the node has no relationships, do not delete the node. Leave this behind for the user to manage.
+    # Oftentimes these are SyncMetadata nodes.
+    if (
+        not node_schema.sub_resource_relationship
+        and not node_schema.other_relationships
+    ):
+        return []
+
+    # Case 1 [Standard]: the node has a sub resource and scoped cleanup is true => clean up stale nodes
+    # of this type, scoped to the sub resource. Continue on to clean up the other_relationships too.
+    if node_schema.sub_resource_relationship and node_schema.scoped_cleanup:
+        queries = _build_cleanup_node_and_rel_queries(
+            node_schema,
+            node_schema.sub_resource_relationship,
+        )
+
+    # Case 2: The node has a sub resource but scoped cleanup is false => this does not make sense
+    # because if have a sub resource, we are implying that we are doing scoped cleanup.
+    elif node_schema.sub_resource_relationship and not node_schema.scoped_cleanup:
         raise ValueError(
-            "Auto-creating a cleanup job for a node_schema without a sub resource relationship is not supported. "
-            f'Please check the class definition of "{node_schema.__class__.__name__}".',
+            f"This is not expected: {node_schema.label} has a sub_resource_relationship but scoped_cleanup=False."
+            "Please check the class definition for this node schema. It doesn't make sense for a node to have a "
+            "sub resource relationship and an unscoped cleanup. Doing this will cause all stale nodes of this type "
+            "to be deleted regardless of the sub resource they are attached to."
         )
 
-    result = _build_cleanup_node_and_rel_queries(node_schema, node_schema.sub_resource_relationship)
+    # Case 3: The node has no sub resource but scoped cleanup is true => do not delete any nodes, but clean up stale relationships.
+    # Return early.
+    elif not node_schema.sub_resource_relationship and node_schema.scoped_cleanup:
+        queries = []
+        other_rels = (
+            node_schema.other_relationships.rels
+            if node_schema.other_relationships
+            else []
+        )
+        for rel in other_rels:
+            query = _build_cleanup_rel_query_no_sub_resource(node_schema, rel)
+            queries.append(query)
+        return queries
+
+    # Case 4: The node has no sub resource and scoped cleanup is false => clean up the stale nodes. Continue on to clean up the other_relationships too.
+    else:
+        queries = [_build_cleanup_node_query_unscoped(node_schema)]
+
     if node_schema.other_relationships:
         for rel in node_schema.other_relationships.rels:
-            # [0] is the delete node query, [1] is the delete relationship query. We only want the latter.
-            _, rel_query = _build_cleanup_node_and_rel_queries(node_schema, rel)
-            result.append(rel_query)
+            if node_schema.scoped_cleanup:
+                # [0] is the delete node query, [1] is the delete relationship query. We only want the latter.
+                _, rel_query = _build_cleanup_node_and_rel_queries(node_schema, rel)
+                queries.append(rel_query)
+            else:
+                queries.append(_build_cleanup_rel_queries_unscoped(node_schema, rel))
+
+    return queries
+
+
+def _build_cleanup_rel_query_no_sub_resource(
+    node_schema: CartographyNodeSchema,
+    selected_relationship: CartographyRelSchema,
+) -> str:
+    """
+    Helper function to delete stale relationships for node_schemas that have no sub resource relationship defined.
+    """
+    if node_schema.sub_resource_relationship:
+        raise ValueError(
+            f"Expected {node_schema.label} to not exist. "
+            "This function is intended for node_schemas without sub_resource_relationships.",
+        )
+    # Ensure the node is attached to the sub resource and delete the node
+    query_template = Template(
+        """
+        MATCH (n:$node_label)
+        $selected_rel_clause
+        WHERE r.lastupdated <> $UPDATE_TAG
+        WITH r LIMIT $LIMIT_SIZE
+        DELETE r;
+        """,
+    )
+    return query_template.safe_substitute(
+        node_label=node_schema.label,
+        selected_rel_clause=_build_selected_rel_clause(selected_relationship),
+    )
 
-    return result
+
+def _build_match_statement_for_cleanup(node_schema: CartographyNodeSchema) -> str:
+    """
+    Helper function to build a MATCH statement for a given node schema for cleanup.
+    """
+    if not node_schema.sub_resource_relationship and not node_schema.scoped_cleanup:
+        template = Template("MATCH (n:$node_label)")
+        return template.safe_substitute(
+            node_label=node_schema.label,
+        )
+
+    # if it has a sub resource relationship defined, we need to match on the sub resource to make sure we only delete
+    # nodes that are attached to the sub resource.
+    template = Template(
+        "MATCH (n:$node_label)$sub_resource_link(:$sub_resource_label{$match_sub_res_clause})"
+    )
+    sub_resource_link = ""
+    sub_resource_label = ""
+    match_sub_res_clause = ""
+
+    if node_schema.sub_resource_relationship:
+        # Draw sub resource rel with correct direction
+        if node_schema.sub_resource_relationship.direction == LinkDirection.INWARD:
+            sub_resource_link_template = Template("<-[s:$SubResourceRelLabel]-")
+        else:
+            sub_resource_link_template = Template("-[s:$SubResourceRelLabel]->")
+        sub_resource_link = sub_resource_link_template.safe_substitute(
+            SubResourceRelLabel=node_schema.sub_resource_relationship.rel_label,
+        )
+        sub_resource_label = node_schema.sub_resource_relationship.target_node_label
+        match_sub_res_clause = _build_match_clause(
+            node_schema.sub_resource_relationship.target_node_matcher,
+        )
+    return template.safe_substitute(
+        node_label=node_schema.label,
+        sub_resource_link=sub_resource_link,
+        sub_resource_label=sub_resource_label,
+        match_sub_res_clause=match_sub_res_clause,
+    )
 
 
 def _build_cleanup_node_and_rel_queries(
-        node_schema: CartographyNodeSchema,
-        selected_relationship: CartographyRelSchema,
+    node_schema: CartographyNodeSchema,
+    selected_relationship: CartographyRelSchema,
 ) -> List[str]:
     """
     Private function that performs the main string template logic for generating cleanup node and relationship queries.
@@ -67,15 +171,6 @@ def _build_cleanup_node_and_rel_queries(
             "verify the node class definition for the relationships that it has.",
         )
 
-    # Draw sub resource rel with correct direction
-    if node_schema.sub_resource_relationship.direction == LinkDirection.INWARD:
-        sub_resource_link_template = Template("<-[s:$SubResourceRelLabel]-")
-    else:
-        sub_resource_link_template = Template("-[s:$SubResourceRelLabel]->")
-    sub_resource_link = sub_resource_link_template.safe_substitute(
-        SubResourceRelLabel=node_schema.sub_resource_relationship.rel_label,
-    )
-
     # The cleanup node query must always be before the cleanup rel query
     delete_action_clauses = [
         """
@@ -86,7 +181,9 @@ def _build_cleanup_node_and_rel_queries(
     ]
     # Now clean up the relationships
     if selected_relationship == node_schema.sub_resource_relationship:
-        _validate_target_node_matcher_for_cleanup_job(node_schema.sub_resource_relationship.target_node_matcher)
+        _validate_target_node_matcher_for_cleanup_job(
+            node_schema.sub_resource_relationship.target_node_matcher,
+        )
         delete_action_clauses.append(
             """
             WHERE s.lastupdated <> $UPDATE_TAG
@@ -106,26 +203,99 @@ def _build_cleanup_node_and_rel_queries(
     # Ensure the node is attached to the sub resource and delete the node
     query_template = Template(
         """
-        MATCH (n:$node_label)$sub_resource_link(:$sub_resource_label{$match_sub_res_clause})
+        $match_statement
         $selected_rel_clause
         $delete_action_clause
         """,
     )
     return [
         query_template.safe_substitute(
-            node_label=node_schema.label,
-            sub_resource_link=sub_resource_link,
-            sub_resource_label=node_schema.sub_resource_relationship.target_node_label,
-            match_sub_res_clause=_build_match_clause(node_schema.sub_resource_relationship.target_node_matcher),
+            match_statement=_build_match_statement_for_cleanup(node_schema),
             selected_rel_clause=(
-                "" if selected_relationship == node_schema.sub_resource_relationship
+                ""
+                if selected_relationship == node_schema.sub_resource_relationship
                 else _build_selected_rel_clause(selected_relationship)
             ),
             delete_action_clause=delete_action_clause,
-        ) for delete_action_clause in delete_action_clauses
+        )
+        for delete_action_clause in delete_action_clauses
     ]
 
 
+def _build_cleanup_node_query_unscoped(
+    node_schema: CartographyNodeSchema,
+) -> str:
+    """
+    Generates a cleanup query for a node_schema to allow unscoped cleanup.
+    """
+    if node_schema.scoped_cleanup:
+        raise ValueError(
+            f"_build_cleanup_node_query_for_unscoped_cleanup() failed: '{node_schema.label}' does not have "
+            "scoped_cleanup=False, so we cannot generate a query to clean it up. Please verify that the class "
+            "definition is what you expect.",
+        )
+
+    # The cleanup node query must always be before the cleanup rel query
+    delete_action_clause = """
+        WHERE n.lastupdated <> $UPDATE_TAG
+        WITH n LIMIT $LIMIT_SIZE
+        DETACH DELETE n;
+    """
+
+    # Ensure the node is attached to the sub resource and delete the node
+    query_template = Template(
+        """
+        $match_statement
+        $delete_action_clause
+        """,
+    )
+    return query_template.safe_substitute(
+        match_statement=_build_match_statement_for_cleanup(node_schema),
+        delete_action_clause=delete_action_clause,
+    )
+
+
+def _build_cleanup_rel_queries_unscoped(
+    node_schema: CartographyNodeSchema,
+    selected_relationship: CartographyRelSchema,
+) -> str:
+    """
+    Generates relationship cleanup query for a node_schema with scoped_cleanup=False.
+    """
+    if node_schema.scoped_cleanup:
+        raise ValueError(
+            f"_build_cleanup_node_and_rel_queries_unscoped() failed: '{node_schema.label}' does not have "
+            "scoped_cleanup=False, so we cannot generate a query to clean it up. Please verify that the class "
+            "definition is what you expect.",
+        )
+    if not rel_present_on_node_schema(node_schema, selected_relationship):
+        raise ValueError(
+            f"_build_cleanup_node_query(): Attempted to build cleanup query for node '{node_schema.label}' and "
+            f"relationship {selected_relationship.rel_label} but that relationship is not present on the node. Please "
+            "verify the node class definition for the relationships that it has.",
+        )
+
+    # The cleanup node query must always be before the cleanup rel query
+    delete_action_clause = """WHERE r.lastupdated <> $UPDATE_TAG
+        WITH r LIMIT $LIMIT_SIZE
+        DELETE r;
+        """
+
+    # Ensure the node is attached to the sub resource and delete the node
+    query_template = Template(
+        """
+        $match_statement
+        $selected_rel_clause
+        $delete_action_clause
+        """,
+    )
+    return query_template.safe_substitute(
+        match_statement=_build_match_statement_for_cleanup(node_schema),
+        selected_rel_clause=_build_selected_rel_clause(selected_relationship),
+        delete_action_clause=delete_action_clause,
+    )
+
+
 def _build_selected_rel_clause(selected_relationship: CartographyRelSchema) -> str:
     """
     Draw selected relationship with correct direction. Returns a string that looks like either
@@ -136,8 +306,12 @@ def _build_selected_rel_clause(selected_relationship: CartographyRelSchema) -> s
         selected_rel_template = Template("<-[r:$SelectedRelLabel]-")
     else:
         selected_rel_template = Template("-[r:$SelectedRelLabel]->")
-    selected_rel = selected_rel_template.safe_substitute(SelectedRelLabel=selected_relationship.rel_label)
-    selected_rel_clause_template = Template("""MATCH (n)$selected_rel(:$other_node_label)""")
+    selected_rel = selected_rel_template.safe_substitute(
+        SelectedRelLabel=selected_relationship.rel_label,
+    )
+    selected_rel_clause_template = Template(
+        """MATCH (n)$selected_rel(:$other_node_label)""",
+    )
     selected_rel_clause = selected_rel_clause_template.safe_substitute(
         selected_rel=selected_rel,
         other_node_label=selected_relationship.target_node_label,
@@ -161,3 +335,49 @@ def _validate_target_node_matcher_for_cleanup_job(tgm: TargetNodeMatcher):
                 f"{key} has set_in_kwargs=False, please check by reviewing the full stack trace to know which object"
                 f"this message was raised from. Debug information: PropertyRef name = {prop_ref.name}.",
             )
+
+
+def build_cleanup_query_for_matchlink(rel_schema: CartographyRelSchema) -> str:
+    """
+    Generates a cleanup query for a matchlink relationship.
+    :param rel_schema: The CartographyRelSchema object to generate a query. This CartographyRelSchema object
+    - Must have a source_node_matcher and source_node_label defined
+    - Must have a CartographyRelProperties object where _sub_resource_label and _sub_resource_id are defined
+    :return: A Neo4j query used to clean up stale matchlink relationships.
+    """
+    if not rel_schema.source_node_matcher:
+        raise ValueError(
+            f"No source node matcher found for {rel_schema.rel_label}; returning empty list."
+        )
+
+    query_template = Template(
+        """
+        MATCH (from:$source_node_label)$rel_direction[r:$rel_label]$rel_direction_end(to:$target_node_label)
+        WHERE r.lastupdated <> $UPDATE_TAG
+            AND r._sub_resource_label = $sub_resource_label
+            AND r._sub_resource_id = $sub_resource_id
+        WITH r LIMIT $LIMIT_SIZE
+        DELETE r;
+        """
+    )
+
+    # Determine which way to point the arrow. INWARD is toward the source, otherwise we go toward the target.
+    if rel_schema.direction == LinkDirection.INWARD:
+        rel_direction = "<-"
+        rel_direction_end = "-"
+    else:
+        rel_direction = "-"
+        rel_direction_end = "->"
+
+    # Small hack: avoid type-checking errors by converting the rel_schema to a dict.
+    rel_props_as_dict = _asdict_with_validate_relprops(rel_schema)
+
+    return query_template.safe_substitute(
+        source_node_label=rel_schema.source_node_label,
+        target_node_label=rel_schema.target_node_label,
+        rel_label=rel_schema.rel_label,
+        rel_direction=rel_direction,
+        rel_direction_end=rel_direction_end,
+        sub_resource_label=rel_props_as_dict["_sub_resource_label"],
+        sub_resource_id=rel_props_as_dict["_sub_resource_id"],
+    )

cartography/graph/job.py

@@ -13,9 +13,11 @@ from typing import Union
 import neo4j
 
 from cartography.graph.cleanupbuilder import build_cleanup_queries
+from cartography.graph.cleanupbuilder import build_cleanup_query_for_matchlink
 from cartography.graph.statement import get_job_shortname
 from cartography.graph.statement import GraphStatement
 from cartography.models.core.nodes import CartographyNodeSchema
+from cartography.models.core.relationships import CartographyRelSchema
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +34,7 @@ def _get_identifiers(template: string.Template) -> List[str]:
             filter(
                 lambda v: v is not None,
                 (
-                    mo.group('named') or mo.group('braced')
+                    mo.group("named") or mo.group("braced")
                     for mo in template.pattern.finditer(template.template)
                 ),
             ),
@@ -71,7 +73,12 @@ class GraphJob:
     A job that will run against the cartography graph. A job is a sequence of statements which execute sequentially.
     """
 
-    def __init__(self, name: str, statements: List[GraphStatement], short_name: Optional[str] = None):
+    def __init__(
+        self,
+        name: str,
+        statements: List[GraphStatement],
+        short_name: Optional[str] = None,
+    ):
         # E.g. "Okta intel module cleanup"
         self.name = name
         self.statements: List[GraphStatement] = statements
@@ -100,7 +107,11 @@ class GraphJob:
                     e,
                 )
                 raise
-        log_msg = f"Finished job {self.short_name}" if self.short_name else f"Finished job {self.name}"
+        log_msg = (
+            f"Finished job {self.short_name}"
+            if self.short_name
+            else f"Finished job {self.name}"
+        )
         logger.info(log_msg)
 
     def as_dict(self) -> Dict:
@@ -114,43 +125,55 @@ class GraphJob:
         }
 
     @classmethod
-    def from_json(cls, blob: str, short_name: Optional[str] = None) -> 'GraphJob':
+    def from_json(
+        cls, blob: Union[str, dict], short_name: Optional[str] = None
+    ) -> "GraphJob":
         """
-        Create a job from a JSON blob.
+        Create a job from a JSON dict or blob.
         """
-        data: Dict = json.loads(blob)
+        data = json.loads(blob) if isinstance(blob, str) else blob
         statements = _get_statements_from_json(data, short_name)
         name = data["name"]
         return cls(name, statements, short_name)
 
     @classmethod
     def from_node_schema(
-            cls,
-            node_schema: CartographyNodeSchema,
-            parameters: Dict[str, Any],
-    ) -> 'GraphJob':
+        cls,
+        node_schema: CartographyNodeSchema,
+        parameters: Dict[str, Any],
+        iterationsize: int = 100,
+    ) -> "GraphJob":
         """
         Create a cleanup job from a CartographyNodeSchema object.
         For a given node, the fields used in the node_schema.sub_resource_relationship.target_node_node_matcher.keys()
         must be provided as keys and values in the params dict.
+        :param iterationsize: The number of items to process in each iteration. Defaults to 100.
         """
         queries: List[str] = build_cleanup_queries(node_schema)
 
         expected_param_keys: Set[str] = get_parameters(queries)
         actual_param_keys: Set[str] = set(parameters.keys())
         # Hacky, but LIMIT_SIZE is specified by default in cartography.graph.statement, so we exclude it from validation
-        actual_param_keys.add('LIMIT_SIZE')
+        actual_param_keys.add("LIMIT_SIZE")
 
         missing_params: Set[str] = expected_param_keys - actual_param_keys
 
         if missing_params:
             raise ValueError(
                 f'GraphJob is missing the following expected query parameters: "{missing_params}". Please check the '
-                f'value passed to `parameters`.',
+                f"value passed to `parameters`.",
             )
 
         statements: List[GraphStatement] = [
-            GraphStatement(query, parameters=parameters, iterative=True, iterationsize=100) for query in queries
+            GraphStatement(
+                query,
+                parameters=parameters,
+                iterative=True,
+                iterationsize=iterationsize,
+                parent_job_name=node_schema.label,
+                parent_job_sequence_num=idx,
+            )
+            for idx, query in enumerate(queries, start=1)
         ]
 
         return cls(
@@ -160,7 +183,49 @@ class GraphJob:
         )
 
     @classmethod
-    def from_json_file(cls, file_path: Union[str, Path]) -> 'GraphJob':
+    def from_matchlink(
+        cls,
+        rel_schema: CartographyRelSchema,
+        sub_resource_label: str,
+        sub_resource_id: str,
+        update_tag: int,
+        iterationsize: int = 100,
+    ) -> "GraphJob":
+        """
+        Create a cleanup job from a CartographyRelSchema object (specifically, a MatchLink).
+        This is used for cleaning up stale links between nodes created by load_rels(). Do not use for other purposes.
+
+        Other notes:
+        - For a given rel_schema, the fields used in the rel_schema.properties._sub_resource_label.name and
+        rel_schema.properties._sub_resource_id.name must be provided as keys and values in the params dict.
+        - The rel_schema must have a source_node_matcher and target_node_matcher.
+        :param iterationsize: The number of items to process in each iteration. Defaults to 100.
+        """
+        cleanup_link_query = build_cleanup_query_for_matchlink(rel_schema)
+        logger.debug(f"Cleanup query: {cleanup_link_query}")
+
+        parameters = {
+            "UPDATE_TAG": update_tag,
+            "_sub_resource_label": sub_resource_label,
+            "_sub_resource_id": sub_resource_id,
+        }
+
+        statement = GraphStatement(
+            cleanup_link_query,
+            parameters=parameters,
+            iterative=True,
+            iterationsize=iterationsize,
+            parent_job_name=rel_schema.rel_label,
+        )
+
+        return cls(
+            f"Cleanup {rel_schema.rel_label} between {rel_schema.source_node_label} and {rel_schema.target_node_label}",
+            [statement],
+            rel_schema.rel_label,
+        )
+
+    @classmethod
+    def from_json_file(cls, file_path: Union[str, Path]) -> "GraphJob":
         """
         Create a job from a JSON file.
         """
@@ -168,16 +233,23 @@ class GraphJob:
             data: Dict = json.load(j_file)
 
         job_shortname: str = get_job_shortname(file_path)
-        statements: List[GraphStatement] = _get_statements_from_json(data, job_shortname)
+        statements: List[GraphStatement] = _get_statements_from_json(
+            data,
+            job_shortname,
+        )
         name: str = data["name"]
         return cls(name, statements, job_shortname)
 
     @classmethod
     def run_from_json(
-        cls, neo4j_session: neo4j.Session, blob: str, parameters: Dict, short_name: Optional[str] = None,
+        cls,
+        neo4j_session: neo4j.Session,
+        blob: Union[str, dict],
+        parameters: Dict,
+        short_name: Optional[str] = None,
     ) -> None:
         """
-        Run a job from a JSON blob. This will deserialize the job and execute all statements sequentially.
+        Run a job from a JSON dict or blob. This will deserialize the job and execute all statements sequentially.
         """
         if not parameters:
             parameters = {}
@@ -187,7 +259,12 @@ class GraphJob:
         job.run(neo4j_session)
 
     @classmethod
-    def run_from_json_file(cls, file_path: Union[str, Path], neo4j_session: neo4j.Session, parameters: Dict) -> None:
+    def run_from_json_file(
+        cls,
+        file_path: Union[str, Path],
+        neo4j_session: neo4j.Session,
+        parameters: Dict,
+    ) -> None:
         """
         Run a job from a JSON file. This will deserialize the job and execute all statements sequentially.
         """
@@ -200,14 +277,21 @@ class GraphJob:
         job.run(neo4j_session)
 
 
-def _get_statements_from_json(blob: Dict, short_job_name: Optional[str] = None) -> List[GraphStatement]:
+def _get_statements_from_json(
+    blob: Dict,
+    short_job_name: Optional[str] = None,
+) -> List[GraphStatement]:
     """
     Deserialize all statements from the JSON blob.
     """
     statements: List[GraphStatement] = []
     for i, statement_data in enumerate(blob["statements"]):
         # i+1 to make it 1-based and not 0-based to help with log readability
-        statement: GraphStatement = GraphStatement.create_from_json(statement_data, short_job_name, i + 1)
+        statement: GraphStatement = GraphStatement.create_from_json(
+            statement_data,
+            short_job_name,
+            i + 1,
+        )
         statements.append(statement)
 
     return statements