#!/usr/bin/python2.7
"""
Icinga/Nagios plugin to check status of a Livestatus socket
"""
import argparse
import json
import logging
import os
import socket
import sys
import traceback

DEFAULT_HOST = '127.0.0.1'
DEFAULT_PORT = 6557
DEFAULT_TIMEOUT = 2

# MAIN
parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument(
    '-d', '--debug',
    action='store_true',
    help='Show debug messages'
)

parser.add_argument(
    '-v', '--verbose',
    action='store_true',
    help='Show verbose messages'
)

parser.add_argument(
    '-l',
    '--log-file',
    action="store",
    type=str,
    dest="logfile",
    help="Log file path"
)

parser.add_argument(
    '-C', '--console',
    action='store_true',
    help='Also log on console (even if log file is provided)'
)

parser.add_argument(
    '-U', '--unix-socket-path',
    action='store',
    type=str,
    dest='unix_socket_path',
    help='Livestatus UNIX socket path',
)

parser.add_argument(
    '-H', '--host',
    action='store',
    type=str,
    dest='host',
    help='Livestatus host (default: %s)' % DEFAULT_HOST,
    default=DEFAULT_HOST
)

parser.add_argument(
    '-p', '--port',
    action='store',
    type=int,
    dest='port',
    help='Livestatus port (default: %d)' % DEFAULT_PORT,
    default=DEFAULT_PORT
)

parser.add_argument(
    '-t', '--timeout',
    action='store',
    type=float,
    dest='timeout',
    help='Livestatus timeout in second (default: %s)' % DEFAULT_TIMEOUT,
    default=DEFAULT_TIMEOUT
)
options = parser.parse_args()

# Initialize log
log = logging.getLogger()
logformat = logging.Formatter(
    "%(asctime)s - " + os.path.basename(sys.argv[0]) +
    " - %(levelname)s - %(message)s")

if options.debug:
    log.setLevel(logging.DEBUG)
elif options.verbose:
    log.setLevel(logging.INFO)
else:
    log.setLevel(logging.WARNING)

if options.logfile:
    logfile = logging.FileHandler(options.logfile)
    logfile.setFormatter(logformat)
    log.addHandler(logfile)

if not options.logfile or options.console:
    logconsole = logging.StreamHandler()
    logconsole.setFormatter(logformat)
    log.addHandler(logconsole)

peer = "UNKNOWN"
try:
    if options.unix_socket_path:
        peer = options.unix_socket_path
        livestatus = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        livestatus.connect(options.unix_socket_path)
    elif options.host and options.port:
        peer = "%s:%d" % (options.host, options.port)
        livestatus = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        livestatus.connect((options.host, options.port))
    else:
        parser.error(
            'You must provide Livestatus UNIX socket path or host & port')
        sys.exit(3)
    livestatus.settimeout(options.timeout)

    livestatus.send("GET status\nOutputFormat: json\nColumnHeaders: on\n\n")
    livestatus.shutdown(socket.SHUT_WR)
    rawdata = livestatus.makefile().read()
    status = None
    if rawdata:
        data = json.loads(rawdata)
        assert isinstance(data, list) and len(data) == 2, \
            "Invalid Livestatus return:\n%s" % data
        status = dict(
            (data[0][idx], value)
            for idx, value in enumerate(data[1])
        )
except Exception:  # pylint: disable=broad-except
    print(
        'UNKNOWN - Fail to retreive status from Livestatus '
        'from %s\n\nException:\n%s' %
        (peer, traceback.format_exc())
    )
    sys.exit(2)

if isinstance(status, dict) and 'num_hosts' in status:
    print('OK - %s host(s) found\n' % status['num_hosts'])
    for key in sorted(status.keys()):
        print('%s: %s' % (key, status[key]))
    sys.exit(0)

print('UNKNOWN - Invalid Livestatus return\n\n%s' % data)
sys.exit(3)