def create_alert_graph()

in msticpy/nbtools/security_alert_graph.py [0:0]

• 49 lines of code
• 16 McCabe index (conditional complexity)

def create_alert_graph(alert: SecurityAlert):
    """Create a networkx graph from the alert and contained entities."""
    alertentity_graph = nx.Graph(id="AlertGraph")

    alertentity_graph.add_node(
        alert["AlertType"],
        name=alert["AlertType"],
        time=str(alert["StartTimeUtc"]),
        description="Alert: " + alert["AlertDisplayName"],
        color="red",
        node_type="alert",
    )

    os_family = alert.os_family

    # Cycle through entities
    for entity in alert.entities:
        (e_name, e_desc) = _get_name_and_description(entity, os_family)

        alertentity_graph.add_node(
            e_name,
            entitytype=entity["Type"],
            name=e_name,
            description=e_desc,
            color="green",
            node_type="entity",
            source=str(entity),
        )

        # add an edge by default to the alert
        alertentity_graph.add_edge(alert["AlertType"], e_name)

        # Rather than just add edges to the alert, we want to follow the 'natural'
        # relationships between entities and child entities
        # So if this entity has a property that is an entity, we add an edge to it
        # and prune any edge that it might have to the alert
        if isinstance(entity, Entity):
            ent_props = entity.properties
        elif isinstance(entity, dict):
            ent_props = entity
        else:
            continue
        for prop, rel_entity in [
            (p, v) for (p, v) in ent_props.items() if isinstance(v, Entity)
        ]:
            if rel_entity["Type"] == "host":
                # don't add a new edge to the host
                continue

            # get the node id of the related entity and add an edge if it
            # doesn't already exist
            (related_entity, _) = _get_name_and_description(rel_entity)
            if not alertentity_graph.has_edge(related_entity, e_name):
                alertentity_graph.add_edge(
                    e_name,
                    related_entity,
                    description=prop,
                    color="green",
                    weight=1,
                    line_type="SHORT_DASH",
                )

            # if we have a previously created an edge to the alert, remove it
            if alertentity_graph.has_edge(alert["AlertType"], related_entity):
                alertentity_graph.remove_edge(alert["AlertType"], related_entity)

        # if we haven't added an edge to this entity from anything else,
        # add one to the alert
        if not alertentity_graph.neighbors(e_name):
            alertentity_graph.add_edge(alert["AlertType"], e_name)

    return alertentity_graph