<commit_before>from checks import AgentCheck

from collections import defaultdict

from pysnmp.entity.rfc3413.oneliner import cmdgen
from pysnmp.smi.exval import noSuchInstance
import pysnmp.proto.rfc1902 as snmp_type

snmp_counters = [snmp_type.Counter32, snmp_type.Counter64]
snmp_gauges = [snmp_type.Gauge32]

class SnmpCheck(AgentCheck):

    def __init__(self, name, init_config, agentConfig, instances=None):
        AgentCheck.__init__(self, name, init_config, agentConfig, instances)
        self.counter_state = defaultdict(dict)
        self.interface_list = {}
        if instances is not None:
            for instance in instances:
                if 'ip_address' in instance:
                    ip_address = instance["ip_address"]
                    self.interface_list[ip_address] = self.get_interfaces(instance)
                    tags = instance.get("tags",[])
                    tags.append("snmp_device:" + ip_address)

    def get_interfaces(self, instance):

        interface_list = {}

        def get_interfaces_nb():
            result = SnmpCheck.snmp_get(instance, [(("IF-MIB","ifNumber"),0)])[0]
            if noSuchInstance.isSameTypeWith(result[1]):
                return None
            else:
                return int(result[1])

        interface_nb = get_interfaces_nb()
        if interface_nb is not None:
            interfaces_descr_oids = []
            for interface in range(interface_nb):
                interface_index = interface + 1 #SNMP indexes start from 1
                interfaces_descr_oids.append((("IF-MIB","ifDescr"),interface_index))
                interfaces_descr_oids.append((("IF-MIB","ifType"),interface_index))

            interfaces_description = SnmpCheck.snmp_get(instance, interfaces_descr_oids)
            for i in range(interface_nb):
                # order is guaranteed
                descr = str(interfaces_description.pop(0)[1])
                type = int(interfaces_description.pop(0)[1])
                if type != 24:
                    interface_list[i+1] = descr
        else:
            empty_reply = False
            interface_index = 1
            while not empty_reply:
                interfaces_descr_oids = []
                interfaces_descr_oids.append((("IF-MIB","ifDescr"),interface_index))
                interfaces_descr_oids.append((("IF-MIB","ifType"),interface_index))
                interfaces_description = SnmpCheck.snmp_get(instance, interfaces_descr_oids)
                descr = interfaces_description.pop(0)[1]
                if noSuchInstance.isSameTypeWith(descr):
                    empty_reply= True
                else:
                    type = int(interfaces_description.pop(0)[1])
                    if type != 24:
                        interface_list[interface_index] = str(descr)
                    interface_index += 1

        return interface_list

    @staticmethod
    def get_auth_data(instance):
        if "community_string" in instance:
            return cmdgen.CommunityData(instance['community_string'])
        elif "user" in instance:
            user = instance["user"]
            authKey = None
            privKey = None
            authProtocol = None
            privProtocol = None
            if "authKey" in instance:
                authKey = instance["authKey"]
                authProtocol = cmdgen.usmHMACMD5AuthProtocol
            if "privKey" in instance:
                privKey = instance["privKey"]
                authProtocol = cmdgen.usmHMACMD5AuthProtocol
                privProtocol = cmdgen.usmDESPrivProtocol
            if "authProtocol" in instance:
                authProtocol = getattr(cmdgen,instance["authProtocol"])
            if "privProtocol" in instance:
                privProtocol = getattr(cmdgen,instance["privProtocol"])
            return cmdgen.UsmUserData(user, authKey, privKey, authProtocol, privProtocol)
        else:
            raise Exception("An authentication method needs to be provided")

    @staticmethod
    def get_transport_target(instance):
        if "ip_address" not in instance:
            raise Exception("An IP address needs to be specified")
        ip_address = instance["ip_address"]
        port = instance.get("port", 161)
        return cmdgen.UdpTransportTarget((ip_address,port))

    def check(self, instance):
        tags = instance.get("tags",[])
        device_oids = []
        interface_oids = []
        for metric in instance.get('metrics',[]):
            device_oids.append(((metric["MIB"],metric["symbol"]),metric["index"]))
        results = SnmpCheck.snmp_get(instance, device_oids)
        for oid, value in results:
            self.report_as_statsd(instance, oid, value, tags=tags)

        for metric in instance.get('interface_metrics',[]):
            interface_oids.append((metric["MIB"],metric["symbol"]))
        for interface, descr in self.interface_list[instance['ip_address']].items():
            oids = [(oid, interface) for oid in interface_oids]
            interface_results = SnmpCheck.snmp_get(instance, oids)
            for oid, value in interface_results:
                self.report_as_statsd(instance, oid, value, tags = tags + ["interface:"+descr])

    @staticmethod
    def snmp_get(instance, oids):
        """
        Perform a snmp get command to the device that instance
        describe and return a list of tuble (name, values)
        corresponding to the elements in oids.
        """
        transport_target = SnmpCheck.get_transport_target(instance)
        auth_data = SnmpCheck.get_auth_data(instance)<commit_msg>SNMP| Added docstring<commit_after>from checks import AgentCheck

from collections import defaultdict

from pysnmp.entity.rfc3413.oneliner import cmdgen
from pysnmp.smi.exval import noSuchInstance
import pysnmp.proto.rfc1902 as snmp_type

snmp_counters = [snmp_type.Counter32, snmp_type.Counter64]
snmp_gauges = [snmp_type.Gauge32]

class SnmpCheck(AgentCheck):

    def __init__(self, name, init_config, agentConfig, instances=None):
        AgentCheck.__init__(self, name, init_config, agentConfig, instances)
        self.counter_state = defaultdict(dict)
        self.interface_list = {}
        if instances is not None:
            for instance in instances:
                if 'ip_address' in instance:
                    ip_address = instance["ip_address"]
                    self.interface_list[ip_address] = self.get_interfaces(instance)
                    tags = instance.get("tags",[])
                    tags.append("snmp_device:" + ip_address)

    def get_interfaces(self, instance):
        '''
        Return all the network interfaces of an instance to be used to get metrics
        on those interfaces.
        If available, get the number of interfaces from the ifNumber MIB
        and then query all of them in one request for their description.
        If this info is not available, repeatedly query the interface description
        in order to discover them all.
        '''

        interface_list = {}

        def get_interfaces_nb():
            result = SnmpCheck.snmp_get(instance, [(("IF-MIB","ifNumber"),0)])[0]
            if noSuchInstance.isSameTypeWith(result[1]):
                return None
            else:
                return int(result[1])

        interface_nb = get_interfaces_nb()
        if interface_nb is not None:
            interfaces_descr_oids = []
            for interface in range(interface_nb):
                interface_index = interface + 1 #SNMP indexes start from 1
                interfaces_descr_oids.append((("IF-MIB","ifDescr"),interface_index))
                interfaces_descr_oids.append((("IF-MIB","ifType"),interface_index))

            interfaces_description = SnmpCheck.snmp_get(instance, interfaces_descr_oids)
            for i in range(interface_nb):
                # order is guaranteed
                descr = str(interfaces_description.pop(0)[1])
                type = int(interfaces_description.pop(0)[1])
                if type != 24: # ignore localhost loopback
                    interface_list[i+1] = descr
        else:
            # ifNumber is not available, query the interfaces dor description
            # until blank one
            empty_reply = False
            interface_index = 1
            while not empty_reply:
                interfaces_descr_oids = []
                interfaces_descr_oids.append((("IF-MIB","ifDescr"),interface_index))
                interfaces_descr_oids.append((("IF-MIB","ifType"),interface_index))
                interfaces_description = SnmpCheck.snmp_get(instance, interfaces_descr_oids)
                descr = interfaces_description.pop(0)[1]
                if noSuchInstance.isSameTypeWith(descr):
                    empty_reply= True
                else:
                    type = int(interfaces_description.pop(0)[1])
                    if type != 24: # ignore localhost loopback
                        interface_list[interface_index] = str(descr)
                    interface_index += 1

        return interface_list

    @staticmethod
    def get_auth_data(instance):
        if "community_string" in instance:
            # SNMP v1 - SNMP v2
            return cmdgen.CommunityData(instance['community_string'])
        elif "user" in instance:
            # SNMP v3
            user = instance["user"]
            authKey = None
            privKey = None
            authProtocol = None
            privProtocol = None
            if "authKey" in instance:
                authKey = instance["authKey"]
                authProtocol = cmdgen.usmHMACMD5AuthProtocol
            if "privKey" in instance:
                privKey = instance["privKey"]
                authProtocol = cmdgen.usmHMACMD5AuthProtocol
                privProtocol = cmdgen.usmDESPrivProtocol
            if "authProtocol" in instance:
                authProtocol = getattr(cmdgen,instance["authProtocol"])
            if "privProtocol" in instance:
                privProtocol = getattr(cmdgen,instance["privProtocol"])
            return cmdgen.UsmUserData(user, authKey, privKey, authProtocol, privProtocol)
        else:
            raise Exception("An authentication method needs to be provided")

    @staticmethod
    def get_transport_target(instance):
        if "ip_address" not in instance:
            raise Exception("An IP address needs to be specified")
        ip_address = instance["ip_address"]
        port = instance.get("port", 161) # Default SNMP port
        return cmdgen.UdpTransportTarget((ip_address,port))

    def check(self, instance):
        tags = instance.get("tags",[])
        device_oids = []
        interface_oids = []
        for metric in instance.get('metrics',[]):
            device_oids.append(((metric["MIB"],metric["symbol"]),metric["index"]))
        results = SnmpCheck.snmp_get(instance, device_oids)
        for oid, value in results:
            self.report_as_statsd(instance, oid, value, tags=tags)

        for metric in instance.get('interface_metrics',[]):
            interface_oids.append((metric["MIB"],metric["symbol"]))
        for interface, descr in self.interface_list[instance['ip_address']].items():
            oids = [(oid, interface) for oid in interface_oids]
            interface_results = SnmpCheck.snmp_get(instance, oids)
            for oid, value in interface_results:
                self.report_as_statsd(instance, oid, value, tags = tags + ["interface:"+descr])

    @staticmethod
    def snmp_get(instance, oids):
        """
        Perform a snmp get command to the device that instance
        describe and return a list of tuble (name, values)
        corresponding to the elements in oids.
        """
        transport_target = SnmpCheck.get_transport_target(instance)
        auth_data = SnmpCheck.get_auth_data(instance)
{'fraction_non_alphanumeric': 0.0899113550021106, 'fraction_numerical': 0.007007176023638666, 'mean_word_length': 2.246368868183064, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1627333', 'n_tokens_mistral': 3272, 'n_tokens_neox': 3109, 'n_words': 745}