<commit_before>from checks import AgentCheck

from collections import defaultdict

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

convention_type_builder = builder.MibBuilder()
(CounterBasedGauge64, ZeroBasedCounter64) = convention_type_builder.importSymbols("HCNUM-TC","CounterBasedGauge64", "ZeroBasedCounter64")

SNMP_COUNTERS = [snmp_type.Counter32.__name__, snmp_type.Counter64.__name__, ZeroBasedCounter64.__name__]
SNMP_GAUGES = [snmp_type.Gauge32.__name__, CounterBasedGauge64.__name__]

def reply_invalid(oid):
    return noSuchInstance.isSameTypeWith(oid) or \
           noSuchObject.isSameTypeWith(oid)

class SnmpCheck(AgentCheck):

    cmd_generator = None

    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 = {}
        mibs_path = None
        if init_config is not None:
            mibs_path = init_config.get("mibs_folder")
        SnmpCheck.create_command_generator(mibs_path)
        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)

    @classmethod
    def create_command_generator(cls, mibs_path=None):
        cls.cmd_generator = cmdgen.CommandGenerator()
        if mibs_path is not None:
            mibBuilder = cls.cmd_generator.snmpEngine.msgAndPduDsp.\
                         mibInstrumController.mibBuilder
            mibSources = mibBuilder.getMibSources() + (
                    builder.DirMibSource(mibs_path),
                    )
            mibBuilder.setMibSources(*mibSources)

    def get_interfaces(self, instance):
        '''
        Return all the network interfaces of an instance to be used to get metrics
        on those interfaces.
        Repeatedly query the interface description
        in order to discover them all.
        '''

        interface_list = {}

        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]
            empty_reply= reply_invalid(descr)
            type = interfaces_description.pop(0)[1]
            if not reply_invalid(type) and int(type) !=24:
                # ignore localhost loopback
                interface_list[interface_index] = str(descr)
                self.log.info("Discovered interface %s" % str(descr))
            interface_index += 1
        return interface_list

    @classmethod
    def get_auth_data(cls, 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"]<commit_msg>SNMP Checks | Better interface discovery<commit_after>from checks import AgentCheck

from collections import defaultdict

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

convention_type_builder = builder.MibBuilder()
(CounterBasedGauge64, ZeroBasedCounter64) = convention_type_builder.importSymbols("HCNUM-TC","CounterBasedGauge64", "ZeroBasedCounter64")

SNMP_COUNTERS = [snmp_type.Counter32.__name__, snmp_type.Counter64.__name__, ZeroBasedCounter64.__name__]
SNMP_GAUGES = [snmp_type.Gauge32.__name__, CounterBasedGauge64.__name__]

# IF-MIB magic values
IF_TABLE_TYPE_POS = 9
IF_TABLE_INDEX_POS = 10
IF_INDEX = 1
IF_DESCR = 2
IF_TYPE = 3
LOCALHOST_INTERFACE = 24

def reply_invalid(oid):
    return noSuchInstance.isSameTypeWith(oid) or \
           noSuchObject.isSameTypeWith(oid)

class SnmpCheck(AgentCheck):

    cmd_generator = None

    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 = {}
        mibs_path = None
        if init_config is not None:
            mibs_path = init_config.get("mibs_folder")
        SnmpCheck.create_command_generator(mibs_path)
        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)

    @classmethod
    def create_command_generator(cls, mibs_path=None):
        cls.cmd_generator = cmdgen.CommandGenerator()
        if mibs_path is not None:
            mibBuilder = cls.cmd_generator.snmpEngine.msgAndPduDsp.\
                         mibInstrumController.mibBuilder
            mibSources = mibBuilder.getMibSources() + (
                    builder.DirMibSource(mibs_path),
                    )
            mibBuilder.setMibSources(*mibSources)

    def get_interfaces(self, instance):
        interface_list = {}
        transport_target = self.get_transport_target(instance)
        auth_data = self.get_auth_data(instance)

        snmp_command = self.cmd_generator.nextCmd
        errorIndication, errorStatus, errorIndex, varBinds = snmp_command(
            auth_data,
            transport_target,
            '.1.3.6.1.2.1.2.2.1.',
            lookupValues = True
            )

        ifTable = defaultdict(dict)
        if errorIndication:
            raise Exception("{0} for instance {1}".format(errorIndication,instance["ip_address"]))
        else:
            if errorStatus:
                raise Exception("{0} for instance {1}".format(errorStatus.prettyPrint(),instance["ip_address"]))
            else:
                for tableRow in varBinds:
                    for name, val in tableRow:
                        ifTable[name.asTuple()[IF_TABLE_INDEX_POS]][name.asTuple()[IF_TABLE_TYPE_POS]] = val

        self.log.debug("Interface Table discovered %s" % ifTable)
        for index in ifTable:
            type = ifTable[index].get(IF_TYPE)
            descr = ifTable[index].get(IF_DESCR)
            if not reply_invalid(type):
                if int(type) != LOCALHOST_INTERFACE and not reply_invalid(descr):
                    interface_list[index] = str(descr)
                    self.log.info("Discovered interface %s" % str(descr))

        return interface_list

    @classmethod
    def get_auth_data(cls, 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"]
{'fraction_non_alphanumeric': 0.08167828488754439, 'fraction_numerical': 0.01144285150598448, 'mean_word_length': 2.51061865189289, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, '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': '5347721', 'n_tokens_mistral': 2261, 'n_tokens_neox': 2168, 'n_words': 453}