#!/usr/bin/env python
import fnmatch
import os
import re
import subprocess
import sys

from argparse import ArgumentParser

spincheck_tmpfile='/var/lib/nagios/daemontools_spincheck'

try:
    import json
except ImportError:
    import simplejson as json


def flatten_list(list):
    if not list:
        return []
    return reduce(lambda x, y: x + y, list)

try:
    try:
        old = json.load(open(spincheck_tmpfile, 'r'))
    except (IOError, ValueError):
        old = {}
    new = {}
    ignored = []
    ignoring = []
    regexes = []

    parser = ArgumentParser()
    parser.add_argument(
        "-i", "--ignore", dest="ignored_services", action='append',
        metavar="SERVICES", nargs='+', default=[],
        help="Services to ignore; comma separated; supports globs (*)"
    )
    parser.add_argument(
        "-r", "--regex", dest="ignored_regexes", action='append',
        metavar="REGEX", nargs='+', default=[],
        help="Regexes of services to ignore"
    )
    parser.add_argument(
        "-v", "--verbose", dest="verbose", action="store_true",
        help="Print all services found, including those without errors",
    )

    parser.epilog = "'-r' and '-i' can be specified multiple times."
    args = parser.parse_args()

    # Split ignored_services
    if args.ignored_services:
        args.ignored_services = flatten_list(args.ignored_services)
        ignoring = map(lambda x: x.split(','), args.ignored_services)
        ignoring = flatten_list(ignoring)

    # Precompile regexes.
    args.ignored_regexes = flatten_list(args.ignored_regexes)
    for regex in args.ignored_regexes:
        regexes.append(re.compile(regex))

    services = subprocess.Popen(
        'sudo /usr/bin/svstat /etc/service/*', shell=True, stdout=subprocess.PIPE
    ).communicate()[0].splitlines()

    running_service_names = []
    down_service_names = []

    for line in services:
        stat = line.split(':')[0]
        service_name = os.path.split(stat)[-1]

        f_ignored = False
        if ignoring or regexes:
            for glob in ignoring:
                if fnmatch.fnmatch(service_name, glob):  # globs supported
                    f_ignored = True
                    ignored.append(service_name)
                    continue

            for regex in regexes:
                if regex.match(service_name):
                    f_ignored = True
                    print "Ignoring %s" % service_name
                    ignored.append(service_name)
                    continue

        if f_ignored:
            continue

        count, oldpid = old.get(service_name, [0, 0])
        match = re.match('.*\(pid ([0-9]*)\)', line)

        if not match:
            down_service_names.append(service_name)
            continue

        running_service_names.append(service_name)

        newpid = int(match.group(1))

        if newpid == oldpid:
            new[service_name] = [0, newpid]
        else:
            new[service_name] = [count + 1, newpid]

    try:
        json.dump(new, open(spincheck_tmpfile, 'w'))
    except (IOError, ValueError) as e:
        print 'CRIT: Error writing spincheck file.\n{}'.format(e)
        sys.exit(2)

    bad = []
    for k, v in new.iteritems():
        if v[0] > 3:
            bad.append(k)

    if len(bad) > 0:
        print ', '.join(bad)
        sys.exit(2)

    running_service_names = ', '.join(running_service_names)
    down_service_names = ', '.join(down_service_names)

    extra_info = ''
    if args.verbose:
        extra_info = '\n\nRunning: %s\n\nDown: %s' % (running_service_names, down_service_names)

    # set(ignored)
    if ignored:
        print 'OK: all workers running. Ignored: %s %s' % (ignored, extra_info)
    else:
        print "OK: all workers running. %s" % (extra_info,)

except Exception, e:
    print repr(e)
    sys.exit(3)
