check_elasticsearch_query

#!/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 argparse
import sys
import urllib.parse
import requests
import urllib3  # type: ignore[import]
from cmk.utils.exceptions import MKGeneralException
import cmk.utils.password_store

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("%s | %s\n" % (msg, perfdata))
    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))
    else:
        raw_response = requests.get(url, json=query)

    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())