Start of work on visualisation

Simple code using python-gd which will generate a PNG of a set of
switches and some connections between.

Change-Id: I39b31c5e335f5dcdbc27e2b5d04434ceb2107bd6
diff --git a/visualisation/switch.py b/visualisation/switch.py
new file mode 100644
index 0000000..0e6ef27
--- /dev/null
+++ b/visualisation/switch.py
@@ -0,0 +1,402 @@
+#!/usr/bin/env python
+
+import gd, os, cStringIO, urllib2
+
+fontlist = [
+#    '/usr/share/fonts/truetype/freefont/FreeSerif.ttf',
+#    '/usr/share/fonts/truetype/freefont/FreeSans.ttf',
+    '/usr/share/fonts/truetype/inconsolata/Inconsolata.otf',
+    '/usr/share/fonts/truetype/freefont/FreeMono.ttf'
+    ]
+
+fontpath = '.'
+for f in fontlist:
+    if os.path.exists(f):
+        fontpath = fontpath + ':' + os.path.dirname(f)
+        FONT = os.path.abspath(f)
+        break
+
+os.environ["GDFONTPATH"] = "fontpath"
+
+# Define some things
+
+# Default font size
+label_font_size = 12
+
+# Number of pixels to leave around switches, horizontally and
+# vertically
+y_gap = 50
+x_gap = 40
+
+# How big a gap to leave between trunk connections
+trunk_gap = 8
+
+# Basic colour definitions used later
+colour_defs = {}
+colour_defs['black'] = (0, 0, 0)
+colour_defs['white'] = (255, 255, 255)
+colour_defs['purple'] = (255, 0, 255)
+colour_defs['blue'] = (0, 0, 255)
+colour_defs['darkgrey'] = (60, 60, 60)
+colour_defs['yellow'] = (255, 255, 0)
+colour_defs['red'] = (255, 0, 0)
+colour_defs['aqua'] = (0, 255, 255)
+
+pallette = {}
+
+# colours for the background
+pallette['bg_colour'] = "purple"
+pallette['transparent_colour'] = "purple"
+
+# switch colours
+pallette['switch_outline_colour'] = "black"
+pallette['switch_fill_colour'] = "darkgrey"
+pallette['switch_label_colour'] = "white"
+
+# verious sets of port colours, matching the "highlight" options in
+# draw_port()
+port_pallette = {}
+port_pallette['normal'] = {}
+port_pallette['normal']['port_box'] = "white"
+port_pallette['normal']['port_bg'] = "black"
+port_pallette['normal']['port_label'] = "white"
+port_pallette['normal']['trace'] = "black"
+
+port_pallette['trunk'] = {}
+port_pallette['trunk']['port_box'] = "white"
+port_pallette['trunk']['port_bg'] = "blue"
+port_pallette['trunk']['port_label'] = "yellow"
+port_pallette['trunk']['trace'] = "blue"
+
+port_pallette['locked'] = {}
+port_pallette['locked']['port_box'] = "white"
+port_pallette['locked']['port_bg'] = "red"
+port_pallette['locked']['port_label'] = "yellow"
+port_pallette['locked']['trace'] = "red"
+
+port_pallette['VLAN'] = {}
+port_pallette['VLAN']['port_box'] = "white"
+port_pallette['VLAN']['port_bg'] = "aqua"
+port_pallette['VLAN']['port_label'] = "black"
+port_pallette['VLAN']['trace'] = "aqua"
+
+# TODO: make colours configurable, add parsing for /etc/X11/rgb.txt to
+# allow people to use arbitrary names?
+
+class Switch:
+
+    port_width = 0
+    port_height = 0
+    text_width = 0
+    text_height = 0
+    legend_width = 0
+    legend_height = 0
+    legend_text_width = 0
+    legend_text_height = 0
+    label_left = 0
+    label_bot = 0
+    total_width = 0
+    total_height = 0
+    num_ports = 0
+    left = 0
+    top = 0
+    name = None
+    
+    def __init__(self, num_ports, name):
+        self.num_ports = num_ports
+        self.name = name
+        self._calc_port_size()
+        self._calc_switch_size()
+        self._calc_legend_size()
+    
+    def _get_label_size(self, label):
+        tmp_im = gd.image((200, 200))
+        (llx, lly, lrx, lry, urx, ury, ulx, uly) = tmp_im.get_bounding_rect(FONT, label_font_size, 0.0, (10, 100), label)
+
+        width = max(lrx, urx) - min(llx, ulx)
+        height = max(lly, lry) - min(uly, ury)
+
+    #    print "\"%s\" takes %d to %d and %d to %d (%dx%d)" % (string, left, right, top, bottom, width, height)
+        return (width, height)
+
+    def _calc_port_size(self):
+        max_width = 0
+        max_height = 0
+    
+        for value in range (0, 99):
+            (width, height) = self._get_label_size(repr(value))
+            max_width = max(max_width, width)
+            max_height = max(max_height, height)
+
+#        print "got max w %d h %d" % (max_width, max_height)
+    
+        self.port_width = max_width + 6
+        self.port_height = max_height + 6
+        self.text_width = max_width
+        self.text_height = max_height
+
+    def _calc_legend_size(self):
+        max_width = 0
+        max_height = 0
+
+        for value in port_pallette.iterkeys():
+            (width, height) = self._get_label_size(repr(value))
+            max_width = max(max_width, width)
+            max_height = max(max_height, height)
+
+        self.legend_width = max_width + self.port_width + 10
+        self.legend_height = 3 + self.port_height + 3
+        self.legend_text_width = max_width
+        self.legend_text_height = max_height
+        self.legend_total_width = 6 + (len(port_pallette) * self.legend_width)
+
+    def _calc_switch_size(self):
+        (label_width, label_height) = self._get_label_size(self.name)
+        num_ports = self.num_ports
+        # Make sure we have an even number for 2 rows
+        if (self.num_ports & 1):
+            num_ports += 1
+        self.label_left = 3 + (num_ports * self.port_width / 2) + 3
+        self.label_bot = self.port_height - 2
+        self.total_width = self.label_left + label_width + 3
+        self.total_height = 3 + (2 * self.port_height) + 3
+
+    def get_dimensions(self):
+        return (self.total_width, self.total_height)        
+
+    def get_legend_dimensions(self):
+        return (self.legend_total_width, self.legend_height)        
+
+    def draw_switch(self, im, left, top):
+        self.left = left
+        self.top = top
+        
+        ulx = left
+        uly = top
+        lrx = left + self.total_width -1
+        lry = top + self.total_height - 1
+        print "ds: drawing rectangle at (%d,%d) to (%d,%d)" % (ulx, uly, lrx, lry)
+        im.rectangle((left, top), (lrx, lry),
+                     im.colorExact(colour_defs[pallette['switch_outline_colour']]),
+                     im.colorExact(colour_defs[pallette['switch_fill_colour']]))
+
+        llx = left + self.label_left
+        lly = top + self.label_bot
+        print "ds: drawing text at (%d,%d)" % (llx, lly)
+        im.string_ttf(FONT, label_font_size, 0.0, (llx, lly), self.name,
+                      im.colorExact(colour_defs[pallette['switch_label_colour']]))
+
+    def draw_port(self, im, portnum, highlight):
+        if portnum < 1 or portnum > self.num_ports:
+            raise Exception("port number out of range")
+        if highlight not in port_pallette.iterkeys():
+            raise Exception("unknown highlight type \"%s\"" % highlight)
+        
+        box_colour = im.colorExact(colour_defs[port_pallette[highlight]['port_box']])
+        box_bg_colour = im.colorExact(colour_defs[port_pallette[highlight]['port_bg']])
+        text_colour = im.colorExact(colour_defs[port_pallette[highlight]['port_label']])
+
+        if (portnum & 1): # top
+            ulx = self.left + 3 + ((portnum-1) * self.port_width / 2)
+            uly = self.top + 3
+        else:
+            ulx = self.left + 3 + ((portnum-2) * self.port_width / 2)
+            uly = self.top + 3 + self.port_height
+        lrx = ulx + self.port_width - 1
+        lry = uly + self.port_height - 1
+            
+        print "dp: drawing rectangle for port %d at (%d,%d) to (%d,%d)" % (portnum, ulx, uly, lrx, lry)
+        im.rectangle((ulx,uly), (lrx,lry), box_colour, box_bg_colour)
+
+        # centre the text
+        (width, height) = self._get_label_size(repr(portnum))
+        llx = ulx + 3 + (self.text_width - width) / 2
+        lly = uly + self.text_height + 1
+        print "dp: drawing text at (%d,%d)" % (llx, lly)
+        im.string_ttf(FONT, label_font_size, 0.0, (llx, lly), repr(portnum), text_colour)
+
+    def draw_default_ports(self, im):
+        for portnum in range(1, self.num_ports + 1):
+            self.draw_port(im, portnum, "normal")
+
+    def draw_legend(self, im, left, top):
+        
+        lrx = left + self.legend_total_width - 1
+        lry = top + self.legend_height - 1
+        print "ds: drawing rectangle at (%d,%d) to (%d,%d)" % (left, top, lrx, lry)
+        im.rectangle((left, top), (lrx, lry),
+                     im.colorExact(colour_defs[pallette['switch_outline_colour']]),
+                     im.colorExact(colour_defs[pallette['switch_fill_colour']]))
+
+        curr_x = left + 3
+        curr_y = top + 3
+
+        for value in sorted(port_pallette):
+            box_colour = im.colorExact(colour_defs[port_pallette[value]['port_box']])
+            box_bg_colour = im.colorExact(colour_defs[port_pallette[value]['port_bg']])
+            text_colour = im.colorExact(colour_defs[port_pallette[value]['port_label']])
+            lrx = curr_x + self.port_width - 1
+            lry = curr_y + self.port_height - 1
+            print "dp: drawing rectangle for legend %s at (%d,%d) to (%d,%d)" % (value, curr_x, curr_y, lrx, lry)
+            im.rectangle((curr_x,curr_y), (lrx,lry), box_colour, box_bg_colour)
+
+            (width, height) = self._get_label_size("#")
+            llx = curr_x + 3 + (self.text_width - width) / 2
+            lly = curr_y + self.text_height + 1
+            print "dp: drawing text at (%d,%d)" % (llx, lly)
+            im.string_ttf(FONT, label_font_size, 0.0, (llx, lly), "#", text_colour)
+            curr_x += self.port_width
+
+            print "ds: drawing text at (%d,%d)" % (curr_x, lly)
+            im.string_ttf(FONT, label_font_size, 0.0, (curr_x + 3, lly), value,
+                          im.colorExact(colour_defs[pallette['switch_label_colour']]))
+            curr_x += self.legend_text_width + 10
+
+    # Get the (x,y) co-ordinates of the middle-bottom of the port box,
+    # so we can draw a connection to it
+    def get_port_location(self, portnum):
+        if portnum > self.num_ports:
+            raise Exception("port number out of range")
+
+        if (portnum & 1): # top
+            ulx = self.left + 3 + ((portnum-1) * self.port_width / 2)
+            uly = self.top + 3
+            mid_edge = ulx + int((self.port_width / 2))
+            return (mid_edge, uly, True)
+        else:
+            ulx = self.left + 3 + ((portnum-2) * self.port_width / 2)
+            uly = self.top + 3 + self.port_height
+            lry = uly + self.port_height
+            mid_edge = ulx + int((self.port_width / 2))
+            return (mid_edge, lry, False)
+
+    def dump_state(self):
+        print "port_width %d" % self.port_width
+        print "port_height %d" % self.port_height
+        print "text_width %d" % self.text_width
+        print "text_height %d" % self.text_height
+        print "label_left %d" % self.label_left
+        print "label_bot %d" % self.label_bot
+        print "total_width %d" % self.total_width
+        print "total_height %d" % self.total_height
+
+def draw_trunk(im, trunknum, node1, node2, colour):
+    print node1
+    print node2
+    for node in (node1, node2):
+        (x1,y1,top) = node
+        x2 = x1
+        if (top):
+            y2 = y1 - (trunk_gap * (trunknum + 1))
+        else:
+            y2 = y1 + (trunk_gap * (trunknum + 1))
+        # Quick hack - use 2-pixel wide rectangles as thick lines :-)
+        # First line, vertically up/down from the port
+        im.rectangle((x1-1,y1), (x2,y2), im.colorExact(colour_defs[colour]))
+        # Now draw horizontally across to the left margin space
+        x3 = trunk_gap * (trunknum + 1)
+        im.rectangle((x3, y2), (x2,y2+1), im.colorExact(colour_defs[colour]))
+    # Now join up the trunks vertically
+    (x1,y1,top1) = node1
+    if (top1):
+        y1 -= trunk_gap * (trunknum + 1)
+    else:
+        y1 += trunk_gap * (trunknum + 1)
+    (x2,y2,top2) = node2
+    if (top2):
+        y2 -= trunk_gap * (trunknum + 1)
+    else:
+        y2 += trunk_gap * (trunknum + 1)
+    x1 = trunk_gap * (trunknum + 1)
+    im.rectangle((x1, y1), (x1+1,y2), im.colorExact(colour_defs[colour]))
+
+def simple():
+
+    switch = {}
+    trunk = {}
+    size_x = {}
+    size_y = {}
+
+    current_trunk_num = 0
+
+    switch[0] = Switch(48, "lngswitch01")
+    switch[1] = Switch(24, "lngswitch02")
+    switch[2] = Switch(52, "lngswitch03")
+
+    x = 0
+    y = y_gap
+    
+    for i in range (0, 3):
+        (size_x[i], size_y[i]) = switch[i].get_dimensions()
+        x = max(x, size_x[i])
+        y += size_y[i] + y_gap
+
+    # Add space for the legend
+    y = y_gap + y + 10
+    (legend_width, legend_height) = switch[0].get_legend_dimensions()
+    x = max(x, legend_width)
+    x = x_gap + x + x_gap
+     
+    im = gd.image((x,y))
+
+    # Allocate our colours in the image's colour map
+    for key in colour_defs.iterkeys():
+        im.colorAllocate((colour_defs[key][0], colour_defs[key][1], colour_defs[key][2]))
+    
+    im.fill((0,0), im.colorExact(colour_defs[pallette['bg_colour']]))
+    im.colorTransparent(im.colorExact(colour_defs[pallette['transparent_colour']]))
+    im.interlace(0)
+
+    curr_y = y_gap
+    switch[0].draw_switch(im, x_gap, curr_y)
+    switch[0].draw_default_ports(im)
+    switch[0].draw_port(im, 2, "VLAN")
+    switch[0].draw_port(im, 5, "locked")
+    switch[0].draw_port(im, 11, "trunk")
+    switch[0].draw_port(im, 14, "trunk")
+    curr_y += size_y[0] + y_gap
+
+    switch[1].draw_switch(im, x_gap, curr_y)
+    switch[1].draw_default_ports(im)
+    switch[1].draw_port(im, 5, "VLAN")
+    switch[1].draw_port(im, 8, "locked")
+    switch[1].draw_port(im, 13, "trunk")
+    switch[1].draw_port(im, 16, "trunk")
+    curr_y += size_y[2] + y_gap
+
+    switch[2].draw_switch(im, x_gap, curr_y)
+    switch[2].draw_default_ports(im)
+    switch[2].draw_port(im, 1, "trunk")
+    switch[2].draw_port(im, 2, "locked")
+    switch[2].draw_port(im, 14, "trunk")
+    switch[2].draw_port(im, 19, "VLAN")
+    curr_y += size_y[2] + y_gap
+
+    switch[0].draw_legend(im, x_gap, curr_y)
+
+    # Now let's try and draw some trunks!
+    draw_trunk(im, 0,
+               switch[0].get_port_location(11),
+               switch[1].get_port_location(16),
+               port_pallette['trunk']['trace'])
+    draw_trunk(im, 1,
+               switch[1].get_port_location(13),
+               switch[2].get_port_location(1),
+               port_pallette['trunk']['trace'])
+    draw_trunk(im, 2,
+               switch[0].get_port_location(44),
+               switch[2].get_port_location(14),
+               port_pallette['trunk']['trace'])
+               
+    
+    f=open("xx.png","w")
+    im.writePng(f)
+    f.close()
+
+try:
+    FONT
+except NameError:
+    print "no fonts found"
+    sys.exit(1)
+
+simple()