|
|
# 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
|