check_xmpp

#!/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 socket
import ssl
import os
from enum import Enum
from select import select
from datetime import datetime
from xml.sax.handler import ContentHandler, feature_namespaces
from _ssl import CERT_NONE, CERT_REQUIRED

import nagiosplugin
from defusedxml.sax import make_parser
from shared.common import DaysValidContext, ExceptionContext, CommaSeparatedSummary, NagiosPlugin

LOGGER = logging.getLogger('nagiosplugin')

NS_IETF_XMPP_SASL = 'urn:ietf:params:xml:ns:xmpp-sasl'
NS_IETF_XMPP_TLS = 'urn:ietf:params:xml:ns:xmpp-tls'
NS_IETF_XMPP_STREAMS = 'urn:ietf:params:xml:ns:xmpp-streams'
NS_JABBER_CAPS = 'http://jabber.org/protocol/caps'
NS_ETHERX_STREAMS = 'http://etherx.jabber.org/streams'


class XmppState(Enum):
    NEW = 1
    STREAM_START = 2
    RECEIVED_FEATURES = 3
    FINISHED = 4
    ERROR = 5
    PROCEED_STARTTLS = 6


class XmppException(Exception):
    def __init__(self, message):
        super(XmppException, self).__init__(message)


class XmppStreamError(object):
    def __init__(self):
        self.text = None
        self.condition = None
        self.data = {}

    def __str__(self):
        if self.text:
            return '%s: %s' % (self.condition, self.text)
        else:
            return self.condition


class XmppResponseHandler(ContentHandler):
    def __init__(self, expect_starttls):
        self.state = XmppState.NEW
        self.starttls = False
        self.tls_required = False
        self.expect_starttls = expect_starttls
        self.capabilities = {}
        self.sasl_mechanisms = []

        self.incoming_elements = []
        self.seen_elements = set()
        self.stream_information = None
        self.error_instance = None

        super(XmppResponseHandler, self).__init__()

    def startElementNS(self, name, qname, attrs):
        self.incoming_elements.append(name)
        self.seen_elements.add(name)

        if name == (NS_ETHERX_STREAMS, 'stream'):
            self.state = XmppState.STREAM_START
            self.stream_information = dict(
                [(qname, attrs.getValueByQName(qname)) for qname in attrs.getQNames()])
        elif name == (NS_IETF_XMPP_TLS, 'starttls'):
            self.starttls = True
            LOGGER.info('XMPP server offers StartTLS.')
        elif self.incoming_elements[-2] == (NS_IETF_XMPP_TLS, 'starttls') and name == (
                NS_IETF_XMPP_TLS, 'required'):
            self.tls_required = True
            LOGGER.info('XMPP server requires the other side to use StartTLS.')
        elif name == (NS_JABBER_CAPS, 'c'):
            for qname in attrs.getQNames():
                value = attrs.getValueByQName(qname)
                self.capabilities[qname] = value
                LOGGER.info('XMPP server announced capability: %s => %s', qname, value)
        elif name == (NS_ETHERX_STREAMS, 'error'):
            self.state = XmppState.ERROR
            self.error_instance = XmppStreamError()
        elif self.state == XmppState.ERROR and name != (NS_IETF_XMPP_STREAMS, 'text'):
            if name[0] == NS_IETF_XMPP_STREAMS:
                self.error_instance.condition = name[1]
            else:
                self.error_instance.data[name] = {
                    'attrs': dict(
                        [(qname, attrs.getValueByQName(qname)) for qname in attrs.getQNames()])}

    def endElementNS(self, name, qname):
        if name == (NS_ETHERX_STREAMS, 'features'):
            self.state = XmppState.RECEIVED_FEATURES
        elif name == (NS_ETHERX_STREAMS, 'stream'):
            self.state = XmppState.FINISHED
        elif name == (NS_ETHERX_STREAMS, 'error'):
            raise XmppException('XMPP stream error: %s' % self.error_instance)
        elif name == (NS_IETF_XMPP_TLS, 'proceed'):
            self.state = XmppState.PROCEED_STARTTLS
        elif name == (NS_IETF_XMPP_TLS, 'failure'):
            raise XmppException('StartTLS initialization failed.')
        del self.incoming_elements[-1]

    def characters(self, content):
        element = self.incoming_elements[-1]
        if element == (NS_IETF_XMPP_SASL, 'mechanism'):
            self.sasl_mechanisms.append(content)
        elif self.state == XmppState.ERROR:
            if element == (NS_IETF_XMPP_STREAMS, 'text'):
                self.error_instance.text = content
            else:
                self.error_instance.data[element]['text'] = content
        else:
            LOGGER.debug('Ignored XMPP stream content in %s: %s', element, content)

    def is_valid_connection_start(self):
        if self.state != XmppState.RECEIVED_FEATURES:
            raise XmppException('Did not receive complete feature list from XMPP server.')
        elif self.expect_starttls and not self.starttls:
            raise XmppException('Expected StartTLS, which was not offered by XMPP server.')

        if self.stream_information.get('version', None) != '1.0':
            LOGGER.warning('Unknown XMPP stream version: %s', self.stream_information['version'])
            return False

        return True


class Xmpp(nagiosplugin.Resource):
    def __init__(self, server, port, host, use_s2s, use_ipv6, starttls, ca_roots,
                 check_certificates):
        self.server = server
        self.port = port if port else (5269 if use_s2s else 5222)
        self.host = host
        self.use_s2s = use_s2s
        self.use_ipv6 = use_ipv6
        self.starttls = starttls
        self.ca_roots = ca_roots
        self.check_certificates = check_certificates

        self.parser = None
        self.content_handler = None
        self.socket = None
        self.certificate_days_left = None
        self.state = None
        self.cause = None

        self.initialize_parser()
        self.initialize_content_handler()

    def initialize_parser(self):
        self.parser = make_parser()
        self.parser.setFeature(feature_namespaces, True)

    def initialize_content_handler(self):
        self.content_handler = XmppResponseHandler(expect_starttls=self.starttls)
        self.parser.setContentHandler(self.content_handler)

    def getaddrinfo(self):
        if self.use_ipv6 is None:
            address_family = 0
        elif self.use_ipv6 is True:
            address_family = socket.AF_INET6
        else:
            address_family = socket.AF_INET

        return socket.getaddrinfo(self.server, self.port, address_family, socket.SOCK_STREAM,
                                  socket.IPPROTO_TCP)

    @staticmethod
    def open_socket(addrinfo):
        sock = None
        for res in addrinfo:
            address_family, socktype, proto, _, socket_address = res
            try:
                sock = socket.socket(address_family, socktype, proto)
                sock.connect(socket_address)
                break
            except socket.error:
                sock.close()
                sock = None
                continue

        if sock is None:
            raise XmppException('Could not open socket to XMPP server.')
        return sock

    def receive_xmpp_stanza(self, timeout=0.1):
        chunks = []
        while True:
            read_ready, _, _ = select([self.socket], [], [], timeout)
            if self.socket in read_ready:
                data = self.socket.recv(4096)
                if data:
                    chunks.append(data)
                else:
                    break
            else:
                break

        return b''.join(chunks).decode('utf-8')

    def handle_xmpp_stanza(self, message, timeout=0.1, expected_state=None):
        self.socket.sendall((message.encode('utf-8')))
        xml_stanza = self.receive_xmpp_stanza(timeout)
        self.parser.feed(xml_stanza)

        if expected_state is not None and self.content_handler.state != expected_state:
            raise XmppException(
                'XMPP connection has unexpected state: %s' % self.content_handler.state)

    def start_stream(self):
        if self.use_s2s:
            self.handle_xmpp_stanza(('<?xml version="1.0"?><stream:stream to="{host}" '
                                     'xmlns="jabber:server" '
                                     'xmlns:stream="http://etherx.jabber.org/streams" '
                                     'version="1.0">').format(host=self.host),
                                    expected_state=XmppState.RECEIVED_FEATURES)
        else:
            self.handle_xmpp_stanza(('<?xml version="1.0"?><stream:stream to="{host}" '
                                     'xmlns="jabber:client" '
                                     'xmlns:stream="http://etherx.jabber.org/streams" '
                                     'version="1.0">').format(host=self.host),
                                    expected_state=XmppState.RECEIVED_FEATURES)

    def setup_ssl_context(self):
        context = ssl.create_default_context()
        if not self.check_certificates:
            context.check_hostname = False
            context.verify_mode = CERT_NONE
        else:
            context.verify_mode = CERT_REQUIRED
            if self.ca_roots:
                if os.path.isfile(self.ca_roots):
                    kwargs = {'cafile': self.ca_roots}
                elif os.path.isdir(self.ca_roots):
                    kwargs = {'capath': self.ca_roots}
                else:
                    raise XmppException(
                        'Custom CA certificates location does not exist: %s' % self.ca_roots)
                context.load_verify_locations(**kwargs)
            else:
                context.load_default_certs()

            stats = context.cert_store_stats()
            if stats['x509_ca'] == 0:
                LOGGER.warning(
                    'Tried to load CA certificates from default locations, but could not find any '
                    'CA certificates.')
                raise XmppException('Could not load CA certificates from default locations.')
            else:
                LOGGER.debug('Certificate store statistics: %s', stats)

        return context

    def initiate_tls(self):
        self.handle_xmpp_stanza('<starttls xmlns="{xmlns}"/>'.format(xmlns=NS_IETF_XMPP_TLS),
                                expected_state=XmppState.PROCEED_STARTTLS)
        ssl_context = self.setup_ssl_context()

        try:
            self.socket = ssl_context.wrap_socket(self.socket, server_hostname=self.host)
            LOGGER.info('TLS socket setup was successful.')
        except ssl.SSLError:
            raise Exception('SSL error')
        except ssl.CertificateError as err:
            raise Exception('Certificate error %s', err)

        self.starttls = False
        self.parser.reset()
        if self.check_certificates:
            certificate_info = self.socket.getpeercert()
            LOGGER.debug('Received certificate info from XMPP server: %s', certificate_info)
            LOGGER.info('Certificate is valid from %s until %s', certificate_info['notBefore'],
                        certificate_info['notAfter'])
            end_date = ssl.cert_time_to_seconds(certificate_info['notAfter'])
            remaining_time = datetime.fromtimestamp(end_date) - datetime.now()
            self.certificate_days_left = remaining_time.days
            LOGGER.info('Certificate lifetime: %i days left', self.certificate_days_left)

        self.initialize_parser()
        self.initialize_content_handler()
        self.start_stream()

        if not self.content_handler.is_valid_connection_start():
            raise Exception('Invalid XMPP connection start')

    def handle_xmpp(self):
        self.start_stream()
        if not self.content_handler.is_valid_connection_start():
            raise Exception("Invalid XMPP connection response")
        if self.starttls is True or self.content_handler.tls_required:
            self.initiate_tls()
        self.handle_xmpp_stanza('</stream:stream>')

    def probe(self):
        start = datetime.now()
        try:
            self.socket = self.open_socket(self.getaddrinfo())
            try:
                self.handle_xmpp()
            finally:
                self.socket.close()
            self.parser.close()
        except socket.gaierror as err:
            LOGGER.debug('Catched socket.gaierror exception: %s', err)
            yield nagiosplugin.Metric('exception', err)
            return
        except XmppException as err:
            LOGGER.debug('Catched XmppException exception: %s', err)
            yield nagiosplugin.Metric('exception', err)
            return
        end = datetime.now()

        yield nagiosplugin.Metric('time', (end - start).total_seconds(), uom='s', min=0)
        yield nagiosplugin.Metric('certificate_days_left', self.certificate_days_left, uom='d')


class XmppPlugin(NagiosPlugin):
    def declare_arguments(self):
        self.argument_parser.add_argument(
            '-s', '--server', required=True,
            help='Specifies the XMPP server address / domain which should be checked.'
        )
        self.argument_parser.add_argument(
            '-p', '--port', type=int,
            help='Optionally specifies the XMPP port, defaults to TCP 5222 (C2S) / TCP 5269 (S2S).'
        )
        self.argument_parser.add_argument(
            '-H', '--host', required=True,
            help='Specifies the XMPP host which should be used.'
        )
        self.argument_parser.add_argument(
            '--starttls', action='store_true',
            help='Whether to use StartTLS when connecting.'
        )
        self.argument_parser.add_argument(
            '--no-check-certificates', dest='check_certificates',
            default=True, action='store_false',
            help='Do not check whether the XMPP server certificate is valid.'
        )
        self.argument_parser.add_argument(
            '--ca-roots',
            help='Path to a file or directory where CA certificates can be found.'
        )
        self.argument_parser.add_argument(
            '-w', '--warning', metavar='SECONDS', default='',
            help='Return warning if connection setup takes longer than SECONDS'
        )
        self.argument_parser.add_argument(
            '-c', '--critical', metavar='SECONDS', default='',
            help='Return critical if connection setup takes longer than SECONDS'
        )
        self.argument_parser.add_argument(
            '--warning-days', dest='warning_days', type=int, default=0, metavar='DAYS',
            help='Return warning if the remaining certificate lifetime in days is '
                 + 'shorter than DAYS.'
        )
        self.argument_parser.add_argument(
            '--critical-days', dest='critical_days', type=int, default=0, metavar='DAYS',
            help='Return critical if the remaining certificate lifetime in days is '
                 + 'shorter than DAYS.'
        )

        use_s2s = self.argument_parser.add_mutually_exclusive_group()
        use_s2s.add_argument(
            '--s2s', dest='use_s2s', action='store_true',
            help='Server to Server (S2S)'
        )
        use_s2s.add_argument(
            '--c2s', dest='use_s2s', action='store_false',
            help='Client to Server (C2S)'
        )

        use_ipv6 = self.argument_parser.add_mutually_exclusive_group()
        use_ipv6.add_argument(
            '-6', dest='use_ipv6', action='store_true',
            help='Enforce usage of IPv6'
        )
        use_ipv6.add_argument(
            '-4', dest='use_ipv6', action='store_false',
            help='Enforce usage of IPv4'
        )

        self.exclude_from_kwargs += ('warning', 'critical', 'warning_days', 'critical_days')

    def instantiate_check(self):
        return nagiosplugin.Check(
            Xmpp(**self.keyword_arguments),

            nagiosplugin.ScalarContext(
                name='time',
                warning=self.arguments.get('warning'),
                critical=self.arguments.get('critical'),
                fmt_metric='Request took {value}{uom}'
            ),

            DaysValidContext(
                name='certificate_days_left',
                warning_days=self.arguments.get('warning_days'),
                critical_days=self.arguments.get('critical_days'),
                check_lifetime=self.arguments.get('check_certificates'),
                fmt_metric='Certificate valid for {value} days'
            ),

            ExceptionContext(),

            CommaSeparatedSummary()
        )


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