#!/usr/bin/env python

"""
    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 ...
"""

import pvl.args
import pvl.hosts

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

from pysnmp.entity.rfc3413.oneliner import cmdgen as pysnmp
import collections

class SNMPError (Exception) :
    pass

class SNMPEngineError (SNMPError) :
    """
        Internal SNMP Engine error (?)
    """

class SNMPAgent (object) :
    """
        GET SNMP shit from a remote host.
    """

    SNMP_PORT = 161
    snmp_cmdgen = pysnmp.CommandGenerator()

    @classmethod
    def apply (cls, options, host) :
        community = options.snmp_community
        port = cls.SNMP_PORT

        if '@' in host :
            community, host = host.split('@', 1)
        
        if ':' in host :
            host, port = host.rsplit(':', 1)

        return cls(
                pysnmp.CommunityData(community),
                pysnmp.UdpTransportTarget((host, port))
        )

    def __init__ (self, security, transport) :
        self.security = security
        self.transport = transport

    def get (self, *request) :
        """
            request = (
                pysnmp.MibVariable('IF-MIB', 'ifInOctets', 1),
            )
        """

        opts = dict(
                lookupNames     = True,
                lookupValues    = True,
        )

        try :
            error, error_status, error_index, response = self.snmp_cmdgen.getCmd(self.security, self.transport,  *request, **opts)
        except pysnmp.error.PySnmpError as ex :
            raise SNMPEngineError(ex)

        if error :
            raise SNMPEngineError(error)
        
        if error_status :
            raise SNMPError(errorStatus.prettyPrint())
        
        return response
        #for name, value in response :
        #    yield name.getMibSymbol(), value.prettyPrint()

    def walk (self, *request) :
        """
            request = (
                    pysnmp.MibVariable('IF-MIB', 'ifInOctets'),
            )
        """

        opts = dict(
                lookupNames     = True,
                lookupValues    = True,
        )

        try :
            error, error_status, error_index, responses = self.snmp_cmdgen.nextCmd(self.security, self.transport, *request, **opts)
        except pysnmp.error.PySnmpError as ex :
            raise SNMPEngineError(ex)

        if error :
            raise SNMPEngineError(error)
        
        if error_status :
            raise SNMPError(errorStatus.prettyPrint())
        
        return responses
        #for response in responses:
        #    for name, value in response :
        #        yield name.getMibSymbol(), value.prettyPrint()

    def table (self, table) :
        """
            Load a simple table with a running index.

            Given [(attr, oid, type)] yields { attr: type(value) }
        """

        data = collections.defaultdict(dict)

        for row in self.walk(*[oid for attr, oid, type in table]) :
            log.debug("%s", row)

            for (attr, type), (field, value) in zip([(attr, type) for attr, oid, type in table], row) :
                mib, sym, idx = field.getMibSymbol()

                # convert
                if type :
                    value = type(value)

                log.debug("%s::%s.%s: %s: %s", mib, sym, idx, attr, value)
                
                data[idx][attr] = value
            
        return data.items()

from memoized_property import memoized_property

def macaddr (value) :
    """
        Excepts a MAC address from an SNMP OctetString.
    """

    return ':'.join('{octet:02x}'.format(octet=c) for c in value.asNumbers())

class LLDPAgent (SNMPAgent) :
    """
        Query LLDP info from a remote agent.
    """

    LOCAL = (
            ( 'chassis_id',     pysnmp.MibVariable('LLDP-MIB', 'lldpLocChassisId'),     macaddr ),
            ( 'sys_name',       pysnmp.MibVariable('LLDP-MIB', 'lldpLocSysName'),       str     ),
    )
    
    @memoized_property
    def local (self) :
        """
            Describe the local system.
        """

        for idx, data in self.table(self.LOCAL) :
            return data

    PORTS = (
            ( 'port_id',        pysnmp.MibVariable('LLDP-MIB', 'lldpLocPortId'),      str     ),
    )

    @memoized_property
    def ports (self) :
        """
            Describe the local ports.
        """

        ports = { }

        for idx, data in self.table(self.PORTS) :
            port, = idx
            
            ports[int(port)] = data

        return ports

    def port (self, port) :
        return self.ports[port]

    REMOTE_DATA = (
            ( 'chassis_id',     pysnmp.MibVariable('LLDP-MIB', 'lldpRemChassisId'),     macaddr ),
            ( 'sys_name',       pysnmp.MibVariable('LLDP-MIB', 'lldpRemSysName'),       str     ),
            ( 'remote_port',    pysnmp.MibVariable('LLDP-MIB', 'lldpRemPortId'),        str     ),
    )

    def remotes (self) :
        """
            Describe remote systems, indexed by local port.
        """

        for idx, data in self.table(self.REMOTE_DATA) :
            time, port, idx = idx

            yield int(port), data

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('--snmp-community',
            help="SNMPv2 read community")

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

    # input
    #hosts = pvl.hosts.apply(options, args)

    for arg in args :
        lldp = LLDPAgent.apply(options, arg)
        
        print arg, "{chassis_id}({sys_name})".format(**lldp.local)
        
        for port, remote in lldp.remotes() :
            print '\t', lldp.port(port), remote

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