misp-galaxy/tools/mkdocs/modules/universe.py

from modules.galaxy import Galaxy
from modules.cluster import Cluster

from collections import defaultdict, deque


class Universe:
    def __init__(self, add_inbound_relationship=False):
        self.galaxies = {}  # Maps galaxy_name to Galaxy objects
        self.add_inbound_relationship = add_inbound_relationship

    def add_galaxy(self, galaxy_name, json_file_name, authors, description):
        if galaxy_name not in self.galaxies:
            self.galaxies[galaxy_name] = Galaxy(galaxy_name=galaxy_name, json_file_name=json_file_name, authors=authors, description=description)

    def add_cluster(self, galaxy_name, uuid, description, value, meta):
        if galaxy_name in self.galaxies:
            self.galaxies[galaxy_name].add_cluster(uuid=uuid, description=description, value=value, meta=meta)

    def define_relationship(self, cluster_a_id, cluster_b_id):
        cluster_a = None
        cluster_b = None

        # Search for Cluster A and Cluster B in all galaxies
        for galaxy in self.galaxies.values():
            if cluster_a_id in galaxy.clusters:
                cluster_a = galaxy.clusters[cluster_a_id]
            if cluster_b_id in galaxy.clusters:
                cluster_b = galaxy.clusters[cluster_b_id]
            if cluster_a and cluster_b:  # Both clusters found
                break

        # If both clusters are found, define the relationship
        if cluster_a and cluster_b:
            cluster_a.add_outbound_relationship(cluster_b)
            cluster_b.add_inbound_relationship(cluster_a)
        else:
            # If Cluster B is not found, create a private cluster relationship for Cluster A
            if cluster_a:
                private_cluster = Cluster(uuid=cluster_b_id, galaxy=None, description=None, value="Private Cluster", meta=None)
                cluster_a.add_outbound_relationship(private_cluster)
            else:
                raise ValueError(f"Cluster {cluster_a} not found in any galaxy")
    
    def get_relationships_with_levels(self, start_cluster):

        def bfs_with_undirected_relationships(start_cluster):
            visited = set()  # Tracks whether a cluster has been visited
            relationships = defaultdict(lambda: float('inf'))  # Tracks the lowest level for each cluster pair

            queue = deque([(start_cluster, 0)])  # Queue of (cluster, level)

            while queue:
                current_cluster, level = queue.popleft()
                if current_cluster not in visited:
                    visited.add(current_cluster)

                    # Process all relationships regardless of direction
                    if self.add_inbound_relationship:
                        neighbors = current_cluster.outbound_relationships.union(current_cluster.inbound_relationships)
                    else:
                        neighbors = current_cluster.outbound_relationships
                    for neighbor in neighbors:
                        link = frozenset([current_cluster, neighbor])
                        if level + 1 < relationships[link]:
                            relationships[link] = level + 1
                            if neighbor not in visited and neighbor.value != "Private Cluster":
                                queue.append((neighbor, level + 1))
            
            # count = 0
            # Convert the defaultdict to a list of tuples, ignoring direction
            processed_relationships = []
            for link, lvl in relationships.items():
                # Extract clusters from the frozenset; direction is irrelevant
                clusters = list(link)
                if len(clusters) != 2:
                    # count += 1
                    continue
                
                # Arbitrarily choose the first cluster as 'source' for consistency
                if clusters[0].value == "Private Cluster":
                    processed_relationships.append((clusters[1], clusters[0], lvl))
                else:
                
                    processed_relationships.append((clusters[0], clusters[1], lvl))
                # except:
                #     processed_relationships.append((clusters[0], clusters[0], lvl)) # This is wrong just for testing!!!

            # print(f"Count: {count}")
            return processed_relationships

        return bfs_with_undirected_relationships(start_cluster)