<commit_before># std
from collections import defaultdict
from functools import wraps

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

# project
from checks import AgentCheck
from config import _is_affirmative



# Additional types that are not part of the SNMP protocol. cf RFC 2856
(CounterBasedGauge64, ZeroBasedCounter64) = builder.MibBuilder().importSymbols(
    "HCNUM-TC",
    "CounterBasedGauge64",
    "ZeroBasedCounter64")

# Metric type that we support
SNMP_COUNTERS = frozenset([
    snmp_type.Counter32.__name__,
    snmp_type.Counter64.__name__,
    ZeroBasedCounter64.__name__])

SNMP_GAUGES = frozenset([
    snmp_type.Gauge32.__name__,
    snmp_type.Unsigned32.__name__,
    CounterBasedGauge64.__name__,
    snmp_type.Integer.__name__,
    snmp_type.Integer32.__name__])

OID_BATCH_SIZE = 10


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


class SnmpCheck(AgentCheck):

    cmd_generator = None
    # pysnmp default values
    DEFAULT_RETRIES = 5
    DEFAULT_TIMEOUT = 1

    def __init__(self, name, init_config, agentConfig, instances=None):
        AgentCheck.__init__(self, name, init_config, agentConfig, instances)

        # Load Custom MIB directory
        mibs_path = None
        ignore_nonincreasing_oid = False
        if init_config is not None:
            mibs_path = init_config.get("mibs_folder")
            ignore_nonincreasing_oid = _is_affirmative(
                init_config.get("ignore_nonincreasing_oid", False))

        # Create SNMP command generator and aliases
        self.create_command_generator(mibs_path, ignore_nonincreasing_oid)

    def snmp_logger(self, func):
        """
        Decorator to log, with DEBUG level, SNMP commands
        """
        @wraps(func)
        def wrapper(*args, **kwargs):
            self.log.debug("Running SNMP command {0} on OIDS {1}"
                           .format(func.__name__, args[2:]))
            result = func(*args, **kwargs)
            self.log.debug("Returned vars: {0}".format(result[-1]))
            return result
        return wrapper

    def create_command_generator(self, mibs_path, ignore_nonincreasing_oid):
        '''
        Create a command generator to perform all the snmp query.
        If mibs_path is not None, load the mibs present in the custom mibs
        folder. (Need to be in pysnmp format)
        '''
        self.cmd_generator = cmdgen.CommandGenerator()
        self.cmd_generator.ignoreNonIncreasingOid = ignore_nonincreasing_oid
        if mibs_path is not None:
            mib_builder = self.cmd_generator.snmpEngine.msgAndPduDsp.\
                mibInstrumController.mibBuilder
            mib_sources = mib_builder.getMibSources() + \
                (builder.DirMibSource(mibs_path), )
            mib_builder.setMibSources(*mib_sources)

        # Set aliases for snmpget and snmpgetnext with logging
        self.snmpget = self.snmp_logger(self.cmd_generator.getCmd)
        self.snmpgetnext = self.snmp_logger(self.cmd_generator.nextCmd)

    @classmethod
    def get_auth_data(cls, instance):
        '''
        Generate a Security Parameters object based on the instance's
        configuration.
        See http://pysnmp.sourceforge.net/docs/current/security-configuration.html
        '''
        if "community_string" in instance:
            # SNMP v1 - SNMP v2

            # See http://pysnmp.sourceforge.net/docs/current/security-configuration.html
            if int(instance.get("snmp_version", 2)) == 1:
                return cmdgen.CommunityData(instance['community_string'],
                    mpModel=0)
            return cmdgen.CommunityData(instance['community_string'], mpModel=1)

        elif "user" in instance:
            # SNMP v3
            user = instance["user"]
            auth_key = None
            priv_key = None
            auth_protocol = None
            priv_protocol = None
            if "authKey" in instance:
                auth_key = instance["authKey"]
                auth_protocol = cmdgen.usmHMACMD5AuthProtocol
            if "privKey" in instance:
                priv_key = instance["privKey"]
                auth_protocol = cmdgen.usmHMACMD5AuthProtocol
                priv_protocol = cmdgen.usmDESPrivProtocol
            if "authProtocol" in instance:
                auth_protocol = getattr(cmdgen, instance["authProtocol"])
            if "privProtocol" in instance:
                priv_protocol = getattr(cmdgen, instance["privProtocol"])
            return cmdgen.UsmUserData(user,
                                      auth_key,
                                      priv_key,
                                      auth_protocol,
                                      priv_protocol)
        else:
            raise Exception("An authentication method needs to be provided")

    @classmethod
    def get_transport_target(cls, instance, timeout, retries):
        '''
        Generate a Transport target object based on the instance's configuration
        '''
        if "ip_address" not in instance:
            raise Exception("An IP address needs to be specified")
        ip_address = instance["ip_address"]
        port = int(instance.get("port", 161)) # Default SNMP port
        return cmdgen.UdpTransportTarget((ip_address, port), timeout=timeout, retries=retries)

    def raise_on_error_indication(self, error_indication, instance):
        if error_indication:
            message = "{0} for instance {1}".format(error_indication,
                                                    instance["ip_address"])
            instance["service_check_error"] = message
            raise Exception(message)

    def check_table(self, instance, oids, lookup_names, timeout, retries):
        '''
        Perform a snmpwalk on the domain specified by the oids, on the device
        configured in instance.
        lookup_names is a boolean to specify whether or not to use the mibs to
        resolve the name and values.

        Returns a dictionary:
        dict[oid/metric_name][row index] = value
        In case of scalar objects, the row index is just 0
        '''
        # UPDATE: We used to perform only a snmpgetnext command to fetch metric values.
        # It returns the wrong value when the OID passeed is referring to a specific leaf.
        # For example:
        # snmpgetnext -v2c -c public localhost:11111 1.36.1.2.1.25.4.2.1.7.222
        # iso.3.6.1.2.1.25.4.2.1.7.224 = INTEGER: 2
        # SOLUTION: perform a snmget command and fallback with snmpgetnext if not found
        transport_target = self.get_transport_target(instance, timeout, retries)
        auth_data = self.get_auth_data(instance)

        first_oid = 0
        all_binds = []
        results = defaultdict(dict)

        while first_oid < len(oids):

            # Start with snmpget command
            error_indication, error_status, error_index, var_binds = self.snmpget(
                auth_data,
                transport_target,
                *(oids[first_oid:first_oid + OID_BATCH_SIZE]),
                lookupValues=lookup_names,
                lookupNames=lookup_names)

            # Raise on error_indication
            self.raise_on_error_indication(error_indication, instance)

            # Continue on error_status
            if error_status:
                message = "{0} for instance {1}".format(error_status.prettyPrint(),
                                                        instance["ip_address"])
                instance["service_check_error"] = message
                self.log.warning(message)
                continue

            missing_results = []
            complete_results = []

            for var in var_binds:
                result_oid, value = var
                if reply_invalid(value):
                    oid_tuple = result_oid.asTuple()
                    oid = ".".join([str(i) for i in oid_tuple])
                    missing_results.append(oid)
                else:
                    complete_results.append(var)

            if missing_results:
                # If we didn't catch the metric using snmpget, try snmpnext
                error_indication, error_status, error_index, var_binds_table = self.snmpgetnext(
                    auth_data,
                    transport_target,
                    *missing_results,
                    lookupValues=lookup_names,
                    lookupNames=lookup_names)

                # Raise on error_indication
                self.raise_on_error_indication(error_indication, instance)

                # Continue on error_status
                if error_status:
                    message = "{0} for instance {1}".format(error_status.prettyPrint(),
                                                            instance["ip_address"])
                    instance["service_check_error"] = message
                    self.log.warning(message)
                    continue
                for table_row in var_binds_table:
                    complete_results.extend(table_row)

            all_binds.extend(complete_results)
            first_oid = first_oid + OID_BATCH_SIZE
<commit_msg>[snmp] do not infinite loop on `error_status`<commit_after># std
from collections import defaultdict
from functools import wraps

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

# project
from checks import AgentCheck
from config import _is_affirmative



# Additional types that are not part of the SNMP protocol. cf RFC 2856
(CounterBasedGauge64, ZeroBasedCounter64) = builder.MibBuilder().importSymbols(
    "HCNUM-TC",
    "CounterBasedGauge64",
    "ZeroBasedCounter64")

# Metric type that we support
SNMP_COUNTERS = frozenset([
    snmp_type.Counter32.__name__,
    snmp_type.Counter64.__name__,
    ZeroBasedCounter64.__name__])

SNMP_GAUGES = frozenset([
    snmp_type.Gauge32.__name__,
    snmp_type.Unsigned32.__name__,
    CounterBasedGauge64.__name__,
    snmp_type.Integer.__name__,
    snmp_type.Integer32.__name__])

OID_BATCH_SIZE = 10


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


class SnmpCheck(AgentCheck):

    cmd_generator = None
    # pysnmp default values
    DEFAULT_RETRIES = 5
    DEFAULT_TIMEOUT = 1

    def __init__(self, name, init_config, agentConfig, instances=None):
        AgentCheck.__init__(self, name, init_config, agentConfig, instances)

        # Load Custom MIB directory
        mibs_path = None
        ignore_nonincreasing_oid = False
        if init_config is not None:
            mibs_path = init_config.get("mibs_folder")
            ignore_nonincreasing_oid = _is_affirmative(
                init_config.get("ignore_nonincreasing_oid", False))

        # Create SNMP command generator and aliases
        self.create_command_generator(mibs_path, ignore_nonincreasing_oid)

    def snmp_logger(self, func):
        """
        Decorator to log, with DEBUG level, SNMP commands
        """
        @wraps(func)
        def wrapper(*args, **kwargs):
            self.log.debug("Running SNMP command {0} on OIDS {1}"
                           .format(func.__name__, args[2:]))
            result = func(*args, **kwargs)
            self.log.debug("Returned vars: {0}".format(result[-1]))
            return result
        return wrapper

    def create_command_generator(self, mibs_path, ignore_nonincreasing_oid):
        '''
        Create a command generator to perform all the snmp query.
        If mibs_path is not None, load the mibs present in the custom mibs
        folder. (Need to be in pysnmp format)
        '''
        self.cmd_generator = cmdgen.CommandGenerator()
        self.cmd_generator.ignoreNonIncreasingOid = ignore_nonincreasing_oid
        if mibs_path is not None:
            mib_builder = self.cmd_generator.snmpEngine.msgAndPduDsp.\
                mibInstrumController.mibBuilder
            mib_sources = mib_builder.getMibSources() + \
                (builder.DirMibSource(mibs_path), )
            mib_builder.setMibSources(*mib_sources)

        # Set aliases for snmpget and snmpgetnext with logging
        self.snmpget = self.snmp_logger(self.cmd_generator.getCmd)
        self.snmpgetnext = self.snmp_logger(self.cmd_generator.nextCmd)

    @classmethod
    def get_auth_data(cls, instance):
        '''
        Generate a Security Parameters object based on the instance's
        configuration.
        See http://pysnmp.sourceforge.net/docs/current/security-configuration.html
        '''
        if "community_string" in instance:
            # SNMP v1 - SNMP v2

            # See http://pysnmp.sourceforge.net/docs/current/security-configuration.html
            if int(instance.get("snmp_version", 2)) == 1:
                return cmdgen.CommunityData(instance['community_string'],
                    mpModel=0)
            return cmdgen.CommunityData(instance['community_string'], mpModel=1)

        elif "user" in instance:
            # SNMP v3
            user = instance["user"]
            auth_key = None
            priv_key = None
            auth_protocol = None
            priv_protocol = None
            if "authKey" in instance:
                auth_key = instance["authKey"]
                auth_protocol = cmdgen.usmHMACMD5AuthProtocol
            if "privKey" in instance:
                priv_key = instance["privKey"]
                auth_protocol = cmdgen.usmHMACMD5AuthProtocol
                priv_protocol = cmdgen.usmDESPrivProtocol
            if "authProtocol" in instance:
                auth_protocol = getattr(cmdgen, instance["authProtocol"])
            if "privProtocol" in instance:
                priv_protocol = getattr(cmdgen, instance["privProtocol"])
            return cmdgen.UsmUserData(user,
                                      auth_key,
                                      priv_key,
                                      auth_protocol,
                                      priv_protocol)
        else:
            raise Exception("An authentication method needs to be provided")

    @classmethod
    def get_transport_target(cls, instance, timeout, retries):
        '''
        Generate a Transport target object based on the instance's configuration
        '''
        if "ip_address" not in instance:
            raise Exception("An IP address needs to be specified")
        ip_address = instance["ip_address"]
        port = int(instance.get("port", 161)) # Default SNMP port
        return cmdgen.UdpTransportTarget((ip_address, port), timeout=timeout, retries=retries)

    def raise_on_error_indication(self, error_indication, instance):
        if error_indication:
            message = "{0} for instance {1}".format(error_indication,
                                                    instance["ip_address"])
            instance["service_check_error"] = message
            raise Exception(message)

    def check_table(self, instance, oids, lookup_names, timeout, retries):
        '''
        Perform a snmpwalk on the domain specified by the oids, on the device
        configured in instance.
        lookup_names is a boolean to specify whether or not to use the mibs to
        resolve the name and values.

        Returns a dictionary:
        dict[oid/metric_name][row index] = value
        In case of scalar objects, the row index is just 0
        '''
        # UPDATE: We used to perform only a snmpgetnext command to fetch metric values.
        # It returns the wrong value when the OID passeed is referring to a specific leaf.
        # For example:
        # snmpgetnext -v2c -c public localhost:11111 1.36.1.2.1.25.4.2.1.7.222
        # iso.3.6.1.2.1.25.4.2.1.7.224 = INTEGER: 2
        # SOLUTION: perform a snmget command and fallback with snmpgetnext if not found
        transport_target = self.get_transport_target(instance, timeout, retries)
        auth_data = self.get_auth_data(instance)

        first_oid = 0
        all_binds = []
        results = defaultdict(dict)

        while first_oid < len(oids):

            # Start with snmpget command
            error_indication, error_status, error_index, var_binds = self.snmpget(
                auth_data,
                transport_target,
                *(oids[first_oid:first_oid + OID_BATCH_SIZE]),
                lookupValues=lookup_names,
                lookupNames=lookup_names)

            first_oid = first_oid + OID_BATCH_SIZE

            # Raise on error_indication
            self.raise_on_error_indication(error_indication, instance)

            # Continue on error_status
            if error_status:
                message = "{0} for instance {1}".format(error_status.prettyPrint(),
                                                        instance["ip_address"])
                instance["service_check_error"] = message
                self.log.warning(message)
                continue

            missing_results = []
            complete_results = []

            for var in var_binds:
                result_oid, value = var
                if reply_invalid(value):
                    oid_tuple = result_oid.asTuple()
                    oid = ".".join([str(i) for i in oid_tuple])
                    missing_results.append(oid)
                else:
                    complete_results.append(var)

            if missing_results:
                # If we didn't catch the metric using snmpget, try snmpnext
                error_indication, error_status, error_index, var_binds_table = self.snmpgetnext(
                    auth_data,
                    transport_target,
                    *missing_results,
                    lookupValues=lookup_names,
                    lookupNames=lookup_names)

                # Raise on error_indication
                self.raise_on_error_indication(error_indication, instance)

                # Continue on error_status
                if error_status:
                    message = "{0} for instance {1}".format(error_status.prettyPrint(),
                                                            instance["ip_address"])
                    instance["service_check_error"] = message
                    self.log.warning(message)
                    continue
                for table_row in var_binds_table:
                    complete_results.extend(table_row)

            all_binds.extend(complete_results)

{'fraction_non_alphanumeric': 0.07438326109059774, 'fraction_numerical': 0.010702627495050035, 'mean_word_length': 1.9996789727126805, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 8, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6369016', 'n_tokens_mistral': 5283, 'n_tokens_neox': 4954, 'n_words': 1238}