#!/usr/bin/env python3
# Copyright (C) 2019 Checkmk 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 argparse
import sys
import urllib.parse

import requests
import urllib3

import cmk.utils.password_store
from cmk.utils.exceptions import MKGeneralException

cmk.utils.password_store.replace_passwords()
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


def main():
    args = parse_arguments()

    try:
        msg, state, perfdata = handle_request(args)
    except Exception as exc:
        sys.stdout.write("UNKNOWN - %s\n" % exc)
        return 3

    sys.stdout.write(f"{msg} | {perfdata}\n")
    sys.exit(state)


def handle_request(args):
    url = urllib.parse.urlunparse(
        (
            args.protocol,
            "%s:%d" % (args.hostname, args.port),
            "%s/_count" % args.index.replace(" ", ","),
            None,
            None,
            None,
        )
    )

    query = {
        "query": {
            "bool": {
                "must": [
                    {"query_string": {"query": args.pattern}},
                    {"range": {"@timestamp": {"gte": "now-%ds" % args.timerange, "lt": "now"}}},
                ]
            }
        },
    }

    if args.fieldname:
        query["query"]["bool"]["must"][0]["query_string"]["fields"] = args.fieldname.split(" ")

    if args.user and args.password:
        raw_response = requests.get(url, json=query, auth=(args.user, args.password))  # nosec B113
    else:
        raw_response = requests.get(url, json=query)  # nosec B113

    msg, state, perfdata = handle_query(raw_response, args.warn, args.crit, args.timerange)

    return msg, state, perfdata


def handle_query(raw_response, warn, crit, timerange):
    response_data = raw_response.json()

    if "count" not in response_data:
        raise MKGeneralException("Missing section count in raw response data")

    state = 0
    value = response_data["count"]
    perfdata = "count=%s" % value
    msg = "%s messages found" % value

    if crit and warn:
        msg += " (warn/crit at %d/%d)" % (warn, crit)
        if value >= crit:
            state = 2
        elif value >= warn:
            state = 1

    return msg, state, perfdata


def parse_arguments(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument(
        "-u",
        "--user",
        default=None,
        help="Username for elasticsearch login",
    )
    parser.add_argument(
        "-s",
        "--password",
        default=None,
        help="Password for easticsearch login",
    )
    parser.add_argument(
        "-P",
        "--protocol",
        default="https",
        help="Use 'http' or 'https' for connection to elasticsearch (default=https)",
    )
    parser.add_argument(
        "-p",
        "--port",
        type=int,
        default=9200,
        help="Use alternative port (default: 9200)",
    )
    parser.add_argument(
        "-q",
        "--pattern",
        help=("Search pattern"),
    )
    parser.add_argument(
        "-f",
        "--fieldname",
        default=None,
        help=("Fieldname to query"),
    )
    parser.add_argument(
        "-i",
        "--index",
        help=("Index to query"),
        default="_all",
    )
    parser.add_argument(
        "-t",
        "--timerange",
        type=int,
        default=60,
        help=("The timerange to query, eg. x minutes from now."),
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help=("Debug mode: let Python exceptions come through"),
    )
    parser.add_argument(
        "--warn",
        type=int,
        default=None,
        help=("number of log messages above which the check will warn"),
    )
    parser.add_argument(
        "--crit",
        type=int,
        default=None,
        help=("number of log messages above which the check will become critical"),
    )
    parser.add_argument(
        "-H",
        "--hostname",
        help=("Defines the elasticsearch instances to query."),
    )

    return parser.parse_args()


if __name__ == "__main__":
    sys.exit(main())
