#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (C) 2019 tribe29 GmbH - License: GNU General Public License v2
# This file is part of Checkmk (https://checkmk.com). It is subject to the terms and
# conditions defined in the file COPYING, which is part of this source code package.

import getopt
import os
import sys
import time
from typing import Any, Dict, List, NoReturn, Optional

from livestatus import QueryTypes, SingleSiteConnection

# These variable will be substituted at 'make dist' time
__version__ = "2.1.0p23"


def dump_templates():
    check_interval = ""
    if opt_check_interval is not None:
        check_interval = "check_interval          %d" % opt_check_interval

    sys.stdout.write(
        """define host {
    name                    livedump-host
    use                     check_mk_default
    register                0
    active_checks_enabled   0
    passive_checks_enabled  1
    %s
}

define service {
    name                    livedump-service
    register                0
    active_checks_enabled   0
    passive_checks_enabled  1
    check_period            0x0
    %s
}

define command {
    command_name            check-livedump
    command_line            echo "WARN - You did an active check, but this check is passive" ; exit 1
}

define timeperiod {
    timeperiod_name         0x0
    alias                   Never ever
}

"""
        % (check_interval, check_interval)
    )


#   .-Livedump-------------------------------------------------------------.
#   |            _     _               _                                   |
#   |           | |   (_)_   _____  __| |_   _ _ __ ___  _ __              |
#   |           | |   | \ \ / / _ \/ _` | | | | '_ ` _ \| '_ \             |
#   |           | |___| |\ V /  __/ (_| | |_| | | | | | | |_) |            |
#   |           |_____|_| \_/ \___|\__,_|\__,_|_| |_| |_| .__/             |
#   |                                                   |_|                |
#   +----------------------------------------------------------------------+
#   | The actual livedump                                                  |
#   '----------------------------------------------------------------------'

opt_socket: Optional[str] = None


def connect(socket_path: Optional[str]) -> SingleSiteConnection:
    if socket_path is not None:
        return SingleSiteConnection(socket_path)
    if omd_root := os.getenv("OMD_ROOT"):
        return SingleSiteConnection(f"unix:{omd_root}/tmp/run/live")
    bail_out("specify Livestatus socket or set OMD_ROOT")


def run_query(query: QueryTypes) -> List[Dict[str, Any]]:
    if g_connection is None:
        bail_out("No connection")
    return g_connection.query_table_assoc(query)


def livedump_config():
    if opt_dump_templates:
        dump_templates()

    check_commands = set()

    def prepare_row(row):
        check_commands.add(row["check_command"])
        row["contactsstring"] = ",".join(row["contacts"])
        if "contact_groups" in row:
            row["contact_groups"] = ",".join(row["contact_groups"])

    # Dump host config
    query = (
        "GET hosts\n"
        "Columns: name alias address groups check_command icon_image "
        "max_check_attempts contacts contact_groups\n" + opt_host_headers + opt_host_only_headers
    )
    for row in run_query(query):
        prepare_row(row)
        row["groupstring"] = ",".join(row["groups"])
        sys.stdout.write("define host {\n" "  use                livedump-host\n")

        if opt_prefix:
            sys.stdout.write("  host_name          %s%s\n" % (opt_prefix, row["name"]))
        else:
            sys.stdout.write("  host_name          %s\n" % row["name"])

        sys.stdout.write(
            "  alias              %(alias)s\n"
            "  address            %(address)s\n"
            "  host_groups        %(groupstring)s\n"
            "  check_command      %(check_command)s\n"
            "  max_check_attempts %(max_check_attempts)d\n" % row
        )

        if opt_groups:
            sys.stdout.write("  contacts           %s\n" % row["contactsstring"])
        else:
            sys.stdout.write("  contact_groups     %s\n" % row["contact_groups"])

        if opt_host_icon:
            if row.get("icon_image"):
                sys.stdout.write("  icon_image         %s\n" % row["icon_image"])

        sys.stdout.write("}\n\n")

    # Dump service config
    query = (
        "GET services\n"
        "Columns: host_name description groups check_command "
        "max_check_attempts contacts\n" + opt_host_headers + opt_service_headers
    )
    for row in run_query(query):
        prepare_row(row)
        if row["groups"]:
            row["groupstring"] = "service_groups " + ",".join(row["groups"])
        else:
            row["groupstring"] = ""
        row["contactsstring"] = ",".join(row["contacts"])
        sys.stdout.write("define service {\n  use                livedump-service\n")

        if opt_prefix:
            sys.stdout.write("  host_name          %s%s\n" % (opt_prefix, row["host_name"]))
        else:
            sys.stdout.write("  host_name          %s\n" % row["host_name"])

        sys.stdout.write(
            "  description        %(description)s\n"
            "  %(groupstring)s\n"
            "  check_command      check-livedump\n"
            "  contacts           %(contactsstring)s\n"
            "  max_check_attempts %(max_check_attempts)d\n"
            "}\n\n" % row
        )


def livedump_state():
    now = time.time()
    # Dump hosts
    query = (
        "GET hosts\n"
        "Columns: name state plugin_output perf_data latency\n"
        + opt_host_headers
        + opt_host_only_headers
    )

    for row in run_query(query):
        row["now"] = now
        if opt_prefix:
            sys.stdout.write("host_name=%s%s" % (opt_prefix, row["name"]))
        else:
            sys.stdout.write("host_name=%s" % row["name"])

        sys.stdout.write(
            """
check_type=1
check_options=0
reschedule_check
latency=%(latency).2f
start_time=%(now).1f
finish_time=%(now).1f
return_code=%(state)d
output=%(plugin_output)s|%(perf_data)s

"""
            % row
        )

    query = (
        "GET services\n"
        "Columns: host_name description state plugin_output perf_data latency\n"
        + opt_host_headers
        + opt_service_headers
    )

    for row in run_query(query):
        row["now"] = now
        if opt_prefix:
            sys.stdout.write("host_name=%s%s" % (opt_prefix, row["host_name"]))
        else:
            sys.stdout.write("host_name=%s" % row["host_name"])

        sys.stdout.write(
            """
service_description=%(description)s
check_type=1
check_options=0
reschedule_check
latency=%(latency).2f
start_time=%(now).1f
finish_time=%(now).1f
return_code=%(state)d
output=%(plugin_output)s|%(perf_data)s

"""
            % row
        )


# .
#   .-Helpers--------------------------------------------------------------.
#   |                  _   _      _                                        |
#   |                 | | | | ___| |_ __   ___ _ __ ___                    |
#   |                 | |_| |/ _ \ | '_ \ / _ \ '__/ __|                   |
#   |                 |  _  |  __/ | |_) |  __/ |  \__ \                   |
#   |                 |_| |_|\___|_| .__/ \___|_|  |___/                   |
#   |                              |_|                                     |
#   +----------------------------------------------------------------------+
#   | Various helper functions                                             |
#   '----------------------------------------------------------------------'


def omd_site():
    return os.getenv("OMD_SITE")


def verbose(x):
    if opt_verbose:
        sys.stderr.write("%s\n" % x)


def bail_out(x: str) -> NoReturn:
    sys.stderr.write("%s\n" % x)
    sys.exit(1)


def usage():
    sys.stderr.write(
        """Usage: %s [OPTION]...

  -s, --socket S             Connect to Livestatus-socket at S
                             -s tcp:10.11.0.55:6557
                             -s unix:/var/run/nagios/rw/live

  -C, --config               Dump configuration (instead of state)

  -T, --dump-templates       Also dump host/service templates

  -M, --mark-mode            Puts the mode (state/configuration dump) in the first
                             line for use with livedump-ssh-recv and similar scripts

  -p, --prefix P             Add a prefix P to hostnames. Use this option to dump live
                             data from multiple sites with duplicated hostnames

  -O, --host-only-header H   Add header H to host queries only (usually Filter: ...)

  -H, --host-header H        Add header H to host queries (usually Filter: ...)
                             This header is also used in service queries

  -S, --service-header H     Add header H to service queries (usually Filter: ...)

  -i, --interval I           Assume this check interval for hosts/services. This is
                             used in nagios config definition "check_interval".
                             Will be used for staleness checks on the server. This
                             option is only used when dumping the config together
                             with templates (-C -T)

  -G, --include-groups       Use contact groups instead of contacts on dumping config

      --include-host-icon    Add host icon_image to config

  -V, --version              Show version and exit
  -h, --help                 Show this help
  -v, --verbose              Output debug information on stderr
      --debug                Do not catch Python exceptions
"""
        % os.path.split(sys.argv[0])[1]
    )


def print_version():
    sys.stdout.write("This is livedump version %s\n" % __version__)


# .
#   .-main-----------------------------------------------------------------.
#   |                                       _                              |
#   |                       _ __ ___   __ _(_)_ __                         |
#   |                      | '_ ` _ \ / _` | | '_ \                        |
#   |                      | | | | | | (_| | | | | |                       |
#   |                      |_| |_| |_|\__,_|_|_| |_|                       |
#   |                                                                      |
#   +----------------------------------------------------------------------+
#   | Main entry point, getopt, etc.                                       |
#   '----------------------------------------------------------------------'

short_options = "ChVTMp:s:O:H:S:i:vG"
long_options = [
    "config",
    "help",
    "version",
    "dump-templates",
    "mark-mode",
    "prefix=",
    "socket=",
    "host-only-header=",
    "host-header=",
    "service-header=",
    "interval=",
    "verbose",
    "include-groups",
    "include-host-icon",
    "debug",
]

opt_verbose = False
opt_debug = False
opt_mark_mode = False
opt_prefix = ""
opt_host_headers = ""
opt_host_only_headers = ""
opt_service_headers = ""
opt_dump_templates = False
opt_groups = False
opt_host_icon = False
opt_check_interval = None

try:
    opts, args = getopt.getopt(sys.argv[1:], short_options, long_options)
except getopt.GetoptError as err:
    sys.stderr.write("%s\n\n" % err)
    usage()
    sys.exit(1)

for o, a in opts:
    # Docu modes
    if o in ["-h", "--help"]:
        usage()
        sys.exit(0)
    elif o in ["-V", "--version"]:
        print_version()
        sys.exit(0)

    # Modifiers
    elif o in ["-T", "--dump-templates"]:
        opt_dump_templates = True
    elif o in ["-M", "--mark-mode"]:
        opt_mark_mode = True
    elif o in ["-p", "--prefix"]:
        opt_prefix = a
    elif o in ["-s", "--socket"]:
        opt_socket = a
    elif o in ["-O", "--host-only-header"]:
        opt_host_only_headers += a + "\n"
    elif o in ["-H", "--host-header"]:
        opt_host_headers += a + "\n"
    elif o in ["-S", "--service-header"]:
        opt_service_headers += a + "\n"
    elif o in ["-i", "--interval"]:
        opt_check_interval = int(a)
    elif o in ["-v", "--verbose"]:
        opt_verbose = True
    elif o in ["-G", "--include-groups"]:
        opt_groups = True
    elif o in ["--include-host-icon"]:
        opt_host_icon = True
    elif o == "--debug":
        opt_debug = True

# Main modes
try:
    g_connection = connect(opt_socket)
    for o, a in opts:
        if o in ["-C", "--config"]:
            if opt_mark_mode:
                sys.stdout.write("config %s\n" % omd_site())
            livedump_config()
            sys.exit(0)
    if opt_mark_mode:
        sys.stdout.write("status\n")
    livedump_state()

except Exception as e:
    if opt_debug:
        raise
    bail_out(str(e))
