terom@399: #!/usr/bin/env python
terom@399: 
terom@399: """
terom@399:     Requirements:
terom@399:         pydot
terom@399: """
terom@399: 
terom@399: import pvl.args
terom@399: import pvl.hosts
terom@402: from pvl.invoke import merge
terom@399: 
terom@399: import collections
terom@399: import logging; log = logging.getLogger('pvl.hosts-graph')
terom@399: import optparse
terom@399: 
terom@399: COLOR_VLANS = {
terom@399:     1:      'grey',         # pvl-lan
terom@399:     2:      'blue',         # pvl-lan2
terom@399:     3:      'red',          # pvl-san
terom@399:     4:      'green',        # pvl-veturi
terom@399:     7:      'orange',       # pvl-ranssi
terom@399:     8:      'yellow',       # pvl-mgmt
terom@399:     10:     'brown',        # pvl-public
terom@399:     100:    'navyblue',     # pvl-test
terom@399:     103:    'red4',         # pvl-test-san
terom@399:     104:    'red2',         # pvl-ganeti
terom@399:     192:    'purple',       # paivola-services
terom@399:     255:    'magenta',      # pvl-sonera
terom@399: }
terom@399: 
terom@399: class ParseError (Exception) :
terom@399:     def __init__ (self, file, line, msg) :
terom@399:         self.file = file
terom@399:         self.line = line
terom@399:         self.msg = msg
terom@399:   
terom@399:     def __str__ (self) :
terom@399:         return "{self.file}:{self.line}: {self.msg}".format(self=self)
terom@399: 
terom@400: def _parse_snmp_part (part) :
terom@400:     if part.isdigit() :
terom@400:         return int(part)
terom@400:     else :
terom@400:         return part
terom@400: 
terom@404: def _parse_snmp_attr (line) :
terom@400:     for part in line.split() :
terom@400:         yield _parse_snmp_part(part)
terom@400: 
terom@404: def _parse_snmp_value (line) :
terom@404:     if '\t' in line :
terom@404:         key, value = line.split('\t', 1)
terom@402: 
terom@404:         return { _parse_snmp_part(key): _parse_snmp_part(value) }
terom@402: 
terom@402:     else :
terom@404:         return set((_parse_snmp_part(line), ))
terom@404:     
terom@399: def _load_snmp_data (options, file) :
terom@399:     """
terom@399:         Load a data dict generated by pvl.hosts-snmp from a file.
terom@399: 
terom@399:         Yields (host, attr, value)
terom@399:     """
terom@399: 
terom@399:     host = None
terom@399:     attr = None
terom@399:     value = None
terom@399:     
terom@399:     for idx, line in enumerate(file, 1) :
terom@400:         indent = 0
terom@400: 
terom@400:         while line.startswith('\t') :
terom@400:             indent += 1
terom@400:             line = line[1:]
terom@400: 
terom@404:         line = line.lstrip('\t').rstrip('\n')
terom@400: 
terom@399:         if indent == 0 :
terom@399:             host = line
terom@399:             attr = None
terom@399:             value = None
terom@399: 
terom@399:         elif indent == 1 :
terom@404:             attr = tuple(_parse_snmp_attr(line))
terom@404:             value = None
terom@400: 
terom@404:             yield host, attr, None
terom@399: 
terom@399:         elif indent == 2 :
terom@399:             if not attr :
terom@400:                 raise ParseError(file, line, "[%s] %s: value outside of attr" % (host, attr))
terom@400:             
terom@404:             value = _parse_snmp_value(line)
terom@399: 
terom@404:             yield host, attr, value
terom@399: 
terom@399: def load_snmp_data (options, file, hosts) :
terom@399:     """
terom@399:         Load snmp data as dict, from given file path, or stdin.
terom@399:     """
terom@399: 
terom@399:     if file :
terom@399:         file = open(file)
terom@399:     else :
terom@399:         file = sys.stdin
terom@399: 
terom@399:     root = { }
terom@399: 
terom@399:     for host_name, attr, value in _load_snmp_data(options, file) :
terom@402:         host = hosts.get(host_name)
terom@404:         
terom@404:         if value :
terom@404:             log.info("[%s] %s: %s", host, ' '.join(str(a) for a in attr), value)
terom@404:         else :
terom@404:             log.info("[%s] %s", host, ' '.join(str(a) for a in attr),)
terom@399: 
terom@399:         item = root.setdefault(host, { })
terom@400:         
terom@399:         for a in attr[:-1] :
terom@399:             item = item.setdefault(a, {})
terom@399: 
terom@399:         a = attr[-1]
terom@400:         
terom@404:         if value is None :
terom@404:             pass
terom@404: 
terom@404:         elif isinstance(value, set) :
terom@400:             item.setdefault(a, set()).update(value)
terom@399: 
terom@404:         elif isinstance(value, dict) :
terom@404:             item.setdefault(a, dict()).update(value)
terom@404: 
terom@399:         else :
terom@402:             item[a] = value
terom@400:             
terom@399:     return root
terom@399: 
terom@402: def search_vlans (snmp) :
terom@402:     for host in snmp :
terom@402:         ports = set()
terom@402:         vlans_untagged = { }
terom@402:         vlans_tagged = collections.defaultdict(set)
terom@402: 
terom@402:         if 'vlan' not in snmp[host] :
terom@402:             continue
terom@402: 
terom@402:         for vlan in snmp[host]['vlan'] :
terom@402:             for port in snmp[host]['vlan'][vlan].get('tagged', ()) :
terom@402:                 ports.add(port)
terom@402:                 vlans_tagged[port].add(vlan)
terom@402:             
terom@402:             for port in snmp[host]['vlan'][vlan].get('untagged', ()) :
terom@404:                 ports.add(port)
terom@402:                 vlans_untagged[port] = vlan
terom@402:         
terom@402:         out = { }
terom@402:         for port in ports :
terom@402:             untag = vlans_untagged.get(port)
terom@404:             tagged = vlans_tagged.get(port, ())
terom@402: 
terom@402:             log.info("%s: %s: untag=%s tag=%s", host, port, untag, tagged)
terom@402: 
terom@402:             out[port] = (untag, tagged)
terom@402: 
terom@402:         yield host, out
terom@402: 
terom@402: def search_lldp (snmp) :
terom@402:     lldp_hosts = {} 
terom@402: 
terom@402:     # first pass to resolve chassis ID's
terom@402:     for host, host_attrs in snmp.iteritems() :
terom@402:         lldp = host_attrs.get('lldp')
terom@402: 
terom@402:         if not lldp :
terom@402:             continue
terom@402: 
terom@402:         lldp_local = lldp.get('local')
terom@402: 
terom@402:         if not lldp_local :
terom@402:             continue
terom@402: 
terom@402:         lldp_hosts[lldp_local['chassis']] = host
terom@402: 
terom@402:     # second pass to resolve remote
terom@402:     for host, host_attrs in snmp.iteritems() :
terom@402:         lldp = host_attrs.get('lldp')
terom@402: 
terom@402:         if not lldp :
terom@402:             continue
terom@402: 
terom@402:         local_attrs = lldp.get('local')
terom@402: 
terom@402:         for port, port_attrs in snmp[host]['lldp']['port'].iteritems() :
terom@402:             port_local_attrs = port_attrs['local']
terom@402: 
terom@402:             for chassis, remote_attrs in port_attrs['remote'].iteritems() :
terom@402:                 remote_host = lldp_hosts.get(chassis)
terom@402: 
terom@402:                 local = merge(local_attrs, port_local_attrs, port_id=port)
terom@402:                 remote = merge(remote_attrs, chassis=chassis)
terom@402: 
terom@402:                 log.info("%s @ %s %s : %s %s @ %s", host, host.location, local, remote, remote_host.location if remote_host else '', remote_host)
terom@402: 
terom@402:                 yield host, local, remote, remote_host
terom@402: 
terom@399: def apply_graph (options, items, vlans={}) :
terom@399:     import pydot
terom@399: 
terom@399:     dot = pydot.Dot(graph_name='lldp_hosts', graph_type='digraph',
terom@399:             # XXX: breaks multi-edges?
terom@399:             #splines     = 'true',
terom@399: 
terom@399:             sep             = '+25,25',
terom@399:             overlap         = 'scalexy',
terom@399: 
terom@399:             # only applies to loops
terom@399:             nodesep     = 0.5,
terom@399:     )
terom@399:     dot.set_edge_defaults(
terom@399:             labeldistance   = 3.0,
terom@399:             penwidth        = 2.0,
terom@399:     )
terom@399: 
terom@399:     nodes = { }
terom@399:     edges = { }
terom@399:     vlan_colors = { } # { vlan: color }
terom@399: 
terom@399:     for host, local, remote, remote_host in items :
terom@399:         # src
terom@399:         src_name = str(host)
terom@399:         src_label = '"{host.location}"'.format(host=host)
terom@399: 
terom@399:         if src_name in nodes :
terom@399:             src = nodes[src_name]
terom@399:         else :
terom@399:             src = nodes[src_name] = pydot.Node(src_name,
terom@399:                     label       = src_label,
terom@399:                     fontsize    = 18,
terom@399:             )
terom@399:             dot.add_node(src)
terom@399:         
terom@399:         # dst
terom@399:         if remote_host :
terom@399:             dst_name = str(remote_host)
terom@399:             dst_label = '"{host.location}"'.format(host=remote_host)
terom@399:         else :
terom@399:             dst_name = remote['chassis'].replace(':', '-')
terom@399: 
terom@399:             # XXX: pydot is not smart enough to quote this
terom@399:             if remote['sys_name'] :
terom@399:                 dst_label = '"{remote[chassis]} ({remote[sys_name]})"'.format(remote=remote)
terom@399:             else :
terom@399:                 dst_label = '"{remote[chassis]}"'.format(remote=remote)
terom@399:         
terom@399:         if dst_name in nodes :
terom@399:             dst = nodes[dst_name]
terom@399:         else :
terom@399:             dst = nodes[dst_name] = pydot.Node(dst_name,
terom@399:                     label       = dst_label,
terom@399:                     fontsize    = 18,
terom@399:             )
terom@399:             dot.add_node(dst)
terom@399: 
terom@399:         # edges
terom@399:         headlabel = '"{remote[port]}"'.format(remote=remote)
terom@399:         taillabel = '"{local[port]}"'.format(local=local)
terom@399:         fillcolor = 'black'
terom@399:         color = 'black'
terom@399:         untag = tagged = None
terom@399: 
terom@399:         # vlans?
terom@399:         if vlans and host in vlans and local['port_id'] in vlans[host] :
terom@399:             untag, tagged = vlans[host][local['port_id']]
terom@399:             
terom@399:             log.debug("%s#%s: %s+%s", host, local['port_id'], untag, tagged)
terom@399:             
terom@399:             colors = []
terom@399:             for tag in sorted(tagged) :
terom@399:                 if tag in COLOR_VLANS :
terom@399:                     colors.append(COLOR_VLANS[tag])
terom@399:                 elif tag in vlan_colors :
terom@399:                     colors.append(vlan_colors[tag])
terom@399:                 else :
terom@399:                     color = '/paired12/{count}'.format(count=1+len(vlan_colors))
terom@399:                     
terom@399:                     log.info("%s#%s: chosing new vlan %s color %s", host, local['port_id'], tag, color)
terom@399: 
terom@399:                     vlan_colors[tag] = color
terom@399:                     colors.append(color)
terom@399: 
terom@399:             if not untag :
terom@399:                 pass
terom@399:             elif untag in COLOR_VLANS :
terom@399:                 fillcolor = COLOR_VLANS[untag]
terom@399:             elif untag in vlan_colors :
terom@399:                 fillcolor = vlan_colors[untag]
terom@399:             else :
terom@399:                 color = '/paired12/{count}'.format(count=1+len(vlan_colors))
terom@399: 
terom@399:                 log.warn("%s#%s: chosing new vlan %s color %s", host, local['port_id'], untag, color)
terom@399:                     
terom@399:                 fillcolor = vlan_colors[tag] = color
terom@399:             
terom@399:             # first color overrides fillcolor for heads
terom@399:             if colors and untag :
terom@399:                 color = ':'.join([fillcolor] + colors)
terom@399:             elif colors :
terom@399:                 color = ':'.join(colors)
terom@399:             elif fillcolor :
terom@399:                 color = fillcolor
terom@399: 
terom@399:         elif vlans :
terom@399:             # XXX: this happens when LLDP gives us the LACP ports but the VLANS are on the TRK port
terom@399:             log.warn("%s#%s: unknown port for vlans: %s", host, local['port_id'], vlans.get(host))
terom@399: 
terom@399:         # edge
terom@399:         if (src_name, local['port'], dst_name, remote['port']) in edges :
terom@399:             log.warning("%s:%s <- %s:%s: duplicate", src_name, local['port'], dst_name, remote['port'])
terom@399:         
terom@399:         elif (dst_name, remote['port'], src_name, local['port']) in edges :
terom@399:             log.info("%s <-> %s", src_name, dst_name)
terom@399:             edge = edges[(dst_name, remote['port'], src_name, local['port'])]
terom@399: 
terom@399:             if edge.get('headlabel') != taillabel :
terom@399:                 log.warn("%s -> %s: local port mismatch: %s vs %s", src_name, dst_name, local['port'], edge.get('headlabel'))
terom@399:             
terom@399:             if edge.get('taillabel') != headlabel :
terom@399:                 log.warn("%s -> %s: remote port mismatch: %s vs %s", src_name, dst_name, remote['port'], edge.get('taillabel'))
terom@399: 
terom@399:             if edge.get('fillcolor') != fillcolor :
terom@399:                 log.warn("%s#%s -> %s#%s: remote untag mismatch: %s vs %s", src_name, local['port'], dst_name, remote['port'], fillcolor, edge.get('fillcolor'))
terom@399: 
terom@399:             if edge.get('color') != '"' + color + '"' :
terom@399:                 log.warn("%s#%s -> %s#%s: remote tagged mismatch: %s vs %s", src_name, local['port'], dst_name, remote['port'], color, edge.get('color'))
terom@399: 
terom@399:             # mark as bidirectional
terom@399:             edges[(src_name, local['port'], dst_name, remote['port'])] = edge
terom@399:             
terom@399:             edge.set('dir', 'both' if untag else 'none')
terom@399: 
terom@399:             # set second color for tail
terom@399:             if untag :
terom@399:                 edge.set('color', '"{headcolor}:{tailcolor}{tagcolors}"'.format(
terom@399:                     headcolor   = edge.get('fillcolor'),
terom@399:                     tailcolor   = fillcolor,
terom@399:                     tagcolors   = ':' + ':'.join(colors) if colors else '',
terom@399:                 ))
terom@399: 
terom@399:         else :
terom@399:             edge = edges[(src_name, local['port'], dst_name, remote['port'])] = pydot.Edge(src, dst,
terom@399:                     dir         = 'forward' if untag else 'none',
terom@399:                     headlabel   = headlabel,
terom@399:                     taillabel   = taillabel,
terom@399:                     color       = '"{color}"'.format(color=color),
terom@399:                     fillcolor   = fillcolor,
terom@399:             )
terom@399: 
terom@399:             dot.add_edge(edge)
terom@399: 
terom@399:     if options.graph_dot :
terom@399:         dot.write(options.graph_dot)
terom@399: 
terom@399: def main (argv) :
terom@399:     """
terom@399:         Graph network
terom@399:     """
terom@399: 
terom@399:     parser = optparse.OptionParser(main.__doc__)
terom@399:     parser.add_option_group(pvl.args.parser(parser))
terom@399:     parser.add_option_group(pvl.hosts.optparser(parser))
terom@399: 
terom@399:     parser.add_option('--snmp-data', metavar='FILE', default=None,
terom@399:             help="Load snmp data from FILE")
terom@399: 
terom@399:     parser.add_option('--graph-dot', metavar='FILE',
terom@399:             help="Output .dot graph data to file")
terom@399: 
terom@399:     parser.add_option('--graph-vlans', action='store_true', dest='vlans', 
terom@399:             help="Graph all VLANs")
terom@399: 
terom@399:     parser.add_option('--no-vlans', action='store_false', dest='vlans',
terom@399:             help="Do not color VLANs")
terom@399: 
terom@399:     # input
terom@399:     options, args = parser.parse_args(argv[1:])
terom@399:     pvl.args.apply(options)
terom@399:     
terom@399:     # load hosts for correlation
terom@399:     hosts = dict((str(host), host) for host in pvl.hosts.apply(options, args))
terom@399: 
terom@399:     # load raw snmp data
terom@402:     snmp = load_snmp_data(options, options.snmp_data, hosts)
terom@399: 
terom@399:     # TODO: process data into graph
terom@402: 
terom@402:     # XXX: compat
terom@402:     vlans = dict(search_vlans(snmp))
terom@402:     lldp = list(search_lldp(snmp))
terom@399:     
terom@402:     # generate graph
terom@402:     apply_graph(options, lldp, vlans if options.vlans else None)
terom@399: 
terom@399:     return 0
terom@399: 
terom@399: if __name__ == '__main__':
terom@399:     pvl.args.main(main)