#!/usr/bin/env python3
# -*- coding: utf-8; py-indent-offset: 4 -*-
#
# Author:  Linuxfabrik GmbH, Zurich, Switzerland
# Contact: info (at) linuxfabrik (dot) ch
#          https://www.linuxfabrik.ch/
# License: The Unlicense, see LICENSE file.

# https://github.com/Linuxfabrik/monitoring-plugins/blob/main/CONTRIBUTING.md

"""See the check's README for more details.
"""

import argparse  # pylint: disable=C0413
import sys  # pylint: disable=C0413

import lib.base  # pylint: disable=C0413
import lib.db_mysql  # pylint: disable=C0413
import lib.txt  # pylint: disable=C0413
from lib.globals import (STATE_CRIT, STATE_OK,  # pylint: disable=C0413
                          STATE_UNKNOWN, STATE_WARN)


__author__ = 'Linuxfabrik GmbH, Zurich/Switzerland'
__version__ = '2024060601'

DESCRIPTION = """This check connects to a MySQL/MariaDB database and can then run a separate
                 warning and/or critical query against it. The result - the number of items found
                 or a specific number - can be checked against a range expression."""

DEFAULT_DEFAULTS_FILE = '/var/spool/icinga2/.my.cnf'
DEFAULT_DEFAULTS_GROUP = 'client'
DEFAULT_TIMEOUT = 3


def parse_args():
    """Parse command line arguments using argparse.
    """
    parser = argparse.ArgumentParser(description=DESCRIPTION)

    parser.add_argument(
        '-V', '--version',
        action='version',
        version='%(prog)s: v{} by {}'.format(__version__, __author__)
    )

    parser.add_argument(
        '--always-ok',
        help='Always returns OK.',
        dest='ALWAYS_OK',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '-c', '--critical',
        help='Set the CRIT threshold. Supports ranges. ',
        dest='CRIT',
    )

    parser.add_argument(
        '--critical-query',
        help='`SELECT` statement. If its result contains more than one column, the number of rows '
             'is checked against `--critical`, otherwise the single value is used.',
        dest='CRITICAL_QUERY',
    )

    parser.add_argument(
        '--defaults-file',
        help='Specifies a cnf file to read parameters like user, host and password from '
             '(instead of specifying them on the command line), '
             'for example `/var/spool/icinga2/.my.cnf`. '
             'Default: %(default)s',
        dest='DEFAULTS_FILE',
        default=DEFAULT_DEFAULTS_FILE,
    )

    parser.add_argument(
        '--defaults-group',
        help='Group/section to read from in the cnf file. '
             'Default: %(default)s',
        dest='DEFAULTS_GROUP',
        default=DEFAULT_DEFAULTS_GROUP,
    )

    parser.add_argument(
        '--timeout',
        help='Network timeout in seconds. '
             'Default: %(default)s (seconds)',
        dest='TIMEOUT',
        type=int,
        default=DEFAULT_TIMEOUT,
    )

    parser.add_argument(
        '-w', '--warning',
        help='Set the WARN threshold. Supports ranges.',
        dest='WARN',
    )

    parser.add_argument(
        '--warning-query',
        help='`SELECT` statement. If its result contains more than one column, the number of rows '
             'is checked against `--warning`, otherwise the single value is used.',
        dest='WARNING_QUERY',
    )

    return parser.parse_args()


def get_state_and_value(conn, query, threshold, _type):
    """Execute SQL query, get the value and check against the threshold.
    * One row, one column: Check this single value.
    * x rows: Check the number of rows against the threshold.
    """
    state = STATE_OK
    value = 0
    result = []
    shortened = False
    if query:
        result = lib.base.coe(lib.db_mysql.select(conn, query))
        if result:
            if len(result) == 1 and len(result[0]) == 1:
                # one row, one column: could be a "select count(*) from ..." result
                value = list(result[0].values())[0]
            else:
                # a bunch of rows (at least one) with multiple columns, so count them
                value = len(result)
                # shorten the result if there are too many rows
                if len(result) > 10:
                    # shorten the result
                    result = result[0:5] + result[-5:]
                    shortened = True

            if _type == 'warn':
                state = lib.base.get_state(value, threshold, None, _operator='range')
            else:
                state = lib.base.get_state(value, None, threshold, _operator='range')

    return state, value, result, shortened


def main():
    """The main function. Hier spielt die Musik.
    """

    # parse the command line, exit with UNKNOWN if it fails
    try:
        args = parse_args()
    except SystemExit:
        sys.exit(STATE_UNKNOWN)

    if args.WARNING_QUERY is None and args.CRITICAL_QUERY is None:
        lib.base.cu('Nothing to check, no queries provided.')

    # init some vars
    msg = ''
    state = STATE_OK
    perfdata = ''

    mysql_connection = {
        'defaults_file':  args.DEFAULTS_FILE,
        'defaults_group': args.DEFAULTS_GROUP,
        'timeout':        args.TIMEOUT,
    }
    conn = lib.base.coe(lib.db_mysql.connect(mysql_connection))
    lib.base.coe(lib.db_mysql.check_select_privileges(conn))

    # analyze data
    state_warn, cnt_warn, result_warn, shortened_warn = get_state_and_value(
        conn,
        args.WARNING_QUERY,
        args.WARN,
        'warn',
    )
    state = lib.base.get_worst(state, state_warn)
    state_crit, cnt_crit, result_crit, shortened_crit = get_state_and_value(
        conn,
        args.CRITICAL_QUERY,
        args.CRIT,
        'crit',
    )
    state = lib.base.get_worst(state, state_crit)

    lib.db_mysql.close(conn)

    # build the message
    if args.WARNING_QUERY:
        msg = '{} {} from warning query `{}`{}'.format(
            cnt_warn,
            lib.txt.pluralize('result', cnt_warn),
            args.WARNING_QUERY,
            lib.base.state2str(state_warn, prefix=' '),
        )
    if args.WARNING_QUERY and args.CRITICAL_QUERY:
        msg += ' and '
    if args.CRITICAL_QUERY:
        msg += '{} {} from critical query `{}`{}'.format(
            cnt_crit,
            lib.txt.pluralize('result', cnt_crit),
            args.CRITICAL_QUERY,
            lib.base.state2str(state_crit, prefix=' '),
        )
    msg += '\n'

    if shortened_warn:
        msg += '\nAttention: Table below is truncated, showing the 5 first and the 5 last items.\n'
    try:
        keys = result_warn[0].keys()
        headers = keys
        msg += '\n' + lib.base.get_table(result_warn, keys, header=headers)
    except:
        # no results
        pass

    if shortened_crit:
        msg += '\nAttention: Table below is truncated, showing the 5 first and the 5 last items.\n'
    try:
        keys = result_crit[0].keys()
        headers = keys
        msg += '\n' + lib.base.get_table(result_crit, keys, header=headers)
    except:
        # no results
        pass

    perfdata += lib.base.get_perfdata('cnt_warn', cnt_warn, None, args.WARN, None, None, None)
    perfdata += lib.base.get_perfdata('cnt_crit', cnt_crit, None, None, args.CRIT, None, None)

    # over and out
    lib.base.oao(msg, state, perfdata, always_ok=args.ALWAYS_OK)


if __name__ == '__main__':
    try:
        main()
    except Exception:   # pylint: disable=W0703
        lib.base.cu()
