SH-Deployer/apps/SA-ITOA/lib/ITOA/service_tree.py

# Copyright (C) 2005-2024 Splunk Inc. All Rights Reserved.

import uuid
from collections import deque

MAX_SERVICES_SUPPORTED_IN_TREE_VIEW = 2000


def _generate_service_dependency_for_partial_backup(vertices):
    '''
    @param vertices:  dict of service vertex object
    @return: dict of service dependency relationship
    {'service_id_1': {'_key': service_id_1, 'title': service_title_1, 'parents':['service_id_2', 'service_id_3']}
    }
    '''
    output = {}
    for v in list(vertices.values()):
        data = {}
        data['_key'] = v.id
        data['title'] = v.title
        # parse the relationship
        data['parents'] = [i.id for i in v.parents]
        output[v.id] = data
    return output


def generate_subgraphs_json(data, service_id_filter=[], logger=None, partial_backup=False, max_depth_filter=-2,
                            severity_filter_ids=[]):
    """
    Handles a request to create a service tree from the given data.

    @type data: list
    @param data: the raw service data

    @type service_id_filter: list
    @param service_id_filter: the list of filtered service ids

    @type logger: object
    @param logger: the logger to use

    @type partial_backup: bool
    @param partial_backup: if the subgraph is used by partial backup

    @type max_depth_filter: integer
    @param max_depth_filter: maximum depth value filter

    @type severity_filter_ids: list
    @param severity_filter_ids: the list of service ids from the severity filter

    @rtype: list
    @return: the list of subgraphs in the tree
    """
    tree = ServiceTree(logger=logger)

    vertices = tree.parse_services_to_vertices(data)
    if partial_backup:
        return _generate_service_dependency_for_partial_backup(vertices)
    else:
        filtered_vertices = tree.filter_vertices(vertices, service_id_filter, severity_filter_ids)
        subgraphs = tree.create_subgraphs(filtered_vertices)
        return tree.subgraphs_to_json(subgraphs, max_depth_filter)


class ServiceTree(object):
    """
    Class to interact with service trees
    """

    def __init__(self, logger=None):
        """
        @type logger: object
        @param logger: the logger to use
        """
        self.logger = logger

    def parse_services_to_vertices(self, services_data):
        """
        Parse the raw services data, returning a lookup dict from service key
        to a vertex representing the service in the tree.

        @type services_data: dict
        @param services_data: the raw services data

        @rtype: dict
        @return: dict from service key to its vertex in the tree
        """
        if not services_data:
            return {}

        vertices = {}

        for service_data in services_data:
            service_id = service_data.get('_key', None)
            if not service_id:
                continue

            vertex = vertices.setdefault(service_id, ServiceVertex(service_id))
            vertex.title = service_data.get('title', '')

            # Parse the service's dependent services (parents)
            parent_ids = []
            for parent in service_data.get('services_depends_on', []):
                parent_id = parent.get('serviceid', None)
                if not parent_id:
                    continue
                parent_ids.append(parent_id)

            # Set the vertex's parents
            parents = {vertices.setdefault(parent_id, ServiceVertex(parent_id))
                       for parent_id in parent_ids}
            vertex.parents.update(parents)

            # Add this vertex to its parents as a child
            for parent in parents:
                parent.children.add(vertex)

            vertices[service_id] = vertex

        return vertices

    def filter_vertices(self, vertices, vertex_ids, severity_filter_ids=[]):
        """
        Filters the given list of vertices. If the list of vertex_ids is provided,
        only vertices that are connected to a vertex in vertex_ids in a meaningful
        way are included in the output list. If the list of severity ids is provided,
        only vertices that are children of a vertex in this list will be included in the
        output list. If both are provided, the intersection of the search results from each
        list will be returned.

        @type vertices: dict
        @param vertices: the lookup dict from vertex id to service vertex

        @type vertex_ids: list
        @param vertices: the list of vertex ids from the service filter

        @type severity_filter_ids: list
        @param severity_filter_ids: the list of vertex ids that have a minimum severity

        @rtype: dict
        @return: dict from service key to its vertex in the tree
        """
        if not vertex_ids and not severity_filter_ids:
            return vertices

        # Maintain separate structures for keeping track of which vertices we've visited
        # from a particular direction (i.e. either through children or through parents)
        children = set()
        parents = set()

        visited_by_service = set()
        # Traverse the vertices so that only the vertices associated with the filter are visited
        if vertex_ids:
            for vertex_id in vertex_ids:
                vertex = vertices.get(vertex_id, None)
                if not vertex:
                    continue

                vertex.traverse_children(children)
                vertex.traverse_parents(parents)

            visited_by_service.update(children)
            visited_by_service.update(parents)
        else:
            visited_by_service = {vertices[v] for v in vertices}
        children = set()

        visited_by_severity = set()
        # Traverse the vertices so that only the vertices that are children of ones from the severity filter are visited
        if severity_filter_ids:
            for vertex_id in severity_filter_ids:
                vertex = vertices.get(vertex_id, None)
                if not vertex:
                    continue

                vertex.traverse_children(children)
            visited_by_severity.update(children)
        else:
            visited_by_severity = {vertices[v] for v in vertices}

        visited = visited_by_service.intersection(visited_by_severity)

        # Update the vertices edge data with only the visited vertices
        filtered = {}
        for v in visited:
            filtered[v.id] = v
            v.children.intersection_update(visited)
            v.parents.intersection_update(visited)

        return filtered

    def create_subgraphs(self, vertices):
        """
        Groups connected vertices into a subgraph object.

        @type vertices: dict
        @param vertices: the lookup dict from vertex id to service vertex

        @rtype: list
        @return: list of subgraphs
        """
        subgraphs = []
        visited = set()

        for vertex in list(vertices.values()):
            if vertex in visited:
                continue

            subgraph = Subgraph()
            subgraph.populate_from(vertex, visited)
            subgraphs.append(subgraph)

        return subgraphs

    def subgraphs_to_json(self, subgraphs, max_depth_filter=-2):
        """
        Returns a JSON-friendly object representing the list of subgraphs.

        @type subgraphs: list
        @param subgraphs: the list of subgraph objects

        @type max_depth_filter: integer
        @param max_depth_filter: maximum depth value filter

        @rtype: list of dict
        @return: list of subgraphs that can be converted to JSON
        """
        graphs = []
        total = 0
        queue = deque()
        visited = set()
        maxPossibleDepth = -1
        depth = -1
        numVisibleNodes = 0
        visibleCount = 0

        for subgraph in subgraphs:
            # putting vertices with highest degrees first seems to create
            # better looking graphs
            sorted_edges = sorted(subgraph.edges,
                                  key=lambda e: (e[1].degree, e[0].degree,
                                                 e[1].title, e[0].title),
                                  reverse=True)
            edges = [{
                'source': e[1].id,
                'target': e[0].id
            } for e in sorted_edges]

            vertices = [{
                'id': v.id,
                'title': v.title,
                'nodeDepth': v.node_depth
            } for v in sorted(subgraph.vertices, key=lambda x: x.title)]

            filtered_vertices = []
            if max_depth_filter != -2:
                for v in vertices:
                    if v.get('nodeDepth') <= max_depth_filter:
                        filtered_vertices.append(v)
            else:
                filtered_vertices = vertices

            cycle = subgraph.has_cycle()
            graphs.append({
                'id': subgraph.id,
                'edges': edges,
                'vertices': filtered_vertices,
                'has_cycle': cycle
            })

            if cycle and numVisibleNodes < MAX_SERVICES_SUPPORTED_IN_TREE_VIEW:
                numVisibleNodes += len(subgraph.vertices)

            for vertex in subgraph.vertices:
                if vertex.node_depth == 0 and vertex not in visited:
                    visited.add(vertex)
                    queue.append(vertex)

            visibleCount += len(filtered_vertices)
            total += len(vertices)

        if numVisibleNodes >= MAX_SERVICES_SUPPORTED_IN_TREE_VIEW:
            # can't show anything, even nodes in cycles exceed MAX_SERVICES_SUPPORTED_IN_TREE_VIEW
            return {
                'graphs': graphs,
                'totalCount': total,
                'visibleCount': visibleCount,
                'depth': -2,
                'maxPossibleDepth': -2
            }

        while queue:
            numInCurrLevel = len(queue)
            for i in range(numInCurrLevel):
                vertex = queue.popleft()

                if numVisibleNodes < MAX_SERVICES_SUPPORTED_IN_TREE_VIEW:
                    numVisibleNodes += 1

                for child in vertex.parents:
                    if child not in visited:
                        visited.add(child)
                        queue.append(child)
            if numVisibleNodes < MAX_SERVICES_SUPPORTED_IN_TREE_VIEW:
                depth += 1
            maxPossibleDepth += 1

        return {
            'graphs': graphs,
            'totalCount': total,
            'visibleCount': visibleCount,
            'depth': depth,
            'maxPossibleDepth': maxPossibleDepth
        }


class Subgraph(object):
    """
    A subgraph representing a connected graph of vertices in the service tree.
    """

    def __init__(self):
        self.id = str(uuid.uuid4())
        self.vertices = set()
        self.edges = set()

    def populate_from(self, start_vertex, visited):
        """
        Populates this subgraph's vertices and edges data by traversing from
        the given start vertex.

        Also calls self.__assign_node_depth().

        NOTE: The visited set param is modified during this call.

        @type start_vertex: ServiceVertex
        @param start_vertex: the vertex to start traversing from

        @type visited: set
        @param visited: set containing the already visited vertices
        """
        stack = [start_vertex]

        while stack:
            v = stack.pop()

            for child in v.children:
                if child in visited:
                    continue

                self.edges.add((v, child))
                stack.append(child)

            for parent in v.parents:
                if parent in visited:
                    continue

                self.edges.add((parent, v))
                stack.append(parent)

            self.vertices.add(v)
            visited.add(v)

        self.__assign_node_depth()

    def __assign_node_depth(self):
        """
        Private method that assigns node depth to each vertex in the subgraph.

        NOTE: Must only be called from within the populate_from() method.

        If a cycle is detected, all vertices get node_depth = -1.

        If the desired tree from this subgraph
        is a directed graph with arrows flowing "up" from bottom to top
        then "root" vertices would be at the top, and have no outgoing arrows (thus no children).
        NOTE: This is the current expected behavior.

        However, if the desired tree from this subgraph
        is a directed graph with arrows flowing "down" from top to bottom
        then "root" vertices would still be at the top, and have no incoming arrows (thus no parents).
        """
        # for cycles, node_depth = -1 (which is also the default assignment of node_depth in the ServiceVertex
        # constructor)
        if self.has_cycle():
            return

        # find "root" nodes (should have no children) - NOTE: see method description
        roots = set([v for v in self.vertices if len(v.children) == 0])
        if len(roots) == 0:
            print("Every vertex has a child, which means a cycle. self.has_cycle() should have detected this.")
            return

        # do multi-source BFS to assign node_depth to each vertex
        visited = set()
        q = deque()
        depth = 0

        # start with the roots
        for root in roots:
            visited.add(root)
            q.append(root)

        while q:
            level_size = len(q)
            for i in range(level_size):
                curr = q.popleft()
                curr.node_depth = depth

                for v in curr.parents:  # would be curr.children if going "down" - NOTE: see method description
                    if v not in visited:
                        visited.add(v)
                        q.append(v)
            depth += 1

    def has_cycle(self):
        """
        Returns whether this subgraph contains a cycle.

        Algorithm inspired by: https://www.geeksforgeeks.org/detect-cycle-direct-graph-using-colors/

        @rtype: bool
        @return: whether the subgraph contains a cycle
        """
        # Initially put all vertices in an unprocessed state
        unprocessed = self.vertices.copy()

        # Keep track of vertices that haven't been fully processed yet
        processing = set()

        for vertex in self.vertices:
            # Starting from an unprocessed vertex, look for back edges in order to detect cycles
            if vertex in unprocessed and self._dfs_cycle(vertex, unprocessed, processing):
                return True

        return False

    def _dfs_cycle(self, vertex, unprocessed, processing):
        """
        Helper method that performs a DFS travel on a given vertex, detecting cycles.

        @rtype: bool
        @return: whether a cycle's been detected
        """
        if vertex in unprocessed:
            unprocessed.remove(vertex)

        # Mark this vertex as being processed
        processing.add(vertex)

        for child in vertex.children:
            if child in processing or (child in unprocessed and self._dfs_cycle(child, unprocessed, processing)):
                # Found a back edge, cycle detected
                return True

        # This vertex is fully processed, so remove it from a processing state
        processing.remove(vertex)

        return False


class ServiceVertex(object):
    """
    A vertex in the service tree, containing parent and children edge data.
    """

    def __init__(self, id):
        self.id = id
        self.title = ''
        self.children = set()
        self.parents = set()
        self.node_depth = -1

    def _traverse(self, edge_attr, visited):
        """
        Traverses the connected graph through the children of this vertex.

        @type edge_attr: str
        @param edge_attr: the edge attribute name to use for the traversal

        @type visited: set
        @param visited: set containing the already visited vertices
        """
        stack = [self]

        while stack:
            v = stack.pop()
            stack.extend([n for n in getattr(v, edge_attr)
                          if n not in visited])
            visited.add(v)

    def traverse_children(self, visited):
        """
        Traverses the connected graph through the children of this vertex.

        @type visited: set
        @param visited: set containing the already visited vertices
        """
        self._traverse('children', visited)

    def traverse_parents(self, visited):
        """
        Traverses the connected graph through the parents of this vertex.

        @type visited: set
        @param visited: set containing the already visited vertices
        """
        self._traverse('parents', visited)

    @property
    def degree(self):
        """
        Returns the degree of the vertex.

        @rtype: int
        @return: the number of edges connected to the vertex
        """
        return len(self.children) + len(self.parents)

    def __repr__(self):
        return u'[ServiceVertex] %s' % self.id