# coding=utf-8

"""
SNMPCollector is a special collector for collecting data by using SNMP

#### Dependencies

 * pysnmp

"""

import socket
import warnings

import diamond.collector

# pysnmp packages on debian 6.0 use sha and md5 which are deprecated
# packages. there is nothing to be done about it until pysnmp
# updates to use new hashlib module -- ignoring warning for now
old_showwarning = warnings.showwarning
warnings.filterwarnings("ignore", category=DeprecationWarning)

cmdgen = None

try:
    import pysnmp.entity.rfc3413.oneliner.cmdgen as cmdgen
    import pysnmp.debug
except ImportError:
    pysnmp = None
    cmdgen = None

warnings.showwarning = old_showwarning


class SNMPCollector(diamond.collector.Collector):
    def __init__(self, *args, **kwargs):
        super(SNMPCollector, self).__init__(*args, **kwargs)

        if cmdgen is not None:
            self.snmpCmdGen = cmdgen.CommandGenerator()

    def get_default_config_help(self):
        config_help = super(SNMPCollector, self).get_default_config_help()
        config_help.update(
            {
                "timeout": "Seconds before timing out the snmp connection",
                "retries": "Number of times to retry before bailing",
            }
        )
        return config_help

    def get_default_config(self):
        # Initialize default config
        default_config = super(SNMPCollector, self).get_default_config()
        default_config["path_suffix"] = ""
        default_config["path_prefix"] = "systems"
        default_config["timeout"] = 5
        default_config["retries"] = 3
        # Return default config
        return default_config

    def _convert_to_oid(self, s):
        d = s.split(".")
        return tuple([int(x) for x in d])

    def _convert_from_oid(self, oid):
        return ".".join([str(x) for x in oid])

    def collect(self):
        for device in self.config["devices"]:
            host = self.config["devices"][device]["host"]
            port = self.config["devices"][device]["port"]
            community = self.config["devices"][device]["community"]
            self.collect_snmp(device, host, port, community)

    def get(self, oid, host, port, community):
        """
        Perform SNMP get for a given OID
        """
        # Initialize return value
        ret = {}

        # Convert OID to tuple if necessary
        if not isinstance(oid, tuple):
            oid = self._convert_to_oid(oid)

        # Convert Host to IP if necessary
        host = socket.gethostbyname(host)

        # Assemble SNMP Auth Data
        snmp_auth_data = cmdgen.CommunityData("agent-{}".format(community), community)

        # Assemble SNMP Transport Data
        snmp_transport_data = cmdgen.UdpTransportTarget(
            (host, port), int(self.config["timeout"]), int(self.config["retries"])
        )

        # Assemble SNMP Next Command
        result = self.snmpCmdGen.getCmd(snmp_auth_data, snmp_transport_data, oid)
        var_bind = result[3]

        # TODO: Error check

        for o, v in var_bind:
            ret[str(o)] = v.prettyPrint()

        return ret

    def walk(self, oid, host, port, community):
        """
        Perform an SNMP walk on a given OID
        """
        # Initialize return value
        ret = {}

        # Convert OID to tuple if necessary
        if not isinstance(oid, tuple):
            oid = self._convert_to_oid(oid)

        # Convert Host to IP if necessary
        host = socket.gethostbyname(host)

        # Assemble SNMP Auth Data
        snmp_auth_data = cmdgen.CommunityData("agent-{}".format(community), community)

        # Assemble SNMP Transport Data
        snmp_transport_data = cmdgen.UdpTransportTarget(
            (host, port), int(self.config["timeout"]), int(self.config["retries"])
        )

        # Assemble SNMP Next Command
        result_table = self.snmpCmdGen.nextCmd(snmp_auth_data, snmp_transport_data, oid)
        var_bind_table = result_table[3]

        # TODO: Error Check

        for varBindTableRow in var_bind_table:
            for o, v in varBindTableRow:
                ret[str(o)] = v.prettyPrint()

        return ret
