#!/usr/bin/env python

import json
import subprocess
import sys
import argparse
import os
import time
import collections
import re
import glob

def merge(a, b, path=None):
  "merges b into a"
  if path is None: path = []
  for key in b:
    if key in a:
      if isinstance(a[key], dict) and isinstance(b[key], dict):
        merge(a[key], b[key], path + [str(key)])
      elif a[key] == b[key]:
        pass # same leaf value
      else:
        raise Exception('Conflict at %s' % '.'.join(path + [str(key)]))
    else:
      a[key] = b[key]
  return a

def flatten(d, parent_key='', sep='_'):
  items = []
  for k,v in d.items():
    new_key = parent_key + sep + k if parent_key else k
    if isinstance(v, collections.MutableMapping):
      items.extend(flatten(v, new_key).items())
    else:
      items.append((new_key, v))
  return dict(items)

class SensuSpec:

  def __init__(self):
    self.status_count = {
      'ok': 0,
      'warning': 0,
      'critical': 0,
      'unknown': 0
    }

    self.options = {}

    path = os.getenv('PATH').split(':')
    path.append('/etc/sensu/plugins')
    path.append('/usr/lib/nagios/plugins')
    path.append('/usr/lib64/nagios/plugins')

    if os.path.isdir('/opt/sensu/embedded/bin'):
      path.insert(0, '/opt/sensu/embedded/bin')

    os.putenv('PATH', ":".join(path))

  def args(self,args):
    self.options = vars(args)

  def run(self):
    self.read_files()
    self.run_tests()
    self.print_summary()

  def read_files(self):
    fh = open(self.options['path'] + '/client.json')
    self.client_data = flatten(json.load(fh)['client'])
    fh.close()

    fh = open(self.options['path'] + '/checks.json')
    self.check_data = json.load(fh)
    fh.close()

    for file in glob.glob(self.options['path'] + '/checks/*.json'):
      fh = open(file)
      merge(self.check_data, json.load(fh))
      fh.close()

  def run_tests(self):
    for check in self.check_data['checks']:
      self.run_test(check)

  def run_test(self, check):
    data = self.check_data['checks'][check]
    command = data['command']

    sys.stdout.write(check+" ")

    for arg in re.findall(':::(?:.+?):::', command):
      clean = re.sub(':::', '', arg)
      name = re.sub('\.', '_', clean)
      value = self.client_data[name]
      command = re.sub(arg, value, command)

    for i in range(self.options['retry'] + 1):
      process = subprocess.Popen(command, shell=True, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
      exit_code = process.wait()
      (stdout,stderr) = process.communicate()

      if exit_code == 0:
        break
      elif i > 0:
        time.sleep(self.options['sleep'])

    if (exit_code == 0):
      print "OK"
      self.status_count['ok'] += 1
    elif (exit_code == 1):
      print "WARNING"
      print "  Command run: " + command
      print "  Output: " + stdout
      self.status_count['warning'] += 1
    elif (exit_code == 2):
      print "CRITICAL"
      print "  Command run: " + command
      print "  Output: " + stdout
      self.status_count['critical'] += 1
    else:
      print "UNKNOWN"
      print "  Command run: " + command
      print "  Output: " + stdout
      self.status_count['unknown'] += 1

    return exit_code

  def print_summary(self):
    total = sum(self.status_count.values())

    if (total == 0):
      print "No tests run"
      sys.exit(0)
    elif (self.status_count['ok'] == 0):
      print "All tests failed"
      sys.exit(1)
    elif (total != self.status_count['ok']):
      print "Some tests failed"
      sys.exit(1)
    else:
      print "All tests passed"
      sys.exit(0)

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("-p", "--path", help="Path containing sensu client config", default="/etc/sensu/conf.d/")
  parser.add_argument("-r", "--retry", help="Number of times to retry a test", type=int, default=0)
  parser.add_argument("-s", "--sleep", help="Seconds to sleep between test retries", type=float, default=5.0)

  sensu_spec = SensuSpec()
  sensu_spec.args(parser.parse_args())
  sensu_spec.run()
