check_elasticsearch_query

#!/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
from pathlib import Path

import requests
import urllib3

from cmk.utils import password_store

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


def main():
    args = parse_arguments()

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

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


def _make_auth(
    user: str | None,
    password: str | None,
    password_ref: str | None,
) -> tuple[str, str] | None:
    if user is None:
        return None
    if password is not None:
        return (user, password)
    if password_ref is not None:
        pw_id, pw_file = password_ref.split(":", 1)
        return (user, password_store.lookup(Path(pw_file), pw_id))
    return None


def handle_request(args: argparse.Namespace, auth: tuple[str, str] | None) -> tuple[str, int, str]:
    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(" ")

    raw_response = requests.get(url, json=query, auth=auth)  # nosec B113 # BNS:0b0eac

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

    return msg, state, perfdata


def handle_query(
    raw_response: requests.Response, warn: int | None, crit: int | None
) -> tuple[str, int, str]:
    response_data = raw_response.json()

    if "count" not in response_data:
        raise ValueError("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",
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "-s",
        "--password",
        default=None,
        help="Password for easticsearch login. Preferred over --password-id",
    )
    group.add_argument(
        "--password-id",
        default=None,
        help="Password store reference to the 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())