File diff 9b889983cd7b → 49f356923784
conntrackt/utils.py
Show inline comments
 
# Standard library imports.
 
import re
 
import itertools
 

	
 
# Third-party Python library imports.
 
import palette
 
import pydot
 

	
 
# Django imports.
 
from django.template import Context, loader
 

	
 
# Application imports.
 
import iptables
 
from .models import Communication
 

	
 

	
 
def generate_entity_iptables(entity):
 
@@ -52,3 +58,125 @@ def generate_entity_iptables(entity):
 
    content = "%s%s" % (filter, nat)
 

	
 
    return content
 

	
 

	
 
def get_distinct_colors(count, start=palette.Color("#AE1111")):
 
    """
 
    Generates a number of distinct colours, and returns them as a list. The
 
    colours are generated using the HSL (hue, saturation, lightness) model,
 
    where saturation and lightness is kept the same for all colours, with
 
    differing hue. The hue difference between each subsequent color in the list
 
    is kept the same.
 

	
 
    Arguments:
 

	
 
        count - Total number of colours that should be generated.
 

	
 
        start - First colour that should be taken as a start point. All colours
 
        are generated relative to this colour by increasing the hue. Should be
 
        an instance of palette.Color class. Defaults to RGB colour "#AE1111".
 

	
 
    Return:
 

	
 
        List of distinct palette.Color instances.
 
    """
 

	
 
    # Read the HSL from provided Color.
 
    hue, sat, lum = start.hsl["h"], start.hsl["s"], start.hsl["l"]
 

	
 
    # Calculate the step increase.
 
    step = 1 / float(count)
 

	
 
    # Initiate an empty list that will store the generated colours.
 
    colors = []
 

	
 
    # Generate new colour by increasing the hue as long as we haven't generated
 
    # the requested number of colours.
 
    while len(colors) < count:
 
        colors.append(palette.Color(hsl=(hue, sat, lum)))
 
        hue += step
 

	
 
    return colors
 

	
 

	
 
def generate_project_diagram(project):
 
    """
 
    Generates communication diagram for provided project.
 

	
 
    Arguments:
 

	
 
        project - Project for which the diagram should be generated. Instance of
 
        conntrackt.models.Project class.
 

	
 
    Returns:
 

	
 
        Dot diagram (digraph) representing all of the communications in a
 
        project.
 
    """
 

	
 
    # Set-up the graph object.
 
    graph = pydot.Dot(graph_name=project.name, graph_type="digraph", bgcolor="transparent", nodesep="1.5")
 
    # Set-up defaults for the graph nodes.
 
    graph.set_node_defaults(shape="record")
 

	
 
    # Obtain list of all entities in a project.
 
    entities = project.entity_set.all()
 

	
 
    # Set-up dictinary that will contains clusters of entities belonging to same
 
    # location.
 
    clusters = {}
 

	
 
    # Dictinoary for storing mapping between nodes and colours.
 
    node_colors = {}
 

	
 
    # Get distinct colours, one for each node/entity.
 
    colors = get_distinct_colors(entities.count())
 

	
 
    # Created nodes based on entities, and put them into correct cluster.
 
    for entity in entities:
 

	
 
        # Try to get the existing cluster based on location name.
 
        location = entity.location
 
        cluster_name = location.name.replace(" ", "_").lower()
 
        cluster = clusters.get(cluster_name, None)
 

	
 
        # Set-up a new cluster for location encountered for the first time.
 
        if cluster is None:
 
            cluster = pydot.Cluster(graph_name=cluster_name, label=location.name)
 
            clusters[cluster_name] = cluster
 

	
 
        # Fetch a colour that will be associated with the node/entity.
 
        node_color = colors.pop()
 
        node_colors[entity.id] = node_color.hex
 

	
 
        # Determine whether the node label should be black or white based on brightness of node colour.
 
        node_color_brightness = 1 - (node_color.rgb["r"] * 0.299 + node_color.rgb["g"] * 0.587 + node_color.rgb["b"] * 0.114)
 

	
 
        if node_color_brightness < 0.5:
 
            font_color = "black"
 
        else:
 
            font_color = "white"
 

	
 
        # Finally create the node, and add it to location cluster.
 
        node = pydot.Node(entity.name, style="filled", color=node_color.hex, fontcolor=font_color)
 
        cluster.add_node(node)
 

	
 
    # Add clusters to the graph.
 
    for cluster in clusters.values():
 
        graph.add_subgraph(cluster)
 

	
 
    # Get all project communications.
 
    communications = Communication.objects.filter(source__entity__project=project)
 

	
 
    # Add the edges (lines) representing communications, drawing them with same
 
    # colour as the source node/entity.
 
    for comm in communications:
 
        edge_color = node_colors[comm.source.entity.id]
 

	
 
        label = '"%s:%s"' % (comm.protocol, str(comm.port))
 

	
 
        edge = pydot.Edge(comm.source.entity.name, comm.destination.entity.name, label=label, color=edge_color)
 

	
 
        graph.add_edge(edge)
 

	
 
    return graph