check_linux_dns

#!/usr/bin/env python
# SnapServ/monitoring-plugins - Simple and reliable Nagios / Icinga plugins
# Copyright (C) 2016 - 2017 SnapServ Mathis
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import logging
import math
import time
from typing import Dict, Optional  # pylint: disable=unused-import

import nagiosplugin
import dns.resolver
from shared.common import ExceptionContext, NagiosPlugin

LOGGER = logging.getLogger('nagiosplugin')


class DnsTimeout(dns.exception.Timeout):
    supp_kwargs = {'timeout', 'threshold'}  # type: Dict[str]
    fmt = "The DNS operation timed out after {timeout} seconds (>{threshold} seconds)"  # type: str


class Dns(nagiosplugin.Resource):
    def __init__(self, server: Optional[str], query_name: str, rd_type: dns.rdatatype,
                 timeout: float) -> None:
        self.server = server  # type: Optional[str]
        self.query_name = query_name  # type: str
        self.rd_type = rd_type  # type: dns.rdatatype
        self.timeout = timeout  # type: float

    def probe(self) -> nagiosplugin.Metric:
        try:
            # Instantiate a new resolver and execute the query while measuring the time.
            resolver = self.get_resolver()
            LOGGER.info('Querying system nameservers for "%s" with type "%s"',
                        self.query_name, self.rd_type)
            start_time = time.time()
            answer = resolver.query(self.query_name, self.rd_type)
            response_time = time.time() - start_time

            # Get the query result and print it as a debug message
            records = [str(record) for record in answer.rrset]
            LOGGER.debug('Query result: ' + ', '.join(records))

            # As dnspython only supports integer timeouts, we have to manually check the query
            # time too and throw an exception if it should have taken too long.
            if response_time > self.timeout:
                LOGGER.debug('Manually raising timeout exception, as dnspython only supports '
                             + 'whole seconds.')
                raise DnsTimeout(timeout=response_time, threshold=self.timeout)

            # Return the metric to the plugin
            return nagiosplugin.Metric('response_time', response_time, min=0, uom='s')

        # Re-raise timeout exceptions from dnspython with our own class (supports displaying
        # the threshold value) and pass it to our ExceptionContext
        except dns.exception.Timeout as exception:
            return nagiosplugin.Metric('exception', DnsTimeout(
                timeout=exception.kwargs.get('timeout'),
                threshold=self.timeout
            ))

        # Pass exceptions from dnspython to our ExceptionContext for further processing
        except dns.exception.DNSException as exception:
            LOGGER.info('Query has failed with an exception: ' + str(exception))
            return nagiosplugin.Metric('exception', exception)

    def get_resolver(self) -> dns.resolver.Resolver:
        # Use /etc/resolv.conf if no server was specified, otherwise directly use the given server.
        if not self.server:
            resolver = dns.resolver.Resolver(filename='/etc/resolv.conf', configure=True)
        else:
            resolver = dns.resolver.Resolver(configure=False)
            resolver.nameservers = [self.server]

        # dnspython only supports integer timeouts, so we always round the number up - we do
        # our own, manual checks when executing the query to properly handle float timeouts.
        resolver.timeout = resolver.lifetime = math.ceil(self.timeout)
        return resolver


class DnsSummary(nagiosplugin.Summary):
    def __init__(self, server: Optional[str]) -> None:
        self.prefix = (server + ': ') if server else ''  # type: str

    def ok(self, results: nagiosplugin.Results) -> str:
        message = 'Response time is %ss' % round(results['response_time'].metric.value, 3)
        return self.prefix + message

    def problem(self, results: nagiosplugin.Results) -> str:
        message = super().problem(results)
        return self.prefix + message


class DnsPlugin(NagiosPlugin):
    def declare_arguments(self) -> None:
        self.argument_parser.add_argument(
            '-s', '--server', default=None,
            help='Use a specific DNS server for executing the check. Defaults to system DNS '
                 + 'resolvers if not specified.'
        )
        self.argument_parser.add_argument(
            '-q', '--query-name', required=True,
            help='The query name which should be queried by using domain name resolution.'
        )
        self.argument_parser.add_argument(
            '-T', '--rd-type', required=True,
            help='The record type which should be queried, e.g.: SOA, NS, A, AAAA'
        )
        self.argument_parser.add_argument(
            '-t', '--timeout', required=True, type=float, metavar='TIMEOUT',
            help='Abort the query and return CRITICAL when query takes longer than TIMEOUT.'
        )
        self.argument_parser.add_argument(
            '-w', '--warning', metavar='RANGE', default='',
            help='Return WARNING if response time is outside RANGE.'
        )
        self.argument_parser.add_argument(
            '-c', '--critical', metavar='RANGE', default='',
            help='Return CRITICAL if response time is outside RANGE.'
        )
        self.exclude_from_kwargs += ('warning', 'critical')

    def instantiate_check(self) -> nagiosplugin.Check:
        return nagiosplugin.Check(
            Dns(**self.keyword_arguments),

            nagiosplugin.ScalarContext(
                name='response_time',
                warning=self.arguments.get('warning'),
                critical=self.arguments.get('critical')
            ),

            ExceptionContext(),

            DnsSummary(self.arguments.get('server'))
        )


if __name__ == '__main__':
    DnsPlugin().execute()