visualisation/graphics.py

#! /usr/bin/python

#  Copyright 2015 Linaro Limited
#
#  This program 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 2 of the License, or
#  (at your option) any later version.
#
#  This program 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 this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#
#  Visualisation graphics module for VLANd
#
#  This code uses python-gd to generate graphics ready for insertion
#  into our web interface. Example code in the self-test at the
#  bottom.

import gd, os, sys

if __name__ == '__main__':
    vlandpath = os.path.abspath(os.path.normpath(os.path.dirname(sys.argv[0])))
    sys.path.insert(0, vlandpath)
    sys.path.insert(0, "%s/.." % vlandpath)

from errors import InputError

class Graphics:
    """ Code and config for the visualisation graphics module """

    font = None

    # Default font size for the small labels
    small_font_size = 12

    # And the size for the top-level label
    label_font_size = 24

    # Size in pixels of that font, calculated later
    twocharwidth = 0
    charheight = 0

    # How big a gap to leave between trunk connections
    trunk_gap = 8

    # Details of the legend
    legend_width = 0
    legend_height = 0
    legend_text_width = 0
    legend_text_height = 0
    legend_total_width = 0
    legend_box_width = 0
    legend_box_height = 0

    # 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'
    pallette['graphic_label_colour'] = 'black'

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

    im = None

    # TODO: make colours configurable, add maybe parsing for
    # /etc/X11/rgb.txt to allow people to use arbitrary names?

    # Choose a font for our graphics to use. Pass in a list of fonts
    # to be tried, in priority order.
    def set_font(self, fontlist):
        for font in fontlist:
            if os.path.exists(font):
                self.font = os.path.abspath(font)
                break

        # Work out how big we need to be for the biggest possible text
        # in a 2-digit number. Grotty, but we need to know this later.
        for value in range (0, 100):
            (width, height) = self.get_label_size(repr(value), self.small_font_size)
            self.twocharwidth = max(self.twocharwidth, width)
            self.charheight = max(self.charheight, height)

        # Now we can also calulate other stuff
        self._calc_legend_size()

    # Create a canvas and set things up ready for use
    def create_canvas(self, x, y):
        im = gd.image((x, y))

        # Allocate our colours in the image's colour map
        for key in self.colour_defs.iterkeys():
            im.colorAllocate((self.colour_defs[key][0],
                              self.colour_defs[key][1],
                              self.colour_defs[key][2]))

        im.fill((0,0), im.colorExact(self.colour_defs[self.pallette['bg_colour']]))
        im.colorTransparent(im.colorExact(self.colour_defs[self.pallette['transparent_colour']]))
        im.interlace(0)
        self.im = im

    # Using our selected font, what dimensions will a particular piece
    # of text take?
    def get_label_size(self, label, font_size):
        tmp_im = gd.image((200, 200))
        (llx, lly, lrx, lry, urx, ury, ulx, uly) = tmp_im.get_bounding_rect(self.font,
                                                                            font_size,
                                                                            0.0,
                                                                            (10, 100), label)
        width = max(lrx, urx) - min(llx, ulx)
        height = max(lly, lry) - min(uly, ury)
        return (width, height)

    # Draw a trunk connection between two ports
    #
    # Ports are defined as (ulx,uly),(lrx,lry), top): x, y
    # co-ordinates of UL and LR corners, and whether the port is on
    # the top or bottom row of a switch, i.e. does the wire come up or
    # down when it leaves the port.
    def draw_trunk(self, trunknum, node1, node2, colour):
        for node in (node1, node2):
            ((ulx,uly),(lrx,lry),top) = node

            # Work out the co-ordinates for a line vertically up or
            # down from the edge of the port
            x1 = int((ulx + lrx) / 2)
            x2 = x1
            if (top):
                y1 = uly
                y2 = y1 - (self.trunk_gap * (trunknum + 1))
            else:
                y1 = lry
                y2 = y1 + (self.trunk_gap * (trunknum + 1))
            # Quick hack - use 2-pixel wide rectangles as thick lines :-)
            # First line, vertically up/down from the port
            self.im.rectangle((x1-1,y1), (x2,y2), self.im.colorExact(self.colour_defs[colour]))
            # Now draw horizontally across to the left margin space
            x3 = self.trunk_gap * (trunknum + 1)
            self.im.rectangle((x3, y2), (x2,y2+1), self.im.colorExact(self.colour_defs[colour]))

        # Now join up the trunks vertically
        ((ulx1,uly1),(lrx1,lry1),top1) = node1
        if (top1):
            y1 = uly1 - self.trunk_gap * (trunknum + 1)
        else:
            y1 = lry1 + self.trunk_gap * (trunknum + 1)
        ((ulx2,uly2),(lrx2,lry2),top2) = node2
        if (top2):
            y2 = uly2 - self.trunk_gap * (trunknum + 1)
        else:
            y2 = lry2 + self.trunk_gap * (trunknum + 1)
        x1 = self.trunk_gap * (trunknum + 1)
        self.im.rectangle((x1, y1), (x1+1,y2), self.im.colorExact(self.colour_defs[colour]))

    # How big is the legend?
    def _calc_legend_size(self):
        max_width = 0
        max_height = 0

        for value in self.port_pallette.iterkeys():
            (width, height) = self.get_label_size(value, self.small_font_size)
            max_width = max(max_width, width)
            max_height = max(max_height, height)

        (width, height) = self.get_label_size('##', self.small_font_size)
        self.legend_box_width = width + 6
        self.legend_box_height = height + 6
        self.legend_width = max_width + self.legend_box_width + 10
        self.legend_height = 3 + self.legend_box_height + 3
        self.legend_text_width = max_width
        self.legend_text_height = max_height
        self.legend_total_width = 6 + (len(self.port_pallette) * self.legend_width)

    # Return the legend dimensions
    def get_legend_dimensions(self):
        return (self.legend_total_width, self.legend_height)

    # Draw the legend using (left, top) as the top left corner
    def draw_legend(self, left, top):
        lrx = left + self.legend_total_width - 1
        lry = top + self.legend_height - 1
        self.im.rectangle((left, top), (lrx, lry),
                        self.im.colorExact(self.colour_defs[self.pallette['switch_outline_colour']]),
                        self.im.colorExact(self.colour_defs[self.pallette['switch_fill_colour']]))
        curr_x = left + 3
        curr_y = top + 3

        for value in sorted(self.port_pallette):
            box_colour = self.im.colorExact(self.colour_defs[self.port_pallette[value]['port_box']])
            box_bg_colour = self.im.colorExact(self.colour_defs[self.port_pallette[value]['port_bg']])
            text_colour = self.im.colorExact(self.colour_defs[self.port_pallette[value]['port_label']])
            lrx = curr_x + self.legend_box_width - 1
            lry = curr_y + self.legend_box_height - 1
            self.im.rectangle((curr_x,curr_y), (lrx,lry), box_colour, box_bg_colour)

            llx = curr_x + 4
            lly = curr_y + self.legend_box_height - 4
            self.im.string_ttf(self.font, self.small_font_size, 0.0, (llx, lly), '##', text_colour)
            curr_x += self.legend_box_width
            self.im.string_ttf(self.font, self.small_font_size, 0.0, (curr_x + 3, lly), value,
                               self.im.colorExact(self.colour_defs[self.pallette['switch_label_colour']]))
            curr_x += self.legend_text_width + 10

    # Draw the graphic's label using (left, top) as the top left
    # corner with a box around
    def draw_label(self, left, top, label, gap):
        box_colour = self.im.colorExact(self.colour_defs[self.pallette['switch_label_colour']])
        box_bg_colour = self.im.colorExact(self.colour_defs[self.pallette['switch_fill_colour']])
        text_colour = self.im.colorExact(self.colour_defs[self.pallette['switch_label_colour']])
        (width, height) = self.get_label_size(label, self.label_font_size)
        curr_x = left
        curr_y = top
        lrx = curr_x + width + gap
        lry = curr_y + height + 20
        self.im.rectangle((curr_x,curr_y), (lrx,lry), box_colour, box_bg_colour)
        curr_x = left + 10
        curr_y = top + height + 6
        self.im.string_ttf(self.font, self.label_font_size, 0.0, (curr_x, curr_y), label, text_colour)


class Switch:
    """ Code and config for dealing with a switch """
    port_width = 0
    port_height = 0
    text_width = 0
    text_height = 0
    label_left = 0
    label_bot = 0
    total_width = 0
    total_height = 0
    num_ports = 0
    left = None
    top = None
    name = None

    # Set up a new switch instance; calculate all the sizes so we can
    # size our canvas
    def __init__(self, g, num_ports, name):
        self.num_ports = num_ports
        self.name = name
        self._calc_port_size(g)
        self._calc_switch_size(g)

    # How big is a port and the text within it?
    def _calc_port_size(self, g):
        self.text_width = g.twocharwidth
        self.text_height = g.charheight
        # Leave enough space around the text for a nice clear box
        self.port_width = self.text_width + 6
        self.port_height = self.text_height + 6

    # How big is the full switch, including all the ports and the
    # switch name label?
    def _calc_switch_size(self, g):
        (label_width, label_height) = g.get_label_size(self.name, g.small_font_size)
        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 + max(label_height, (2 * self.port_height)) + 3

    # Return the switch dimensions
    def get_dimensions(self):
        return (self.total_width, self.total_height)

    # Draw the basic switch outline and label using (left, top) as the
    # top left corner. The switch object will remember this origin for
    # later use when drawing ports.
    def draw_switch(self, g, left, top):
        self.left = left
        self.top = top
        lrx = left + self.total_width -1
        lry = top + self.total_height - 1
        g.im.rectangle((left, top), (lrx, lry),
                       g.im.colorExact(g.colour_defs[g.pallette['switch_outline_colour']]),
                       g.im.colorExact(g.colour_defs[g.pallette['switch_fill_colour']]))
        llx = left + self.label_left
        lly = top + self.label_bot
        g.im.string_ttf(g.font, g.small_font_size, 0.0, (llx, lly), self.name,
                        g.im.colorExact(g.colour_defs[g.pallette['switch_label_colour']]))

    # Draw a port inside the switch, using a specified colour scheme
    # to denote its type. The switch outline must have been drawn
    # first, for its origin to be set.
    def draw_port(self, g, portnum, highlight):
        if portnum < 1 or portnum > self.num_ports:
            raise InputError('port number out of range')
        if not self.left or not self.top:
            raise InputError('cannot draw ports before switch is drawn')
        if highlight not in g.port_pallette.iterkeys():
            raise InputError('unknown highlight type \"%s\"' % highlight)

        box_colour = g.im.colorExact(g.colour_defs[g.port_pallette[highlight]['port_box']])
        box_bg_colour = g.im.colorExact(g.colour_defs[g.port_pallette[highlight]['port_bg']])
        text_colour = g.im.colorExact(g.colour_defs[g.port_pallette[highlight]['port_label']])

        if (portnum & 1): # odd port number, so top row
            ulx = self.left + 3 + ((portnum-1) * self.port_width / 2)
            uly = self.top + 3
        else: # bottom row
            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
        g.im.rectangle((ulx,uly), (lrx,lry), box_colour, box_bg_colour)

        # centre the text
        (width, height) = g.get_label_size(repr(portnum), g.small_font_size)
        llx = ulx + 3 + (self.text_width - width) / 2
        lly = uly + max(height, self.text_height) + 1
        g.im.string_ttf(g.font, g.small_font_size,
                        0.0, (llx, lly), repr(portnum), text_colour)

    # Quick helper: draw all the ports for a switch in the default
    # colour scheme.
    def draw_default_ports(self, g):
        for portnum in range(1, self.num_ports + 1):
            self.draw_port(g, portnum, 'normal')

    # Get the (x,y) co-ordinates of the UL and LR edges of the port
    # box, and if it's upper row. This lets us so useful things such
    # as draw a connection to that point for a trunk.
    def get_port_location(self, portnum):
        if portnum > self.num_ports:
            raise InputError('port number out of range')

        if (portnum & 1): # odd port number, so top row
            ulx = self.left + 3 + ((portnum-1) * self.port_width / 2)
            uly = self.top
            lrx = ulx + self.port_width
            lry = uly + self.port_height
            return ((ulx,uly), (lrx,lry), True)
        else: # bottom row
            ulx = self.left + 3 + ((portnum-2) * self.port_width / 2)
            uly = self.top + 3 + self.port_height
            lrx = ulx + self.port_width
            lry = uly + self.port_height
            return ((ulx,uly), (lrx,lry), False)

    # Debug: print some of the state of the switch object
    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

# Test harness - generate a PNG using fake data
if __name__ == '__main__':
    gim = Graphics()
    gim.set_font(['/usr/share/fonts/truetype/inconsolata/Inconsolata.otf',
                  '/usr/share/fonts/truetype/freefont/FreeMono.ttf'])
    try:
        gim.font
    except NameError:
        print 'no fonts found'
        sys.exit(1)

    switch = {}
    size_x = {}
    size_y = {}
    switch[0] = Switch(gim, 48, 'lngswitch01')
    switch[1] = Switch(gim, 24, 'lngswitch02')
    switch[2] = Switch(gim, 52, 'lngswitch03')
    label = "VLAN 4jj"

    # Need to set gaps big enough for the number of trunks, at least.
    num_trunks = 3
    y_gap = max(20, 15 * num_trunks)
    x_gap = max(20, 15 * num_trunks)

    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 and the label
    (legend_width, legend_height) = gim.get_legend_dimensions()
    (label_width, label_height) = gim.get_label_size(label, gim.label_font_size)

    x = max(x, legend_width + 2*x_gap + label_width)
    x = x_gap + x + x_gap
    y = y + max(legend_height + y_gap, label_height)

    gim.create_canvas(x, y)

    curr_y = y_gap
    switch[0].draw_switch(gim, x_gap, curr_y)
    switch[0].draw_default_ports(gim)
    switch[0].draw_port(gim, 2, 'VLAN')
    switch[0].draw_port(gim, 5, 'locked')
    switch[0].draw_port(gim, 11, 'trunk')
    switch[0].draw_port(gim, 44, 'trunk')
    curr_y += size_y[0] + y_gap

    switch[1].draw_switch(gim, x_gap, curr_y)
    switch[1].draw_default_ports(gim)
    switch[1].draw_port(gim, 5, 'VLAN')
    switch[1].draw_port(gim, 8, 'locked')
    switch[1].draw_port(gim, 13, 'trunk')
    switch[1].draw_port(gim, 16, 'trunk')
    curr_y += size_y[2] + y_gap

    switch[2].draw_switch(gim, x_gap, curr_y)
    switch[2].draw_default_ports(gim)
    switch[2].draw_port(gim, 1, 'trunk')
    switch[2].draw_port(gim, 2, 'locked')
    switch[2].draw_port(gim, 14, 'trunk')
    switch[2].draw_port(gim, 19, 'VLAN')
    curr_y += size_y[2] + y_gap

    # Now let's try and draw some trunks!
    gim.draw_trunk(0,
                 switch[0].get_port_location(11),
                 switch[1].get_port_location(16),
                 gim.port_pallette['trunk']['trace'])
    gim.draw_trunk(1,
                 switch[1].get_port_location(13),
                 switch[2].get_port_location(1),
                 gim.port_pallette['trunk']['trace'])
    gim.draw_trunk(2,
                 switch[0].get_port_location(44),
                 switch[2].get_port_location(14),
                 gim.port_pallette['trunk']['trace'])

    gim.draw_legend(x_gap, curr_y)
    gim.draw_label(x - label_width - 2*x_gap, curr_y, label, int(x_gap / 2))

    f=open('xx.png','w')
    gim.im.writePng(f)
    f.close()
    print 'Test graphic written to xx.png'