#
# This file is part of pysnmp software.
#
# Copyright (c) 2005-2018, Ilya Etingof <etingof@gmail.com>
# License: http://snmplabs.com/pysnmp/license.html
#
from pysnmp.smi.rfc1902 import *
from pysnmp.hlapi.v3arch.auth import *
from pysnmp.hlapi.v3arch.context import *
from pysnmp.hlapi.v3arch.lcd import *
from pysnmp.hlapi.varbinds import *
from pysnmp.hlapi.v3arch.twisted.transport import *
from pysnmp.entity.rfc3413 import cmdgen
from pysnmp.proto import errind
from pysnmp.proto.api import v2c
from twisted.internet.defer import Deferred
from twisted.python.failure import Failure

__all__ = ['getCmd', 'nextCmd', 'setCmd', 'bulkCmd', 'isEndOfMib']

vbProcessor = CommandGeneratorVarBinds()
lcd = CommandGeneratorLcdConfigurator()

isEndOfMib = lambda varBinds: not v2c.apiPDU.getNextVarBinds(varBinds)


def getCmd(snmpEngine, authData, transportTarget, contextData,
           *varBinds, **options):
    """Performs SNMP GET query.

    Based on passed parameters, prepares SNMP GET packet
    (:RFC:`1905#section-4.2.1`) and schedules its transmission by
    :mod:`twisted` I/O framework at a later point of time.

    Parameters
    ----------
    snmpEngine : :class:`~pysnmp.hlapi.SnmpEngine`
        Class instance representing SNMP engine.

    authData : :class:`~pysnmp.hlapi.CommunityData` or :class:`~pysnmp.hlapi.UsmUserData`
        Class instance representing SNMP credentials.

    transportTarget : :class:`~pysnmp.hlapi.twisted.UdpTransportTarget` or :class:`~pysnmp.hlapi.twisted.Udp6TransportTarget`
        Class instance representing transport type along with SNMP peer address.

    contextData : :class:`~pysnmp.hlapi.ContextData`
        Class instance representing SNMP ContextEngineId and ContextName values.

    \*varBinds : :class:`~pysnmp.smi.rfc1902.ObjectType`
        One or more class instances representing MIB variables to place
        into SNMP request.

    Other Parameters
    ----------------
    \*\*options :
        Request options:

            * `lookupMib` - load MIB and resolve response MIB variables at
              the cost of slightly reduced performance. Default is `True`.

    Returns
    -------
    deferred : :class:`~twisted.internet.defer.Deferred`
        Twisted Deferred object representing work-in-progress. User
        is expected to attach his own `success` and `error` callback
        functions to the Deferred object though
        :meth:`~twisted.internet.defer.Deferred.addCallbacks` method.

    Raises
    ------
    PySnmpError
        Or its derivative indicating that an error occurred while
        performing SNMP operation.

    Notes
    -----
    User `success` callback is called with the following tuple as
    its first argument:

    * errorStatus (str) : True value indicates SNMP PDU error.
    * errorIndex (int) : Non-zero value refers to `varBinds[errorIndex-1]`
    * varBinds (tuple) : A sequence of
      :class:`~pysnmp.smi.rfc1902.ObjectType` class instances representing
      MIB variables returned in SNMP response.

    User `error` callback is called with `errorIndication` object wrapped
    in :class:`~twisted.python.failure.Failure` object.

    Examples
    --------
    >>> from twisted.internet.task import react
    >>> from pysnmp.hlapi.twisted import *
    >>>
    >>> def success(args):
    ...     (errorStatus, errorIndex, varBinds) = args
    ...     print(errorStatus, errorIndex, varBind)
    ...
    >>> def failure(errorIndication):
    ...     print(errorIndication)
    ...
    >>> def run(reactor):
    ...     d = getCmd(SnmpEngine(),
    ...                CommunityData('public'),
    ...                UdpTransportTarget(('demo.snmplabs.com', 161)),
    ...                ContextData(),
    ...                ObjectType(ObjectIdentity('SNMPv2-MIB', 'sysDescr', 0)))
    ...     d.addCallback(success).addErrback(failure)
    ...     return d
    ...
    >>> react(run)
    (0, 0, [ObjectType(ObjectIdentity(ObjectName('1.3.6.1.2.1.1.1.0')), DisplayString('SunOS zeus.snmplabs.com 4.1.3_U1 1 sun4m'))])
    >>>

    """
    def __cbFun(snmpEngine, sendRequestHandle,
                errorIndication, errorStatus, errorIndex,
                varBinds, cbCtx):
        lookupMib, deferred = cbCtx
        if errorIndication:
            deferred.errback(Failure(errorIndication))
        else:
            try:
                varBindsUnmade = vbProcessor.unmakeVarBinds(snmpEngine.cache, varBinds, lookupMib)

            except Exception as e:
                deferred.errback(Failure(e))

            else:
                deferred.callback((errorStatus, errorIndex, varBindsUnmade))

    addrName, paramsName = lcd.configure(
        snmpEngine, authData, transportTarget, contextData.contextName)

    deferred = Deferred()

    cmdgen.GetCommandGenerator().sendVarBinds(
        snmpEngine, addrName, contextData.contextEngineId,
        contextData.contextName,
        vbProcessor.makeVarBinds(snmpEngine.cache, varBinds), __cbFun,
        (options.get('lookupMib', True), deferred)
    )
    return deferred

def setCmd(snmpEngine, authData, transportTarget, contextData,
           *varBinds, **options):
    """Performs SNMP SET query.

    Based on passed parameters, prepares SNMP SET packet
    (:RFC:`1905#section-4.2.5`) and schedules its transmission by
    :mod:`twisted` I/O framework at a later point of time.

    Parameters
    ----------
    snmpEngine : :class:`~pysnmp.hlapi.SnmpEngine`
        Class instance representing SNMP engine.

    authData : :class:`~pysnmp.hlapi.CommunityData` or :class:`~pysnmp.hlapi.UsmUserData`
        Class instance representing SNMP credentials.

    transportTarget : :class:`~pysnmp.hlapi.twisted.UdpTransportTarget` or :class:`~pysnmp.hlapi.twisted.Udp6TransportTarget`
        Class instance representing transport type along with SNMP peer address.

    contextData : :class:`~pysnmp.hlapi.ContextData`
        Class instance representing SNMP ContextEngineId and ContextName values.

    \*varBinds : :class:`~pysnmp.smi.rfc1902.ObjectType`
        One or more class instances representing MIB variables to place
        into SNMP request.

    Other Parameters
    ----------------
    \*\*options :
        Request options:

            * `lookupMib` - load MIB and resolve response MIB variables at
              the cost of slightly reduced performance. Default is `True`.

    Returns
    -------
    deferred : :class:`~twisted.internet.defer.Deferred`
        Twisted Deferred object representing work-in-progress. User
        is expected to attach his own `success` and `error` callback
        functions to the Deferred object though
        :meth:`~twisted.internet.defer.Deferred.addCallbacks` method.

    Raises
    ------
    PySnmpError
        Or its derivative indicating that an error occurred while
        performing SNMP operation.

    Notes
    -----
    User `success` callback is called with the following tuple as
    its first argument:

    * errorStatus (str) : True value indicates SNMP PDU error.
    * errorIndex (int) : Non-zero value refers to `varBinds[errorIndex-1]`
    * varBinds (tuple) : A sequence of
      :class:`~pysnmp.smi.rfc1902.ObjectType` class instances representing
      MIB variables returned in SNMP response.

    User `error` callback is called with `errorIndication` object wrapped
    in :class:`~twisted.python.failure.Failure` object.

    Examples
    --------
    >>> from twisted.internet.task import react
    >>> from pysnmp.hlapi.twisted import *
    >>>
    >>> def success(args):
    ...     (errorStatus, errorIndex, varBinds) = args
    ...     print(errorStatus, errorIndex, varBind)
    ...
    >>> def failure(errorIndication):
    ...     print(errorIndication)
    ...
    >>> def run(reactor):
    ...     d = setCmd(SnmpEngine(),
    ...                CommunityData('public'),
    ...                UdpTransportTarget(('demo.snmplabs.com', 161)),
    ...                ContextData(),
    ...                ObjectType(ObjectIdentity('SNMPv2-MIB', 'sysDescr', 0), 'Linux i386')
    ...     d.addCallback(success).addErrback(failure)
    ...     return d
    ...
    >>> react(run)
    (0, 0, [ObjectType(ObjectIdentity(ObjectName('1.3.6.1.2.1.1.1.0')), DisplayString('Linux i386'))])
    >>>

    """
    def __cbFun(snmpEngine, sendRequestHandle,
                errorIndication, errorStatus, errorIndex,
                varBinds, cbCtx):
        lookupMib, deferred = cbCtx
        if errorIndication:
            deferred.errback(Failure(errorIndication))
        else:
            try:
                varBindsUnmade = vbProcessor.unmakeVarBinds(snmpEngine.cache, varBinds, lookupMib)

            except Exception as e:
                deferred.errback(Failure(e))

            else:
                deferred.callback((errorStatus, errorIndex, varBindsUnmade))

    addrName, paramsName = lcd.configure(
        snmpEngine, authData, transportTarget, contextData.contextName)

    deferred = Deferred()

    cmdgen.SetCommandGenerator().sendVarBinds(
        snmpEngine, addrName, contextData.contextEngineId,
        contextData.contextName,
        vbProcessor.makeVarBinds(snmpEngine.cache, varBinds), __cbFun,
        (options.get('lookupMib', True), deferred)
    )
    return deferred

def nextCmd(snmpEngine, authData, transportTarget, contextData,
            *varBinds, **options):
    """Performs SNMP GETNEXT query.

    Based on passed parameters, prepares SNMP GETNEXT packet
    (:RFC:`1905#section-4.2.2`) and schedules its transmission by
    :mod:`twisted` I/O framework at a later point of time.

    Parameters
    ----------
    snmpEngine : :class:`~pysnmp.hlapi.SnmpEngine`
        Class instance representing SNMP engine.

    authData : :class:`~pysnmp.hlapi.CommunityData` or :class:`~pysnmp.hlapi.UsmUserData`
        Class instance representing SNMP credentials.

    transportTarget : :class:`~pysnmp.hlapi.twisted.UdpTransportTarget` or :class:`~pysnmp.hlapi.twisted.Udp6TransportTarget`
        Class instance representing transport type along with SNMP peer address.

    contextData : :class:`~pysnmp.hlapi.ContextData`
        Class instance representing SNMP ContextEngineId and ContextName values.

    \*varBinds : :class:`~pysnmp.smi.rfc1902.ObjectType`
        One or more class instances representing MIB variables to place
        into SNMP request.

    Other Parameters
    ----------------
    \*\*options :
        Request options:

            * `lookupMib` - load MIB and resolve response MIB variables at
              the cost of slightly reduced performance. Default is `True`.
            * `ignoreNonIncreasingOid` - continue iteration even if response
              MIB variables (OIDs) are not greater then request MIB variables.
              Be aware that setting it to `True` may cause infinite loop between
              SNMP management and agent applications. Default is `False`.

    Returns
    -------
    deferred : :class:`~twisted.internet.defer.Deferred`
        Twisted Deferred object representing work-in-progress. User
        is expected to attach his own `success` and `error` callback
        functions to the Deferred object though
        :meth:`~twisted.internet.defer.Deferred.addCallbacks` method.

    Raises
    ------
    PySnmpError
        Or its derivative indicating that an error occurred while
        performing SNMP operation.

    Notes
    -----
    User `success` callback is called with the following tuple as
    its first argument:

    * errorStatus (str) : True value indicates SNMP PDU error.
    * errorIndex (int) : Non-zero value refers to `varBinds[errorIndex-1]`
    * varBinds (tuple) :
        A sequence of sequences (e.g. 2-D array) of
        :py:class:`~pysnmp.smi.rfc1902.ObjectType` class instances
        representing a table of MIB variables returned in SNMP response.
        Inner sequences represent table rows and ordered exactly the same
        as `varBinds` in request. Response to GETNEXT always contain
        a single row.

    User `error` callback is called with `errorIndication` object wrapped
    in :class:`~twisted.python.failure.Failure` object.

    Examples
    --------
    >>> from twisted.internet.task import react
    >>> from pysnmp.hlapi.twisted import *
    >>>
    >>> def success(args):
    ...     (errorStatus, errorIndex, varBindTable) = args
    ...     print(errorStatus, errorIndex, varBindTable)
    ...
    >>> def failure(errorIndication):
    ...     print(errorIndication)
    ...
    >>> def run(reactor):
    ...     d = nextCmd(SnmpEngine(),
    ...                 CommunityData('public'),
    ...                 UdpTransportTarget(('demo.snmplabs.com', 161)),
    ...                 ContextData(),
    ...                 ObjectType(ObjectIdentity('SNMPv2-MIB', 'system'))
    ...     d.addCallback(success).addErrback(failure)
    ...     return d
    ...
    >>> react(run)
    (0, 0, [[ObjectType(ObjectIdentity(ObjectName('1.3.6.1.2.1.1.1.0')), DisplayString('SunOS zeus.snmplabs.com 4.1.3_U1 1 sun4m'))]])
    >>>

    """
    def __cbFun(snmpEngine, sendRequestHandle,
                errorIndication, errorStatus, errorIndex,
                varBindTable, cbCtx):
        lookupMib, deferred = cbCtx
        if (options.get('ignoreNonIncreasingOid', False) and
                errorIndication and isinstance(errorIndication, errind.OidNotIncreasing)):
            errorIndication = None
        if errorIndication:
            deferred.errback(Failure(errorIndication))
        else:
            try:
                varBindsUnmade = [vbProcessor.unmakeVarBinds(snmpEngine.cache,
                                                             varBindTableRow,
                                                             lookupMib)
                                  for varBindTableRow in varBindTable]

            except Exception as e:
                deferred.errback(Failure(e))

            else:
                deferred.callback((errorStatus, errorIndex, varBindsUnmade))

    addrName, paramsName = lcd.configure(
        snmpEngine, authData, transportTarget, contextData.contextName)

    deferred = Deferred()

    cmdgen.NextCommandGenerator().sendVarBinds(
        snmpEngine, addrName, contextData.contextEngineId,
        contextData.contextName,
        vbProcessor.makeVarBinds(snmpEngine.cache, varBinds), __cbFun,
        (options.get('lookupMib', True), deferred)
    )
    return deferred

def bulkCmd(snmpEngine, authData, transportTarget, contextData,
            nonRepeaters, maxRepetitions, *varBinds, **options):
    """Performs SNMP GETBULK query.

    Based on passed parameters, prepares SNMP GETNEXT packet
    (:RFC:`1905#section-4.2.3`) and schedules its transmission by
    :mod:`twisted` I/O framework at a later point of time.

    Parameters
    ----------
    snmpEngine : :class:`~pysnmp.hlapi.SnmpEngine`
        Class instance representing SNMP engine.

    authData : :class:`~pysnmp.hlapi.CommunityData` or :class:`~pysnmp.hlapi.UsmUserData`
        Class instance representing SNMP credentials.

    transportTarget : :class:`~pysnmp.hlapi.twisted.UdpTransportTarget` or :class:`~pysnmp.hlapi.twisted.Udp6TransportTarget`
        Class instance representing transport type along with SNMP peer address.

    contextData : :class:`~pysnmp.hlapi.ContextData`
        Class instance representing SNMP ContextEngineId and ContextName values.

    nonRepeaters : int
        One MIB variable is requested in response for the first
        `nonRepeaters` MIB variables in request.

    maxRepetitions : int
        `maxRepetitions` MIB variables are requested in response for each
        of the remaining MIB variables in the request (e.g. excluding
        `nonRepeaters`). Remote SNMP engine may choose lesser value than
        requested.

    \*varBinds : :class:`~pysnmp.smi.rfc1902.ObjectType`
        One or more class instances representing MIB variables to place
        into SNMP request.

    Other Parameters
    ----------------
    \*\*options :
        Request options:

            * `lookupMib` - load MIB and resolve response MIB variables at
              the cost of slightly reduced performance. Default is `True`.
            * `ignoreNonIncreasingOid` - continue iteration even if response
              MIB variables (OIDs) are not greater then request MIB variables.
              Be aware that setting it to `True` may cause infinite loop between
              SNMP management and agent applications. Default is `False`.

    Returns
    -------
    deferred : :class:`~twisted.internet.defer.Deferred`
        Twisted Deferred object representing work-in-progress. User
        is expected to attach his own `success` and `error` callback
        functions to the Deferred object though
        :meth:`~twisted.internet.defer.Deferred.addCallbacks` method.

    Raises
    ------
    PySnmpError
        Or its derivative indicating that an error occurred while
        performing SNMP operation.

    Notes
    -----
    User `success` callback is called with the following tuple as
    its first argument:

    * errorStatus (str) : True value indicates SNMP PDU error.
    * errorIndex (int) : Non-zero value refers to `varBinds[errorIndex-1]`
    * varBindTable (tuple):
      A sequence of sequences (e.g. 2-D array) of
      :py:class:`~pysnmp.smi.rfc1902.ObjectType` class instances representing a
      table of MIB variables returned in SNMP response, with up to
      ``maxRepetitions`` rows, i.e. ``len(varBindTable) <= maxRepetitions``.

      For ``0 <= i < len(varBindTable)`` and ``0 <= j < len(varBinds)``,
      ``varBindTable[i][j]`` represents:

      - For non-repeaters (``j < nonRepeaters``), the first lexicographic
        successor of ``varBinds[j]``, regardless the value of ``i``, or an
        :py:class:`~pysnmp.smi.rfc1902.ObjectType` instance with the
        :py:obj:`~pysnmp.proto.rfc1905.endOfMibView` value if no such successor
        exists;
      - For repeaters (``j >= nonRepeaters``), the ``i``-th lexicographic
        successor of ``varBinds[j]``, or an
        :py:class:`~pysnmp.smi.rfc1902.ObjectType` instance with the
        :py:obj:`~pysnmp.proto.rfc1905.endOfMibView` value if no such successor
        exists.

      See :rfc:`3416#section-4.2.3` for details on the underlying
      ``GetBulkRequest-PDU`` and the associated ``GetResponse-PDU``, such as
      specific conditions under which the server may truncate the response,
      causing ``varBindTable`` to have less than ``maxRepetitions`` rows.

    User `error` callback is called with `errorIndication` object wrapped
    in :class:`~twisted.python.failure.Failure` object.

    Examples
    --------
    >>> from twisted.internet.task import react
    >>> from pysnmp.hlapi.twisted import *
    >>>
    >>> def success(args):
    ...     (errorStatus, errorIndex, varBindTable) = args
    ...     print(errorStatus, errorIndex, varBindTable)
    ...
    >>> def failure(errorIndication):
    ...     print(errorIndication)
    ...
    >>> def run(reactor):
    ...     d = bulkCmd(SnmpEngine(),
    ...                 CommunityData('public'),
    ...                 UdpTransportTarget(('demo.snmplabs.com', 161)),
    ...                 ContextData(),
    ...                 0, 2,
    ...                 ObjectType(ObjectIdentity('SNMPv2-MIB', 'system'))
    ...     d.addCallback(success).addErrback(failure)
    ...     return d
    ...
    >>> react(run)
    (0, 0, [[ObjectType(ObjectIdentity(ObjectName('1.3.6.1.2.1.1.1.0')), DisplayString('SunOS zeus.snmplabs.com 4.1.3_U1 1 sun4m'))], [ObjectType(ObjectIdentity(ObjectName('1.3.6.1.2.1.1.2.0')), ObjectIdentifier('1.3.6.1.4.1.424242.1.1'))]])
    >>>

    """
    def __cbFun(snmpEngine, sendRequestHandle,
                errorIndication, errorStatus, errorIndex,
                varBindTable, cbCtx):
        lookupMib, deferred = cbCtx
        if options.get('ignoreNonIncreasingOid', False) and errorIndication and \
                isinstance(errorIndication, errind.OidNotIncreasing):
            errorIndication = None
        if errorIndication:
            deferred.errback(Failure(errorIndication))
        else:
            try:
                varBindsUnmade = [vbProcessor.unmakeVarBinds(snmpEngine.cache,
                                                             varBindTableRow,
                                                             lookupMib)
                                  for varBindTableRow in varBindTable]

            except Exception as e:
                deferred.errback(Failure(e))

            else:
                deferred.callback((errorStatus, errorIndex, varBindsUnmade))

    addrName, paramsName = lcd.configure(
        snmpEngine, authData, transportTarget, contextData.contextName)

    deferred = Deferred()

    cmdgen.BulkCommandGenerator().sendVarBinds(
        snmpEngine, addrName, contextData.contextEngineId,
        contextData.contextName, nonRepeaters, maxRepetitions,
        vbProcessor.makeVarBinds(snmpEngine.cache, varBinds),
        __cbFun,
        (options.get('lookupMib', True), deferred)
    )
    return deferred