#!/usr/bin/env python

'''
This plugin checks the number of good ntp servers associated with chronyd

Default parameters are 
WARN if less than 2 associations are active
CRITICAL if 0 associations
'''

import argparse
import logging
import subprocess
import nagiosplugin

SOURCE_STATES = { '*': 'selected', '+': 'combined', '-': 'excluded', '?': 'lost contact', 'x': 'falseticker', '~': 'too variable' }
SOURCE_MODES = { '^': 'server', '=': 'peer', '#': 'local' }
SOURCE_STATE_OK = [ '*', '+', '-' ]

_log = logging.getLogger('nagiosplugin')

class Chronyd(nagiosplugin.Resource):
    '''
    This probe counts the number of good ntp sources
    currently associated with the chrony daemon, by parsing the output of 'chronyc sources'

    * A good source is
    * either a server, peer, or connected clock
    * connectivity has not been lost
    * is not a falseticker
    '''
    def probe(self):

        sources, all_sources = 0, 0
        try:
            output = subprocess.check_output(['chronyc', '-n', 'sources'])
        except subprocess.CalledProcessError:
            _log.warning('chronyc sources returned nonzero!')

        for line in output.splitlines():
            if '===' in line: continue
            if line[0] not in SOURCE_MODES: continue
            all_sources += 1
            ms, ref, stratum, poll, r, lastrx = line.split()[:6]
            lastsample = ' '.join(line.split()[6:])
            m, s = list(ms)
            mode = SOURCE_MODES.get(m)
            state = SOURCE_STATES.get(s)
            reach = bin(int(r, base=8))

            _log.debug("{} {} ok: {} ({}), stratum {}, history {}, lastrx {}".format(mode, ref, s in SOURCE_STATE_OK, state, stratum, reach, lastrx))
            if s in SOURCE_STATE_OK:
                sources += 1

        _log.debug("{} out of {} sources ok".format(sources, all_sources))

        try:
            output = subprocess.check_output(['chronyc', '-n', 'tracking'])
        except subprocess.CalledProcessError:
            _log.warning('chronyc tracking returned nonzero!')

        for line in output.splitlines():
            if line.split()[0] == 'RMS':
                offset = float(line.split()[3])
                _log.debug("offset found: %f" % offset)

        return (nagiosplugin.Metric('sources',
                                    sources, min=0, context='sources'),
                nagiosplugin.Metric('offset',
                                    offset, min=0, context='offset'))


@nagiosplugin.guarded
def main():
    argp = argparse.ArgumentParser(description=__doc__)
    argp.add_argument('-v', '--verbose', action='count', default=0)
    argp.add_argument('-w', '--warning', metavar='RANGE', default='2:',
                      help='return warning if outside RANGE')
    argp.add_argument('-c', '--critical', metavar='RANGE', default='1:',
                      help='return critical if outside RANGE')
    args = argp.parse_args()

    check = nagiosplugin.Check(Chronyd(),
                               nagiosplugin.ScalarContext('sources',
                                                          args.warning, args.critical,
                                                          fmt_metric="{value} active sources"),
                               nagiosplugin.ScalarContext('offset'))

    check.main(args.verbose)


if __name__ == '__main__':
    main()
