# HG changeset patch
# User Tero Marttila <terom@paivola.fi>
# Date 1395079291 -7200
# Node ID 75158fd28784f249ca7d88597f2bfcc942861a02
# Parent  9e1abcf47d27e53d99a39695d72f68e05a15c29f
pvl.hosts-lldp: --graph-dot support

diff -r 9e1abcf47d27 -r 75158fd28784 bin/pvl.hosts-lldp
--- a/bin/pvl.hosts-lldp	Mon Mar 17 17:46:48 2014 +0200
+++ b/bin/pvl.hosts-lldp	Mon Mar 17 20:01:31 2014 +0200
@@ -2,15 +2,7 @@
 
 """
     Requirements:
-        pysnmp
-        pysnmp-mibs
-        memoized-property
-
-    Setup:
-        ./opt/bin/build-pysnmp-mib -o usr/mibs/LLDP-MIB.py etc/mibs/LLDP-MIB
-
-    Run:
-        PYSNMP_MIB_DIRS=usr/mibs/ ./opt/bin/python ...
+        pydot
 """
 
 import pvl.args
@@ -85,6 +77,90 @@
 
             yield host, merge(agent.local, port), remote, remote_host
 
+def apply_graph (options, items) :
+    import pydot
+
+    dot = pydot.Dot(graph_name='lldp_hosts', graph_type='digraph',
+            splines     = 'true',
+
+            sep             = '+25,25',
+            overlap         = 'scalexy',
+
+            # only applies to loops
+            nodesep     = 0.5,
+    )
+    dot.set_edge_defaults(
+            labeldistance   = 3.0,
+    )
+
+    nodes = { }
+    edges = { }
+
+    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)
+            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)
+            dot.add_node(dst)
+
+        # edge
+        headlabel = '"{remote[port]}"'.format(remote=remote)
+        taillabel = '"{local[port]}"'.format(local=local)
+
+        if (src_name, dst_name) in edges :
+            log.warning("%s <- %s: duplicate", src_name, dst_name)
+        
+        elif (dst_name, src_name) in edges :
+            log.info("%s <-> %s", src_name, dst_name)
+            edge = edges[(dst_name, src_name)]
+
+            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'))
+
+            # mark as bidirectional
+            edges[(src_name, dst_name)] = edge
+            
+            edge.set('dir', 'both')
+
+        else :
+            edge = edges[(src_name, dst_name)] = pydot.Edge(src, dst,
+                    dir         = 'back',
+                    headlabel   = headlabel,
+                    taillabel   = taillabel,
+            )
+
+            dot.add_edge(edge)
+
+    if options.graph_dot :
+        dot.write(options.graph_dot)
+
 def main (argv) :
     """
         SNMP polling.
@@ -95,6 +171,8 @@
     parser.add_option_group(pvl.hosts.optparser(parser))
     parser.add_option_group(pvl.snmp.snmp.options(parser))
 
+    parser.add_option('--graph-dot',
+            help="Output .dot graph data")
 
     options, args = parser.parse_args(argv[1:])
     pvl.args.apply(options)
@@ -103,11 +181,17 @@
     hosts = pvl.hosts.apply(options, args)
     
     # apply
-    for host, local, remote, remote_host in apply_hosts_lldp(options, hosts) :
-        if remote_host :
-            print "{host:30} {host.location:>30} {local[port]:>25} <-> {remote[port]:<25} {remote_host.location:>30} # {remote[chassis]} ({remote_host})".format(host=host, local=local, remote=remote, remote_host=remote_host)
-        else :
-            print "{host:30} {host.location:>30} {local[port]:>25} <-- {remote[port]:<25} {empty:30} # {remote[chassis]} ({remote[sys_name]})".format(host=host, local=local, remote=remote, empty='')
+    items = apply_hosts_lldp(options, hosts)
+
+    # print
+    if options.graph_dot :
+        apply_graph(options, items)
+    else :
+        for host, local, remote, remote_host in items :
+            if remote_host :
+                print "{host:30} {host.location:>30} {local[port]:>25} <-> {remote[port]:<25} {remote_host.location:>30} # {remote[chassis]} ({remote_host})".format(host=host, local=local, remote=remote, remote_host=remote_host)
+            else :
+                print "{host:30} {host.location:>30} {local[port]:>25} <-- {remote[port]:<25} {empty:30} # {remote[chassis]} ({remote[sys_name]})".format(host=host, local=local, remote=remote, empty='')
 
 if __name__ == '__main__':
     pvl.args.main(main)