#
# snimpy -- Interactive SNMP tool
#
# Copyright (C) Vincent Bernat <bernat@luffy.cx>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

"""
This module is a low-level interface to build SNMP requests, send
them and receive answers. It is built on top of pysnmp_ but the
exposed interface is far simpler. It is also far less complete and
there is an important dependency to the :mod:`basictypes` module for
type coercing.

.. _pysnmp: http://pysnmp.sourceforge.net/
"""

import re
import socket
import inspect
import threading
import asyncio
import ipaddress
import pysnmp.hlapi.v3arch.asyncio as v3
from pyasn1.type import univ
from pysnmp.proto import rfc1902, rfc1905
from pysnmp.smi import error
from pysnmp.smi.rfc1902 import ObjectType, ObjectIdentity


class SNMPException(Exception):
    """SNMP related base exception. All SNMP exceptions are inherited from
    this one. The inherited exceptions are named after the name of the
    corresponding SNMP error.
    """


class SNMPTooBig(SNMPException):
    pass


class SNMPNoSuchName(SNMPException):
    pass


class SNMPBadValue(SNMPException):
    pass


class SNMPReadOnly(SNMPException):
    pass


# Dynamically build remaining (v2) exceptions
for name, obj in inspect.getmembers(error):
    if name.endswith("Error") and \
       inspect.isclass(obj) and \
       issubclass(obj, error.MibOperationError) and \
       obj != error.MibOperationError:
        name = str("SNMP{}".format(name[:-5]))
        globals()[name] = type(name, (SNMPException,), {})
del name
del obj


class _LoopGuard:
    """Cancel pending tasks and close the event loop on thread exit.

    pysnmp's AsyncioDispatcher creates a background handle_timeout()
    task that runs forever.  When a thread exits, this task must be
    properly cancelled and awaited — otherwise asyncio emits "Task
    was destroyed but it is pending!" warnings.

    This guard is stored in thread-local data alongside the engine
    and event loop.  Because nothing else references it, it is the
    first object to reach refcount 0 when the thread-local dict is
    cleared, so its __del__ runs before the engine's dispatcher
    tries to clean up the same tasks."""

    def __init__(self, loop):
        self._loop = loop

    def __del__(self):
        loop = self._loop
        if loop.is_closed():
            return
        pending = asyncio.all_tasks(loop)
        for task in pending:
            task.cancel()
        if pending:
            try:
                loop.run_until_complete(
                    asyncio.gather(*pending, return_exceptions=True))
            except RuntimeError:
                pass
        loop.close()


class _SnimpyEngine:
    """Manage per-thread SNMP backends and event loop."""

    _tls = threading.local()

    @classmethod
    def loop(cls):
        """Return the per-thread event loop, creating one if needed."""
        if not hasattr(cls._tls, "loop"):
            cls._tls.loop = asyncio.new_event_loop()
            asyncio.set_event_loop(cls._tls.loop)
            cls._tls.guard = _LoopGuard(cls._tls.loop)
        return cls._tls.loop

    @classmethod
    def engine(cls):
        """Return the per-thread v3arch SnmpEngine."""
        cls.loop()
        if not hasattr(cls._tls, "engine"):
            cls._tls.engine = v3.SnmpEngine()
        return cls._tls.engine


class Session:

    """SNMP session. An instance of this object will represent an SNMP
    session. From such an instance, one can get information from the
    associated agent."""

    def _run(self, coro):
        """Run an async coroutine synchronously."""
        return _SnimpyEngine.loop().run_until_complete(coro)

    def __init__(self, host,
                 community="public", version=2,
                 secname=None,
                 authprotocol=None,
                 authpassword=None,
                 privprotocol=None,
                 privpassword=None,
                 contextname=None,
                 bulk=40,
                 none=False):
        """Create a new SNMP session.

        :param host: The hostname or IP address of the agent to
            connect to. Optionally, the port can be specified
            separated with a double colon.
        :type host: str
        :param community: The community to transmit to the agent for
            authorization purpose. This parameter is ignored if the
            specified version is 3.
        :type community: str
        :param version: The SNMP version to use to talk with the
            agent. Possible values are `1`, `2` (community-based) or
            `3`.
        :type version: int
        :param secname: Security name to use for SNMPv3 only.
        :type secname: str
        :param authprotocol: Authorization protocol to use for
            SNMPv3. This can be `None` or one of the strings `SHA`,
            `MD5`, `SHA224`, `SHA256`, `SHA384` or `SHA512`.
        :type authprotocol: None or str
        :param authpassword: Authorization password if authorization
            protocol is not `None`.
        :type authpassword: str
        :param privprotocol: Privacy protocol to use for SNMPv3. This
            can be `None` or either the string `AES`, `AES128`,
            `AES192`, `AES256` or `3DES`.
        :type privprotocol: None or str
        :param privpassword: Privacy password if privacy protocol is
            not `None`.
        :type contextname: str
        :param contextname: Context name for SNMPv3 messages.
        :type privpassword: str
        :param bulk: Max repetition value for `GETBULK` requests. Set
            to `0` to disable.
        :type bulk: int
        :param none: When enabled, will return None for not found
            values (instead of raising an exception)
        :type none: bool
        """
        self._host = host
        self._version = version
        self._none = none
        if version == 1 and none:
            raise ValueError("None-GET requests not compatible with SNMPv1")

        # Put authentication stuff in self._auth and select backend
        if version in [1, 2]:
            self._auth = v3.CommunityData(community, mpModel=version - 1)
            self._contextdata = v3.ContextData()
            self._cmd_args = (_SnimpyEngine.engine(), self._auth)
            self._get_cmd = v3.get_cmd
            self._set_cmd = v3.set_cmd
            self._walk_cmd = v3.walk_cmd
            self._bulk_walk_cmd = v3.bulk_walk_cmd
            UdpTarget = v3.UdpTransportTarget
            Udp6Target = v3.Udp6TransportTarget
        elif version == 3:
            if secname is None:
                secname = community
            try:
                authprotocol = {
                    None: v3.usmNoAuthProtocol,
                    "MD5": v3.usmHMACMD5AuthProtocol,
                    "SHA": v3.usmHMACSHAAuthProtocol,
                    "SHA1": v3.usmHMACSHAAuthProtocol,
                    "SHA224": v3.usmHMAC128SHA224AuthProtocol,
                    "SHA256": v3.usmHMAC192SHA256AuthProtocol,
                    "SHA384": v3.usmHMAC256SHA384AuthProtocol,
                    "SHA512": v3.usmHMAC384SHA512AuthProtocol,
                }[authprotocol]
            except KeyError:
                raise ValueError("{} is not an acceptable authentication "
                                 "protocol".format(authprotocol))
            try:
                privprotocol = {
                    None: v3.usmNoPrivProtocol,
                    "DES": v3.usmDESPrivProtocol,
                    "3DES": v3.usm3DESEDEPrivProtocol,
                    "AES": v3.usmAesCfb128Protocol,
                    "AES128": v3.usmAesCfb128Protocol,
                    "AES192": v3.usmAesCfb192Protocol,
                    "AES256": v3.usmAesCfb256Protocol,
                }[privprotocol]
            except KeyError:
                raise ValueError("{} is not an acceptable privacy "
                                 "protocol".format(privprotocol))
            self._auth = v3.UsmUserData(secname,
                                        authpassword,
                                        privpassword,
                                        authprotocol,
                                        privprotocol)
            if contextname:
                contextdata = v3.ContextData(
                    contextName=rfc1902.OctetString(contextname))
            else:
                contextdata = v3.ContextData()
            self._cmd_args = (_SnimpyEngine.engine(), self._auth)
            self._contextdata = contextdata
            self._get_cmd = v3.get_cmd
            self._set_cmd = v3.set_cmd
            self._walk_cmd = v3.walk_cmd
            self._bulk_walk_cmd = v3.bulk_walk_cmd
            UdpTarget = v3.UdpTransportTarget
            Udp6Target = v3.Udp6TransportTarget
        else:
            raise ValueError("unsupported SNMP version {}".format(version))

        # Put transport stuff into self._transport
        mo = re.match(r'^(?:'
                      r'\[(?P<ipv6>[\d:A-Fa-f]+)\]|'
                      r'(?P<ipv4>[\d\.]+)|'
                      r'(?P<any>.*?))'
                      r'(?::(?P<port>\d+))?$',
                      host)
        if mo.group("port"):
            port = int(mo.group("port"))
        else:
            port = 161
        if mo.group("ipv6"):
            self._transport = self._run(
                Udp6Target.create((mo.group("ipv6"), port)))
        elif mo.group("ipv4"):
            self._transport = self._run(
                UdpTarget.create((mo.group("ipv4"), port)))
        else:
            results = socket.getaddrinfo(mo.group("any"),
                                         port,
                                         0,
                                         socket.SOCK_DGRAM,
                                         socket.IPPROTO_UDP)
            if [x for x in results if x[0] == socket.AF_INET]:
                self._transport = self._run(
                    UdpTarget.create((mo.group("any"), port)))
            else:
                self._transport = self._run(
                    Udp6Target.create((mo.group("any"), port)))
        self._cmd_args += (self._transport,)

        # Bulk stuff
        self.bulk = bulk

    def _check_exception(self, value):
        """Check if the given ASN1 value is an exception"""
        if isinstance(value, rfc1905.NoSuchObject):
            raise SNMPNoSuchObject("No such object was found")  # noqa: F821
        if isinstance(value, rfc1905.NoSuchInstance):
            raise SNMPNoSuchInstance("No such instance exists")  # noqa: F821
        if isinstance(value, rfc1905.EndOfMibView):
            raise SNMPEndOfMibView("End of MIB was reached")  # noqa: F821

    def _check_error(self, errorIndication, errorStatus):
        """Check for SNMP protocol errors in response"""
        if errorIndication:
            self._check_exception(errorIndication)
            raise SNMPException(str(errorIndication))
        if errorStatus:
            exc = str(errorStatus.prettyPrint())
            exc = re.sub(r'\W+', '', exc)
            exc = "SNMP{}".format(exc[0].upper() + exc[1:])
            if str(exc) in globals():
                raise globals()[exc]
            raise SNMPException(errorStatus.prettyPrint())

    def _convert(self, value):
        """Convert a PySNMP value to some native Python type"""
        try:
            # With PySNMP 4.3+, an OID is a ObjectIdentity. We try to
            # extract it while being compatible with earlier releases.
            value = value.getOid()
        except AttributeError:
            pass
        if self._none:
            if isinstance(value, rfc1905.NoSuchObject):
                return None
            if isinstance(value, rfc1905.NoSuchInstance):
                return None
        self._check_exception(value)
        convertors = {rfc1902.Integer: int,
                      rfc1902.Integer32: int,
                      rfc1902.OctetString: bytes,
                      rfc1902.IpAddress: ipaddress.IPv4Address,
                      rfc1902.Counter32: int,
                      rfc1902.Counter64: int,
                      rfc1902.Gauge32: int,
                      rfc1902.Unsigned32: int,
                      rfc1902.TimeTicks: int,
                      rfc1902.Bits: str,
                      rfc1902.Opaque: str,
                      rfc1902.univ.ObjectIdentifier: tuple,
                      # v1arch returns raw pyasn1 types
                      univ.Integer: int,
                      univ.OctetString: bytes,
                      univ.ObjectIdentifier: tuple}
        for cl, fn in convertors.items():
            if isinstance(value, cl):
                return fn(value)
        raise NotImplementedError("unable to convert {}".format(repr(value)))

    def get(self, *oids):
        """Retrieve an OID value using GET.

        :param oids: a list of OID to retrieve. An OID is a tuple.
        :return: a list of tuples with the retrieved OID and the raw value.
        """
        objecttypes = [ObjectType(ObjectIdentity(oid)) for oid in oids]
        errorIndication, errorStatus, errorIndex, varBinds = self._run(
            self._get_cmd(*self._cmd_args, self._contextdata,
                          *objecttypes, lookupMib=False))
        self._check_error(errorIndication, errorStatus)
        results = [(tuple(name), self._convert(val))
                   for name, val in varBinds]
        if not results:
            raise SNMPException("empty answer")
        return tuple(results)

    async def _walk_async(self, *oids):
        """Collect results from GETNEXT-based walk."""
        results = []
        for oid in oids:
            walker = self._walk_cmd(
                *self._cmd_args, self._contextdata,
                ObjectType(ObjectIdentity(oid)),
                lookupMib=False, lexicographicMode=False)
            async for result in walker:
                errorIndication, errorStatus, errorIndex, varBinds = result
                self._check_error(errorIndication, errorStatus)
                for name, val in varBinds:
                    results.append((tuple(name), val))
        return results

    async def _bulkwalk_async(self, bulk, *oids):
        """Collect results from GETBULK-based walk."""
        results = []
        for oid in oids:
            walker = self._bulk_walk_cmd(
                *self._cmd_args, self._contextdata, 0, bulk,
                ObjectType(ObjectIdentity(oid)),
                lookupMib=False, lexicographicMode=False)
            async for result in walker:
                errorIndication, errorStatus, errorIndex, varBinds = result
                self._check_error(errorIndication, errorStatus)
                for name, val in varBinds:
                    results.append((tuple(name), val))
        return results

    def walkmore(self, *oids):
        """Retrieve OIDs values using GETBULK or GETNEXT. The method is called
        "walk" but this is either a GETBULK or a GETNEXT. The later is
        only used for SNMPv1 or if bulk has been disabled using
        :meth:`bulk` property.

        :param oids: a list of OID to retrieve. An OID is a tuple.
        :return: a list of tuples with the retrieved OID and the raw value.

        """
        if self._version == 1 or not self.bulk:
            results = self._run(self._walk_async(*oids))
        else:
            try:
                results = self._run(self._bulkwalk_async(self.bulk, *oids))
            except SNMPTooBig:
                # Let's try to ask for less values. We will never increase
                # bulk again. We cannot increase it just after the walk
                # because we may end up requesting everything twice (or
                # more).
                nbulk = self.bulk / 2 or False
                if nbulk != self.bulk:
                    self.bulk = nbulk
                    return self.walk(*oids)
                raise
        return tuple([(oid, self._convert(val)) for oid, val in results])

    def walk(self, *oids):
        """Walk from given OIDs but don't return any "extra" results. Only
        results in the subtree will be returned.

        :param oid: OIDs used as a start point
        :return: a list of tuples with the retrieved OID and the raw value.
        """
        return ((noid, result)
                for oid in oids
                for noid, result in self.walkmore(oid)
                if (len(noid) >= len(oid) and
                    noid[:len(oid)] == oid[:len(oid)]))

    def set(self, *args):
        """Set an OID value using SET. This function takes an odd number of
        arguments. They are working by pair. The first member is an
        OID and the second one is :class:`basictypes.Type` instace
        whose `pack()` method will be used to transform into the
        appropriate form.

        :return: a list of tuples with the retrieved OID and the raw value.
        """
        if len(args) % 2 != 0:
            raise ValueError("expect an even number of arguments for SET")
        objecttypes = [ObjectType(ObjectIdentity(oid), val.pack())
                       for oid, val in zip(args[0::2], args[1::2])]
        errorIndication, errorStatus, errorIndex, varBinds = self._run(
            self._set_cmd(*self._cmd_args, self._contextdata,
                          *objecttypes, lookupMib=False))
        self._check_error(errorIndication, errorStatus)
        results = [(tuple(name), self._convert(val))
                   for name, val in varBinds]
        if not results:
            raise SNMPException("empty answer")
        return tuple(results)

    def __repr__(self):
        return "{}(host={},version={})".format(
            self.__class__.__name__,
            self._host,
            self._version)

    @property
    def timeout(self):
        """Get timeout value for the current session.

        :return: Timeout value in microseconds.
        """
        return self._transport.timeout * 1000000

    @timeout.setter
    def timeout(self, value):
        """Set timeout value for the current session.

        :param value: Timeout value in microseconds.
        """
        value = int(value)
        if value <= 0:
            raise ValueError("timeout is a positive integer")
        self._transport.timeout = value / 1000000.

    @property
    def retries(self):
        """Get number of times a request is retried.

        :return: Number of retries for each request.
        """
        return self._transport.retries

    @retries.setter
    def retries(self, value):
        """Set number of times a request is retried.

        :param value: Number of retries for each request.
        """
        value = int(value)
        if value < 0:
            raise ValueError("retries is a non-negative integer")
        self._transport.retries = value

    @property
    def bulk(self):
        """Get bulk settings.

        :return: `False` if bulk is disabled or a non-negative integer
            for the number of repetitions.
        """
        return self._bulk

    @bulk.setter
    def bulk(self, value):
        """Set bulk settings.

        :param value: `False` to disable bulk or a non-negative
            integer for the number of allowed repetitions.
        """
        if value is False:
            self._bulk = False
            return
        value = int(value)
        if value <= 0:
            raise ValueError("{} is not an appropriate value "
                             "for max repeater parameter".format(
                                 value))
        self._bulk = value
