#!/usr/bin/env python
"""
   Retrieve a subtree of lexicographically greater Object IDs starting from
   user specifed SNMP Object ID (using GETWALK or GETBULK PDUs) from arbitrary
   SNMP agent.

   Copyright 1999-2004 by Ilya Etingof <ilya@glas.net>. See LICENSE for
   details.
"""
from sys import exit, argv, stderr
from string import find
from pysnmp.version import verifyVersionRequirement

verifyVersionRequirement(3, 4)

from pysnmp.mapping.udp.role import Manager
from pysnmp.proto.api import alpha
from pysnmpap.cli import options, ucd
from pysnmpap.error import SnmpApplicationError

class ConfSnmpVersion(options.ConfOptionValue):
    def help(self): return 'SNMP protocol version to use [%s]' % self
confSnmpVersion = ConfSnmpVersion(alpha.protoVersionId1)    

class ConfReportType(options.ConfOptionFlag):
    def help(self): return 'report variables types [%s]' % self
confReportType = ConfReportType(0)

class ConfReportQty(options.ConfOptionFlag):
    def help(self): return 'print the number of variables found [%s]' % self
confReportQty = ConfReportQty(0)

class ConfInclOid(options.ConfOptionFlag):
    def help(self): return 'include given OID in search range [%s]' % self
confInclOid = ConfInclOid(0)

class ConfIncrOid(options.ConfOptionFlag):
    def help(self): return 'check returned OIDs are increasing [%s]' % self
confIncrOid = ConfIncrOid(1)

def formatUsage():
    tspDsp = Manager()
    ver = alpha.protoVersions[confSnmpVersion()]
    getNextReq = ver.Message()
    getNextReq.apiAlphaSetPdu(ver.pduTypes[nextPduType]())
    appUsage = '[ -h ] [ -v <version> ] [ -C <p|i|c|t> ]'
    appOpts =           ' -h             print this usage note.\n'
    appOpts = appOpts + ' -v <version>   ' + confSnmpVersion.help() + '\n'
    appOpts = appOpts + ' -C <p|i|c|t>   application specific options\n'
    appOpts = appOpts + '     p:         ' + confReportQty.help() + '\n'
    appOpts = appOpts + '     i:         ' + confInclOid.help() + '\n'
    appOpts = appOpts + '     c:         ' + confIncrOid.help() + '\n'
    appOpts = appOpts + '     t:         ' + confReportType.help() + '\n'
    usage = 'Usage: ' + argv[0] + ' ' + appUsage + ' ' + \
            tspDsp.cliUcdGetUsage() + ' ' + getNextReq.cliUcdGetUsage()
    options = 'Application options:\n' + appOpts + \
              'Transport options:\n' + tspDsp.cliUcdGetOptsUsage(' ') + \
              'Protocol options:\n' + getNextReq.cliUcdGetOptsUsage()
    return usage + '\n' + options

# Figure out operation mode (not GNUish)
if find(argv[0], 'pysnmpwalk') != -1:
    nextPduType = alpha.getNextRequestPduType
elif find(argv[0], 'pysnmpbulkwalk') != -1:
    nextPduType = alpha.getBulkRequestPduType
    confSnmpVersion(alpha.protoVersionId2c)    
else:
    stderr.write('Cant determine operation mode (walk/bulk) from tool name: %s' % argv[0])
    exit(1)

# Parse app options
idx = 1; newArgv = []
while idx < len(argv):
    if argv[idx] == '-h':
        stderr.write(formatUsage() + '\n')
        exit(1)
    if argv[idx] == '-v':
        if argv[idx+1][:1] != 'v': argv[idx+1] = 'v'+argv[idx+1]
        if argv[idx+1] == 'v1':
            confSnmpVersion(alpha.protoVersionId1)
        elif argv[idx+1] == 'v2c':
            confSnmpVersion(alpha.protoVersionId2c)
        else:
            stderr.write('Unsupported SNMP protocol version: %s\n%s\n' % \
                         (argv[idx+1], formatUsage()))
            exit(-1)
        if nextPduType == alpha.getBulkRequestPduType and \
               confSnmpVersion != alpha.protoVersionId2c:
            stderr.write('Operation not supported by this SNMP version\n')
            exit(1)
        idx = idx + 2
        continue
    if argv[idx] == '-C':
        try:
            rest = ''
            for c in list(argv[idx+1]):
                if c == 'p': confReportQty(1)
                elif c == 'i': confInclOid(1)
                elif c == 'c': confIncrOid(0)
                elif c == 't': confReportType(1)
                else: rest = rest + c
        except IndexError, why:
            stderr.write('Missing value to %s\n%s\n' % \
                         (argv[idx], formatUsage()))
            exit(-1)
        if rest: newArgv.append(argv[idx] + ' ' + rest)
        idx = idx + 2
        continue
    newArgv.append(argv[idx])
    idx = idx + 1

# Proto & transport objects
tspDsp = Manager()
ver = alpha.protoVersions[confSnmpVersion()]
getNextReq = req = ver.Message(); rsp = ver.Message()
getNextReq.apiAlphaSetPdu(ver.pduTypes[nextPduType]())

# Parse transport and protocol options
try:
    newArgv = tspDsp.cliUcdSetArgs(newArgv)
    tempArgv = getNextReq.cliUcdSetArgs(newArgv)
    if confInclOid:
        getReq = req = ver.Message()
        getReq.apiAlphaSetPdu(ver.GetRequestPdu())
        getReq.cliUcdSetArgs(newArgv)
    newArgv = tempArgv
except SnmpApplicationError, why:
    stderr.write('Error: %s\n%s\n' % (why, formatUsage()))
    exit(-1)

if len(newArgv):
    stderr.write('Extra parameters given: %s\n%s\n' \
                 % (newArgv, formatUsage()))
    exit(-1)

headVars = map(lambda x: ver.ObjectIdentifier(x.apiAlphaGetOidVal()[0].get()),
               getNextReq.apiAlphaGetPdu().apiAlphaGetVarBindList())

prtOidCount = 0

def cbFun(wholeMsg, srcAddr, (req, rsp, headVars)):
    global prtOidCount
    rsp.berDecode(wholeMsg)
    if not req.apiAlphaMatch(rsp): return
    
    reqPdu = req.apiAlphaGetPdu(); rspPdu = rsp.apiAlphaGetPdu()
    
    # Handle errors
    errorStatus = rspPdu.apiAlphaGetErrorStatus()
    if errorStatus != 0 and errorStatus != 2:
        errorIndex = rspPdu.apiAlphaGetErrorIndex()-1
        varBindList = rspPdu.apiAlphaGetVarBindList()
        raise SnmpApplicationError('%s at %s' % (errorStatus, \
                                                 varBindList[errorIndex.get()]))

    # Browse and report SNMP table
    tableIndices = apply(reqPdu.apiAlphaGetTableIndices, [rspPdu] + headVars)
    for rowIndices in tableIndices:
        for idx in filter(lambda x: x!=-1, rowIndices):
            oid, val = rspPdu.apiAlphaGetVarBindList()[idx].apiAlphaGetOidVal()
            if confIncrOid and not confInclOid and \
                   nextPduType != alpha.getBulkRequestPduType:
                pOid, pVal = reqPdu.apiAlphaGetVarBindList()[idx].apiAlphaGetOidVal()
                if pOid >= oid:
                    raise SnmpApplicationError('OID not increasing: %s >= %s'\
                                               % (oid, pOid))
            if confReportType:
                print oid, '=', val
            else:
                print oid.get(), '=', repr(val.get())

            if confReportQty: prtOidCount = prtOidCount + 1

    # GET req processing stops here
    if confInclOid: return 1
    
    # Remove completed SNMP table columns
    if not confInclOid:
        map(lambda idx, headVars=headVars: headVars.__delitem__(idx), \
            filter(lambda x: x==-1, tableIndices[-1]))

    if len(headVars) == 0: raise SnmpApplicationError()

    # Prepare to request next SNMP table row
    apply(reqPdu.apiAlphaSetVarBindList, \
          map(lambda (o, v): (o.get(), None), \
              map(lambda cellIdx, varBindList=rspPdu.apiAlphaGetVarBindList():\
                  varBindList[cellIdx].apiAlphaGetOidVal(), \
                  filter(lambda x: x!=-1, tableIndices[-1]))))

    req.apiAlphaGetPdu().apiAlphaGetRequestId().inc(1)
    return 1

try:
    while 1:
        tspDsp.sendAndReceive(req.berEncode(), (None, 0), \
                              (cbFun, (req, rsp, headVars)))
        if confInclOid:
            req = getNextReq
            confInclOid(0)
        
except SnmpApplicationError, why:
    if why: stderr.write('Application error: %s\n' % why)

except Exception, why:
    stderr.write('pysnmp error: %s\n' % why)

if confReportQty: print 'Variables found: %d' % prtOidCount
