bin/pvl.hosts-graph
author Tero Marttila <terom@paivola.fi>
Tue, 18 Mar 2014 23:53:16 +0200
changeset 400 41dd2a867e0a
parent 399 aadf76a05ec1
child 402 3b6fb4ae56fb
permissions -rwxr-xr-x
pvl.hosts-graph: parse new attribute-values format
#!/usr/bin/env python

"""
    Requirements:
        pydot
"""

import pvl.args
import pvl.hosts

import collections
import logging; log = logging.getLogger('pvl.hosts-graph')
import optparse

COLOR_VLANS = {
    1:      'grey',         # pvl-lan
    2:      'blue',         # pvl-lan2
    3:      'red',          # pvl-san
    4:      'green',        # pvl-veturi
    7:      'orange',       # pvl-ranssi
    8:      'yellow',       # pvl-mgmt
    10:     'brown',        # pvl-public
    100:    'navyblue',     # pvl-test
    103:    'red4',         # pvl-test-san
    104:    'red2',         # pvl-ganeti
    192:    'purple',       # paivola-services
    255:    'magenta',      # pvl-sonera
}

class ParseError (Exception) :
    def __init__ (self, file, line, msg) :
        self.file = file
        self.line = line
        self.msg = msg
  
    def __str__ (self) :
        return "{self.file}:{self.line}: {self.msg}".format(self=self)

def _parse_snmp_part (part) :
    if part.isdigit() :
        return int(part)
    else :
        return part

def _parse_snmp_line (line) :
    for part in line.split() :
        yield _parse_snmp_part(part)

def _load_snmp_data (options, file) :
    """
        Load a data dict generated by pvl.hosts-snmp from a file.

        Yields (host, attr, value)
    """

    host = None
    attr = None
    value = None
    values = None
    
    for idx, line in enumerate(file, 1) :
        indent = 0

        while line.startswith('\t') :
            indent += 1
            line = line[1:]

        line = line.strip()

        if '\t' in line :
            args = line.split('\t', 1)

            line = args.pop(0)
            args = tuple(args)
        else :
            args = None

        if indent == 0 :
            host = line
            attr = None
            value = None
            values = None

        elif indent == 1 :
            if attr and not value and not values :
                raise ParseError(file, line, "[%s] %s: no value" % (host, attr))

            attr = tuple(_parse_snmp_line(line))

            if args :
                value = tuple(tuple(_parse_snmp_line(arg)) for arg in args)
            else :
                value = None
            
            values = None

            if value :
                yield host, attr, value

        elif indent == 2 :
            if not attr :
                raise ParseError(file, line, "[%s] %s: value outside of attr" % (host, attr))
            
            if value :
                raise ParseError(file, line, "[%s] %s: values with value" % (host, attr))

            values = _parse_snmp_part(line)

            yield host, attr, set((values, ))

def load_snmp_data (options, file, hosts) :
    """
        Load snmp data as dict, from given file path, or stdin.
    """

    if file :
        file = open(file)
    else :
        file = sys.stdin

    root = { }

    for host_name, attr, value in _load_snmp_data(options, file) :
        host = hosts[host_name]

        log.info("[%s] %s%s", host, ' '.join(str(a) for a in attr), (': ' + str(value)) if value else '')

        item = root.setdefault(host, { })
        
        for a in attr[:-1] :
            item = item.setdefault(a, {})

        a = attr[-1]
        
        if isinstance(value, set) :
            item.setdefault(a, set()).update(value)

        else :
            item[a] = dict(value)
            
    return root

def apply_graph (options, items, vlans={}) :
    import pydot

    dot = pydot.Dot(graph_name='lldp_hosts', graph_type='digraph',
            # XXX: breaks multi-edges?
            #splines     = 'true',

            sep             = '+25,25',
            overlap         = 'scalexy',

            # only applies to loops
            nodesep     = 0.5,
    )
    dot.set_edge_defaults(
            labeldistance   = 3.0,
            penwidth        = 2.0,
    )

    nodes = { }
    edges = { }
    vlan_colors = { } # { vlan: color }

    for host, local, remote, remote_host in items :
        # src
        src_name = str(host)
        src_label = '"{host.location}"'.format(host=host)

        if src_name in nodes :
            src = nodes[src_name]
        else :
            src = nodes[src_name] = pydot.Node(src_name,
                    label       = src_label,
                    fontsize    = 18,
            )
            dot.add_node(src)
        
        # dst
        if remote_host :
            dst_name = str(remote_host)
            dst_label = '"{host.location}"'.format(host=remote_host)
        else :
            dst_name = remote['chassis'].replace(':', '-')

            # XXX: pydot is not smart enough to quote this
            if remote['sys_name'] :
                dst_label = '"{remote[chassis]} ({remote[sys_name]})"'.format(remote=remote)
            else :
                dst_label = '"{remote[chassis]}"'.format(remote=remote)
        
        if dst_name in nodes :
            dst = nodes[dst_name]
        else :
            dst = nodes[dst_name] = pydot.Node(dst_name,
                    label       = dst_label,
                    fontsize    = 18,
            )
            dot.add_node(dst)

        # edges
        headlabel = '"{remote[port]}"'.format(remote=remote)
        taillabel = '"{local[port]}"'.format(local=local)
        fillcolor = 'black'
        color = 'black'
        untag = tagged = None

        # vlans?
        if vlans and host in vlans and local['port_id'] in vlans[host] :
            untag, tagged = vlans[host][local['port_id']]
            
            log.debug("%s#%s: %s+%s", host, local['port_id'], untag, tagged)
            
            colors = []
            for tag in sorted(tagged) :
                if tag in COLOR_VLANS :
                    colors.append(COLOR_VLANS[tag])
                elif tag in vlan_colors :
                    colors.append(vlan_colors[tag])
                else :
                    color = '/paired12/{count}'.format(count=1+len(vlan_colors))
                    
                    log.info("%s#%s: chosing new vlan %s color %s", host, local['port_id'], tag, color)

                    vlan_colors[tag] = color
                    colors.append(color)

            if not untag :
                pass
            elif untag in COLOR_VLANS :
                fillcolor = COLOR_VLANS[untag]
            elif untag in vlan_colors :
                fillcolor = vlan_colors[untag]
            else :
                color = '/paired12/{count}'.format(count=1+len(vlan_colors))

                log.warn("%s#%s: chosing new vlan %s color %s", host, local['port_id'], untag, color)
                    
                fillcolor = vlan_colors[tag] = color
            
            # first color overrides fillcolor for heads
            if colors and untag :
                color = ':'.join([fillcolor] + colors)
            elif colors :
                color = ':'.join(colors)
            elif fillcolor :
                color = fillcolor

        elif vlans :
            # XXX: this happens when LLDP gives us the LACP ports but the VLANS are on the TRK port
            log.warn("%s#%s: unknown port for vlans: %s", host, local['port_id'], vlans.get(host))

        # edge
        if (src_name, local['port'], dst_name, remote['port']) in edges :
            log.warning("%s:%s <- %s:%s: duplicate", src_name, local['port'], dst_name, remote['port'])
        
        elif (dst_name, remote['port'], src_name, local['port']) in edges :
            log.info("%s <-> %s", src_name, dst_name)
            edge = edges[(dst_name, remote['port'], src_name, local['port'])]

            if edge.get('headlabel') != taillabel :
                log.warn("%s -> %s: local port mismatch: %s vs %s", src_name, dst_name, local['port'], edge.get('headlabel'))
            
            if edge.get('taillabel') != headlabel :
                log.warn("%s -> %s: remote port mismatch: %s vs %s", src_name, dst_name, remote['port'], edge.get('taillabel'))

            if edge.get('fillcolor') != fillcolor :
                log.warn("%s#%s -> %s#%s: remote untag mismatch: %s vs %s", src_name, local['port'], dst_name, remote['port'], fillcolor, edge.get('fillcolor'))

            if edge.get('color') != '"' + color + '"' :
                log.warn("%s#%s -> %s#%s: remote tagged mismatch: %s vs %s", src_name, local['port'], dst_name, remote['port'], color, edge.get('color'))

            # mark as bidirectional
            edges[(src_name, local['port'], dst_name, remote['port'])] = edge
            
            edge.set('dir', 'both' if untag else 'none')

            # set second color for tail
            if untag :
                edge.set('color', '"{headcolor}:{tailcolor}{tagcolors}"'.format(
                    headcolor   = edge.get('fillcolor'),
                    tailcolor   = fillcolor,
                    tagcolors   = ':' + ':'.join(colors) if colors else '',
                ))

        else :
            edge = edges[(src_name, local['port'], dst_name, remote['port'])] = pydot.Edge(src, dst,
                    dir         = 'forward' if untag else 'none',
                    headlabel   = headlabel,
                    taillabel   = taillabel,
                    color       = '"{color}"'.format(color=color),
                    fillcolor   = fillcolor,
            )

            dot.add_edge(edge)

    if options.graph_dot :
        dot.write(options.graph_dot)

def main (argv) :
    """
        Graph network
    """

    parser = optparse.OptionParser(main.__doc__)
    parser.add_option_group(pvl.args.parser(parser))
    parser.add_option_group(pvl.hosts.optparser(parser))

    parser.add_option('--snmp-data', metavar='FILE', default=None,
            help="Load snmp data from FILE")

    parser.add_option('--graph-dot', metavar='FILE',
            help="Output .dot graph data to file")

    parser.add_option('--graph-vlans', action='store_true', dest='vlans', 
            help="Graph all VLANs")

    parser.add_option('--no-vlans', action='store_false', dest='vlans',
            help="Do not color VLANs")

    # input
    options, args = parser.parse_args(argv[1:])
    pvl.args.apply(options)
    
    # load hosts for correlation
    hosts = dict((str(host), host) for host in pvl.hosts.apply(options, args))

    # load raw snmp data
    data = load_snmp_data(options, options.snmp_data, hosts)

    import pprint; pprint.pprint(data)

    # TODO: process data into graph
    
    # TODO: generate graph
    #apply_graph(options, items, vlans if options.vlans else None)

    return 0

if __name__ == '__main__':
    pvl.args.main(main)