Files
@ 2d83b9633ce7
Branch filter:
Location: conntrackt/conntrackt/utils.py
2d83b9633ce7
7.4 KiB
text/x-python
CONNT-20: Moved out the formatter function to utils. Updated docs. Styling fixes. Added tests for the new mixin. Added tests for the formatter function. Switched to different import of models in utils in order to avoid circular imports.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 | # -*- coding: utf-8 -*-
#
# Copyright (C) 2013 Branko Majic
#
# This file is part of Django Conntrackt.
#
# Django Conntrackt is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# Django Conntrackt is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# Django Conntrackt. If not, see <http://www.gnu.org/licenses/>.
#
# Standard library imports.
import re
import itertools
# Third-party Python library imports.
import palette
import pydot
# Django imports.
from django.template import Context, loader
from django.utils.html import format_html
from django.utils.text import capfirst
# Application imports.
import iptables
import models
def generate_entity_iptables(entity):
"""
Generates full iptables rules for the supplied entity. The generated rules
can be fed directly to the iptables-restore utility.
Arguments:
entity - An Entity instance for which the iptables rules should be
generated.
Returns:
String containing the iptables rules for entity.
"""
# Fetch list of incoming communications.
incoming = entity.incoming_communications()
# Set-up the nat table.
nat = iptables.Table("nat")
for chain in ("PREROUTING", "INPUT", "OUTPUT", "POSTROUTING"):
nat.add_chain(iptables.Chain(chain, "ACCEPT"))
# Set-up the filter table INPUT chain.
filter = iptables.Table("filter")
input = iptables.Chain("INPUT", "DROP")
input.add_rule(iptables.LoopbackRule())
input.add_rule(iptables.RelatedRule())
for communication in incoming:
source = "%s/%s" % (communication.source.address, communication.source.netmask)
destination = "%s/%s" % (communication.destination.address, communication.destination.netmask)
input.add_rule(iptables.Rule(source, destination, communication.protocol, communication.port, communication.description))
filter.add_chain(input)
# Set-up empty chains.
filter.add_chain(iptables.Chain("OUTPUT", "ACCEPT"))
filter.add_chain(iptables.Chain("FORWARD", "DROP"))
# Construct the iptables file using the two tables.
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 = models.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
def list_formatter_callback(obj):
"""
Creates model object representation in format:
MODEL_NAME: OBJECT_REPRESENTATION
If passed object has a callable get_absolute_url method, the
instance representation will be surrouned by an HTML anchor
(<a></a>) where target is set to value of the get_absolute_url()
method call.
Arguments:
obj - Model object whose representation should be returned.
Returns:
String represenation of passed model object.
"""
try:
return format_html('<strong>{0}</strong>: <a href="{1}">{2}</a>', capfirst(obj._meta.verbose_name), obj.get_absolute_url(), str(obj))
except AttributeError:
return format_html('<strong>{0}</strong>: {1}', capfirst(obj._meta.verbose_name), str(obj))
|