#!/usr/bin/env python

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

import collections
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.
    """
    
    # second pass to discver links
    for host, agent in hosts_lldp(options, hosts) :
        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)

            log.info("%s: %s: %s", host, port, remote)

            yield host, agent.local, port, remote

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
            
        for port in set(vlan_untagged) | set(vlan_tagged) :
            yield host, port, vlan_untagged.get(port), tuple(vlan_tagged[port])

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)
        
        for (port, vlan), ethers in ports.iteritems() :
            yield host, vlan, port, ethers

def apply_hosts (options, hosts) :
    """
        Gather data on given hosts...

            (host, key, value)
    """
    
    if options.scan or options.scan_lldp :
        # discover node/port graph
        for host, local, port, remote in apply_hosts_lldp(options, hosts) :
            # XXX: duplicates galore
            yield host, ('lldp', 'local'), {
                    'chassis':  local['chassis'],
                    'sys_name': local['sys_name'],
            }
            
            # XXX: duplicates galore
            yield host, ('lldp', 'port', port['port_id'], 'local'), {
                    'port':     port['port'],
            }
        
            yield host, ('lldp', 'port', port['port_id'], 'remote', remote['chassis']), {
                    'sys_name':     remote['sys_name'],
                    'port':         remote['port'],
            }
    
    if options.scan or options.scan_vlan :
        # discover vlan ports
        for host, port, untag, tagged in apply_hosts_vlan(options, hosts) :
            if untag :
                yield host, ('vlan', untag, 'untagged'), set((port, ))

            for tag in tagged :
                yield host, ('vlan', tag, 'tagged'), set((port, ))
    
    if options.scan or options.scan_bridge :
        # discover edge nodes
        for host, vlan, port, ethers in apply_hosts_bridge(options, hosts) :
            if vlan :
                yield host, ('vlan', vlan, 'bridge', port), set(ethers)
            else :
                yield host, ('bridge', port), set(ethers)

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('--scan', action='store_true')
    parser.add_option('--scan-lldp', action='store_true')
    parser.add_option('--scan-vlan', action='store_true')
    parser.add_option('--scan-bridge', action='store_true')

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

    # gather SNMP data from hosts
    hosts = pvl.hosts.apply(options, args)
    
    data = collections.defaultdict(dict)

    for host, attr, values in apply_hosts(options, hosts) :
        if isinstance(values, set) :
            log.info("[%s] %s + %s", host, attr, values)

            data[host].setdefault(attr, set()).update(values)

        elif isinstance(values, dict) :
            log.info("[%s] %s = %s", host, attr, values)

            data[host][attr] = values

        else :
            raise Exception("No value for [%s] %s" % (host, attr))

    # output
    for host, attrs in sorted(data.items()) :
        print "{host}@{domain}".format(host=host, domain=host.domain)

        for attr, value in sorted(attrs.items()) :
            print "\t{attr}".format(attr=' '.join(str(a) for a in attr))

            if isinstance(value, set) :
                for v in sorted(value) :
                    print "\t\t{value}".format(value=v)
            elif isinstance(value, dict) :
                for k, v in sorted(value.items()) :
                    print "\t\t{key}\t{value}".format(key=k, value=v)
            else :
                raise Exception("[%s] %s: invalid value: %s" % (host, attr, value))

    return 0

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