#!/usr/bin/env python

"""
    Requirements:
        pydot
"""

import pvl.args
import pvl.hosts
from pvl.invoke import merge
from pvl.snmp import snmp, lldp, vlan, bridge

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

def hosts_snmp (options, hosts) :
    """
        Discover SNMP-supporting hosts.

        Yields Host, snmpdata
    """

    for host in hosts :
        host_snmp = host.extensions.get('snmp')

        if not host_snmp :
            log.debug("%s: skip non-snmp host", host)
            continue

        elif host.down :
            log.debug("%s: skip down host", host)
            continue

        else :
            log.debug("%s: %s", host, host_snmp)

        yield host, host_snmp

def hosts_lldp (options, hosts) :
    """
        Discover LLDP-supporting hosts.

        Yields Host, LLDPAgent
    """

    for host, host_snmp in hosts_snmp(options, hosts) :
        agent = lldp.LLDPAgent.apply(options, host.fqdn(), community=host_snmp.get('community'))

        try :
            local = agent.local
        except snmp.SNMPError as ex :
            log.warning("%s: %s", host, ex)
            continue
        
        if not local :
            log.info("%s: no lldp support", host)
            continue

        log.info("%s: %s", host, local)

        if local['sys_name'] != host.host :
            log.warning("%s: SNMP sys_name mismatch: %s", host, local['sys_name'])

        yield host, agent

def hosts_vlan (options, hosts) :
    """
        Discover VLAN-supporting hosts.

        Yields Host, VLANAgent
    """

    for host, host_snmp in hosts_snmp(options, hosts) :
        agent = vlan.VLANAgent.apply(options, host.fqdn(), community=host_snmp.get('community'))

        try :
            count = agent.count
        except snmp.SNMPError as ex :
            log.warning("%s: %s", host, ex)
            continue

        log.info("%s: %s", host, count)

        yield host, agent

def hosts_bridge (options, hosts) :
    """
        Discover Bridge-supporting hosts.

        Yields Host, BridgeAgent
    """

    for host, host_snmp in hosts_snmp(options, hosts) :
        agent = bridge.BridgeAgent.apply(options, host.fqdn(), community=host_snmp.get('community'))

        try :
            agent.ping()
        except snmp.SNMPError as ex :
            log.warning("%s: %s", host, ex)
            continue

        log.info("%s", host)

        yield host, agent


def apply_hosts_lldp (options, hosts) :
    """
        Query host LLDP info.
    """
    
    _hosts_lldp = list(hosts_lldp(options, hosts))
    hosts_by_chassis = { }
    
    # first pass to discover hosts
    for host, agent in _hosts_lldp :
        chassis = agent.local['chassis']
        log.info("%s: %s", host, chassis)

        hosts_by_chassis[chassis] = host
    
    # second pass to discver links
    for host, agent in _hosts_lldp :
        try :
            remotes = list(agent.remotes())
        except snmp.SNMPError as ex :
            log.warn("%s: broken lldp remotes: %s", host, ex)
            continue

        for port, remote in remotes :
            port = agent.port(port)

            remote_chassis = remote['chassis']
            remote_host = hosts_by_chassis.get(remote_chassis)
            
            log.info("%s: %s: %s (%s)", host, port, remote, remote_host)

            yield host, merge(agent.local, port), remote, remote_host

import collections

def apply_hosts_vlan (options, hosts) :
    """
        Query host VLAN ports.

        Yields host, { port: (untagged, [tagged]) }
    """

    _hosts_vlan = list(hosts_vlan(options, hosts))

    for host, agent in _hosts_vlan :
        # only one untagged vlan / port
        vlan_untagged = { }
            
        # multiple taggd vlans / port
        vlan_tagged = collections.defaultdict(set)

        for vlan, (tagged, untagged) in agent.vlan_ports() :
            log.info("%s: %s: %s + %s", host, vlan, tagged, untagged)
            
            for port in tagged :
                vlan_tagged[port].add(vlan)
            
            for port in untagged :
                if port in vlan_untagged :
                    log.warning("%s: duplicate untagged vlan %s for port %s on vlan %s", host, vlan, port, vlan_untagged[port])

                vlan_untagged[port] = vlan
            
        # pack into {port: (untagged, [tagged]) }
        yield host, dict(
            (
                    port, (vlan_untagged.get(port), tuple(vlan_tagged[port]))
            ) for port in set(vlan_untagged) | set(vlan_tagged)
        )

def apply_hosts_bridge (options, hosts) :
    """
        Query host bridge tables.

        Yields host, { port: (macs) }
    """

    for host, agent in hosts_bridge(options, hosts) :
        ports = collections.defaultdict(list)

        try :
            vlan_fdb_ports = list(agent.vlan_fdb_ports())
        except snmp.SNMPError as ex :
            log.warn("%s: broken dot1q fdb: %s", host, ex)
            continue

        if vlan_fdb_ports :
            log.info("%s: have dot1q ports", host)

            for ether, port, vlan in agent.vlan_fdb_ports() :
                if not port :
                    # XXX: unknown?
                    continue

                ports[(port, vlan)].append(ether)
        else :
            try :
                fdb_ports = list(agent.fdb_ports())
            except snmp.SNMPError as ex :
                log.warn("%s: broken dot1q fdb: %s", host, ex)
                continue
            
            # fallback to dot1d fdb
            log.info("%s: fallback to dot1d", host)

            for ether, port in agent.fdb_ports() :
                if not port :
                    # XXX: unknown?
                    continue

                ports[(port, None)].append(ether)

        yield host, ports

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
}

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) :
    """
        SNMP polling.
    """

    parser = optparse.OptionParser(main.__doc__)
    parser.add_option_group(pvl.args.parser(parser))
    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")

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

    # input
    options, args = parser.parse_args(argv[1:])
    pvl.args.apply(options)

    hosts = pvl.hosts.apply(options, args)
    
    # lookup host-port-vlan mappings
    if options.vlans is False :
        vlans = None
    else :
        vlans = dict(apply_hosts_vlan(options, hosts))

    # discover node/port graph
    items = apply_hosts_lldp(options, hosts)

    # discover edge nodes
    leafs = apply_hosts_bridge(options, hosts)

    # print
    if options.graph_dot :
        apply_graph(options, items, vlans)
    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='')
       
        for host, ports in vlans.iteritems() :
            for port, (untag, tagged) in ports.iteritems() :
                print "{host:30} {host.location:>30} {port:25} {untag}{tagged}".format(host=host, port=port,
                        untag       = '({untag}) '.format(untag=untag) if untag else '',
                        tagged      = ' '.join('<{tag}>'.format(tag=tag) for tag in tagged),
                )

        for host, ports in leafs :
            for (port, vlan), ethers in ports.iteritems() :
                print "{host:30} {host.location:>30} {port:25} <-- {vlan} # {ethers}".format(
                        host        = host,
                        port        = port,
                        vlan        = '<{vlan}>'.format(vlan=vlan) if vlan else '',
                        ethers      = ' '.join(ethers),
                )

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