pvl.hosts-graph: split out graph-building for dot-building
authorTero Marttila <terom@paivola.fi>
Sun, 30 Mar 2014 14:21:57 +0300
changeset 405 97b436f9363a
parent 404 78e3d83135ab
child 406 92a4de88b86f
pvl.hosts-graph: split out graph-building for dot-building
bin/pvl.hosts-graph
--- a/bin/pvl.hosts-graph	Wed Mar 19 01:08:56 2014 +0200
+++ b/bin/pvl.hosts-graph	Sun Mar 30 14:21:57 2014 +0300
@@ -111,9 +111,9 @@
         host = hosts.get(host_name)
         
         if value :
-            log.info("[%s] %s: %s", host, ' '.join(str(a) for a in attr), value)
+            log.debug("[%s] %s: %s", host, ' '.join(str(a) for a in attr), value)
         else :
-            log.info("[%s] %s", host, ' '.join(str(a) for a in attr),)
+            log.debug("[%s] %s", host, ' '.join(str(a) for a in attr),)
 
         item = root.setdefault(host, { })
         
@@ -136,78 +136,170 @@
             
     return root
 
-def search_vlans (snmp) :
-    for host in snmp :
-        ports = set()
-        vlans_untagged = { }
-        vlans_tagged = collections.defaultdict(set)
+def host_vlans (host, host_vlans) :
+    """
+        {vlan: { tagged/untagged: [port] } } -> (port, (untag, [tag])).
+    """
 
-        if 'vlan' not in snmp[host] :
-            continue
+    ports = set()
+    vlans_untagged = { }
+    vlans_tagged = collections.defaultdict(set)
 
-        for vlan in snmp[host]['vlan'] :
-            for port in snmp[host]['vlan'][vlan].get('tagged', ()) :
-                ports.add(port)
-                vlans_tagged[port].add(vlan)
+    for vlan, vlan_attrs in host_vlans.iteritems() :
+        for port in vlan_attrs.get('tagged', ()) :
+            ports.add(port)
+            vlans_tagged[port].add(vlan)
+        
+        for port in vlan_attrs.get('untagged', ()) :
+            ports.add(port)
+            vlans_untagged[port] = vlan
+    
+    for port in ports :
+        untag = vlans_untagged.get(port)
+        tagged = vlans_tagged.get(port, ())
+
+        log.debug("%s: %s: untag=%s tag=%s", host, port, untag, tagged)
+        
+        yield port, (untag, tagged)
+
+def build_graph (snmp, hosts) :
+    """
+        Combine given snmp data and { host: Host } into
+            { node: label }
+            { (remote, remote_port, local_port, local): (local_untag, tagged, remote_untag) }
+    """
+
+    nodes = { }
+    links = { }
+    
+    # first scan: lldp hosts
+    for host, host_attrs in snmp.iteritems() :
+        lldp = host_attrs.get('lldp')
+
+        if lldp :
+            lldp_local = lldp['local']
             
-            for port in snmp[host]['vlan'][vlan].get('untagged', ()) :
-                ports.add(port)
-                vlans_untagged[port] = vlan
-        
-        out = { }
-        for port in ports :
-            untag = vlans_untagged.get(port)
-            tagged = vlans_tagged.get(port, ())
-
-            log.info("%s: %s: untag=%s tag=%s", host, port, untag, tagged)
-
-            out[port] = (untag, tagged)
-
-        yield host, out
-
-def search_lldp (snmp) :
-    lldp_hosts = {} 
-
-    # first pass to resolve chassis ID's
+            nodes[lldp_local['chassis']] = host.location
+    
+    # second scan: lldp remotes
     for host, host_attrs in snmp.iteritems() :
         lldp = host_attrs.get('lldp')
 
         if not lldp :
             continue
 
-        lldp_local = lldp.get('local')
-
-        if not lldp_local :
-            continue
-
-        lldp_hosts[lldp_local['chassis']] = host
-
-    # second pass to resolve remote
-    for host, host_attrs in snmp.iteritems() :
-        lldp = host_attrs.get('lldp')
-
-        if not lldp :
-            continue
-
-        local_attrs = lldp.get('local')
+        local = lldp['local']['chassis']
 
-        for port, port_attrs in snmp[host]['lldp']['port'].iteritems() :
-            port_local_attrs = port_attrs['local']
-
-            for chassis, remote_attrs in port_attrs['remote'].iteritems() :
-                remote_host = lldp_hosts.get(chassis)
+        if 'vlan' in host_attrs :
+            vlans = dict(host_vlans(host, host_attrs['vlan']))
+        else :
+            vlans = None
+        
+        for port, port_attrs in lldp.get('port', { }).iteritems() :
+            local_port = port_attrs['local']['port']
 
-                local = merge(local_attrs, port_local_attrs, port_id=port)
-                remote = merge(remote_attrs, chassis=chassis)
-
-                log.info("%s @ %s %s : %s %s @ %s", host, host.location, local, remote, remote_host.location if remote_host else '', remote_host)
+            for remote, remote_attrs in port_attrs['remote'].iteritems() :
+                if remote not in nodes :
+                    nodes[remote] = remote_attrs['sys_name']
 
-                yield host, local, remote, remote_host
+                remote_port = remote_attrs['port']
+                
+                # local vlans
+                if vlans :
+                    port_vlans = vlans.get(port)
+                else :
+                    port_vlans = None
 
-def apply_graph (options, items, vlans={}) :
-    import pydot
+                if port_vlans :
+                    local_untag, local_tagged = port_vlans
+                
+                # bidirectional mappings
+                forward = (local, local_port, remote_port, remote)
+                reverse = (remote, remote_port, local_port, local)
 
-    dot = pydot.Dot(graph_name='lldp_hosts', graph_type='digraph',
+                if reverse not in links :
+                    links[forward] = (local_untag, local_tagged, None)
+                else :
+                    remote_untag, remote_tagged, _ = links[reverse]
+
+                    # merge
+                    if remote_untag != local_untag :
+                        log.warning("%s:%s untag %s <=> %s untag %s:%s",
+                                nodes[local], local_port, local_untag,
+                                remote_untag, nodes[remote], remote_port
+                        )
+
+                    if remote_tagged != local_tagged :
+                        log.warning("%s:%s tagged %s <-> %s tagged %s:%s",
+                                nodes[local], local_port, ':'.join(str(x) for x in sorted(local_tagged)),
+                                ':'.join(str(x) for x in sorted(remote_tagged)), nodes[remote], remote_port
+                        )
+
+                    links[reverse] = (remote_untag, remote_tagged, local_untag)
+
+    return nodes, links
+
+class GraphVlans (object) :
+    """
+        Maintain vlan -> dot style/color mappings
+    """
+
+    SERIES = 'paired12'
+    NONE = 'black'
+
+    def __init__ (self, vlans=None) :
+        if vlans :
+            self.vlans = dict(vlans)
+        else :
+            self.vlans = { }
+    
+    def color (self, vlan) :
+        if vlan in self.vlans :
+            return self.vlans[vlan]
+        
+        # alloc
+        color = '/{series}/{index}'.format(series=self.SERIES, index=len(self.vlans) + 1)
+
+        self.vlans[vlan] = color
+
+        return color
+
+def dot_quote (value) :
+    """
+        Quote a dot value.
+    """
+
+    return '"{value}"'.format(value=value)
+
+def dot (*line, **attrs) :
+    """
+        Build dot-syntax:
+            *line {
+                *line [**attrs];
+            }
+    """
+
+    if line and attrs :
+        return ''.join(('\t', ' '.join(str(x) for x in line), ' [',
+            ', '.join('{name}="{value}"'.format(name=name, value=value) for name, value in attrs.iteritems()),
+        ']', ';'))
+    elif line :
+        return ' '.join(line) + ' {'
+    else :
+        return '}'
+
+def build_dot (options, nodes, links, type='digraph', vlans=None) :
+    """
+        Construct a dot description of the given node/links graph.
+    """
+
+    if vlans is True :
+        vlans = { }
+
+    yield dot(type, 'verkko')
+
+    # defaults
+    yield dot('graph',
             # XXX: breaks multi-edges?
             #splines     = 'true',
 
@@ -217,149 +309,60 @@
             # only applies to loops
             nodesep     = 0.5,
     )
-    dot.set_edge_defaults(
-            labeldistance   = 3.0,
-            penwidth        = 2.0,
+    yield dot('edge',
+        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)
+    yield dot('node',
+        fontsize        = 18,
+    )
+    
+    # nodes
+    for node, node_label in nodes.iteritems() :
+        yield dot(dot_quote(node), label=node_label)
 
-                    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'])]
+    # links
+    for (local, local_port, remote_port, remote), (local_untag, tagged, remote_untag) in links.iteritems() :
+        if vlans :
+            head_color = vlans.color(local_untag) if local_untag else None
+            tail_color = vlans.color(remote_untag) if remote_untag else None
+            line_colors = [vlans.color(tag) for tag in sorted(tagged)]
+        else :
+            head_color = GraphVlans.NONE if local_untag else None
+            tail_color = GraphVlans.NONE if remote_untag else None
+            line_colors = []
 
-            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')
+        if head_color and tail_color :
+            dir = 'both'
+            colors = [head_color, tail_color] + line_colors
+        elif head_color :
+            dir = 'forward'
+            colors = [head_color] + line_colors
+        elif tail_color :
+            dir = 'back'
+            colors = [vlans.NONE, tail_color] + line_colors
+        else :
+            dir = 'none'
+            colors = line_colors
 
-            # 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 '',
-                ))
+        yield dot(dot_quote(local), '->', dot_quote(remote),
+            taillabel   = local_port,
+            headlabel   = remote_port,
+            dir         = dir,
 
-        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,
-            )
+            fillcolor   = 'black',
+            color       = ':'.join(colors),
+        )
 
-            dot.add_edge(edge)
+    yield dot()
 
-    if options.graph_dot :
-        dot.write(options.graph_dot)
+def apply_dot (options, file, dot) :
+    """
+        Output dot file for given graphbits
+    """
+
+    for line in dot :
+        file.write(line + '\n')
 
 def main (argv) :
     """
@@ -376,10 +379,10 @@
     parser.add_option('--graph-dot', metavar='FILE',
             help="Output .dot graph data to file")
 
-    parser.add_option('--graph-vlans', action='store_true', dest='vlans', 
+    parser.add_option('--graph-vlans', action='store_true', dest='graph_vlans', 
             help="Graph all VLANs")
 
-    parser.add_option('--no-vlans', action='store_false', dest='vlans',
+    parser.add_option('--no-vlans', action='store_false', dest='graph_vlans',
             help="Do not color VLANs")
 
     # input
@@ -392,14 +395,21 @@
     # load raw snmp data
     snmp = load_snmp_data(options, options.snmp_data, hosts)
 
-    # TODO: process data into graph
+    # process data into graph
+    nodes, links = build_graph(snmp, hosts)
+    
+    # process graph into dot
+    if options.graph_vlans is False :
+        graph_vlans = None
+    else :
+        graph_vlans = GraphVlans()
 
-    # XXX: compat
-    vlans = dict(search_vlans(snmp))
-    lldp = list(search_lldp(snmp))
-    
-    # generate graph
-    apply_graph(options, lldp, vlans if options.vlans else None)
+    if options.graph_dot :
+        # process to dot
+        dot = build_dot(options, nodes, links, vlans=graph_vlans)
+        
+        # write out
+        apply_dot(options, open(options.graph_dot, 'w'), dot)
 
     return 0