#!/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
import sys

import lib.args
import lib.base
import lib.lftest
import lib.shell
import lib.txt
from lib.globals import STATE_OK, STATE_UNKNOWN

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

DESCRIPTION = """Checks the number of currently banned IP addresses across all fail2ban jails.
Reports the total ban count and a per-jail breakdown. Alerts when the number of
banned IPs in any jail exceeds the configured thresholds. Requires root or sudo."""

DEFAULT_CRIT = 10000
DEFAULT_WARN = 2500


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

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

    parser.add_argument(
        '--always-ok',
        help=lib.args.help('--always-ok'),
        dest='ALWAYS_OK',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '-c',
        '--critical',
        help='CRIT threshold for the number of banned IPs per jail. '
        'Default: %(default)s',
        dest='CRIT',
        type=int,
        default=DEFAULT_CRIT,
    )

    parser.add_argument(
        '--test',
        help=lib.args.help('--test'),
        dest='TEST',
        type=lib.args.csv,
    )

    parser.add_argument(
        '-w',
        '--warning',
        help='WARN threshold for the number of banned IPs per jail. '
        'Default: %(default)s',
        dest='WARN',
        type=int,
        default=DEFAULT_WARN,
    )

    args, _ = parser.parse_known_args()
    return args


def main():
    """The main function. This is where the magic happens."""

    # parse the command line
    try:
        args = parse_args()
    except SystemExit:
        sys.exit(STATE_UNKNOWN)

    # Pad args.TEST to 3 elements so the per-command fixture
    # loads below that splice `args.TEST[1]` and `args.TEST[2]`
    # into custom lib.lftest.test() calls do not go out of range
    # when the user passes `--test=path` without trailing commas.
    if args.TEST is not None:
        while len(args.TEST) < 3:
            args.TEST.append('')

    # fetch data
    # fail2ban-client ping
    if args.TEST is None:
        stdout, _stderr, retc = lib.base.coe(
            lib.shell.shell_exec('fail2ban-client ping')
        )
    else:
        # do not call the command, put in test data
        stdout, _stderr, retc = lib.lftest.test(
            [args.TEST[0] + '-ping', args.TEST[1], args.TEST[2]]
        )
    if retc != 0:
        lib.base.cu('Problem while testing if the fail2ban server is alive.')

    # fail2ban-client status
    if args.TEST is None:
        stdout, _stderr, retc = lib.base.coe(
            lib.shell.shell_exec('fail2ban-client status')
        )
    else:
        # do not call the command, put in test data
        stdout, _stderr, retc = lib.lftest.test(
            [args.TEST[0] + '-status', args.TEST[1], args.TEST[2]]
        )
    if retc != 0:
        lib.base.cu('Problem while testing the status of the fail2ban server.')

    # extract the jail list
    jail_list = lib.txt.extract_str(stdout, 'Jail list:', '\n').strip()
    if not jail_list:
        lib.base.cu('No jails found.')
    jail_list = jail_list.split(', ')

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

    # analyze data
    #     for each jail_name:
    #         get status of jail_name
    for jail in jail_list:
        if args.TEST is None:
            stdout, _stderr, retc = lib.base.coe(
                lib.shell.shell_exec(f'fail2ban-client status {jail}')
            )
        else:
            # do not call the command, put in test data
            stdout, _stderr, retc = lib.lftest.test(
                [args.TEST[0] + f'-status-{jail}', args.TEST[1], args.TEST[2]],
            )
        if retc != 0:
            lib.base.oao(
                f'Problem while testing the status of the jail "{jail}" on the fail2ban server.',
                STATE_UNKNOWN,
            )

        f2b_currently_banned = lib.txt.extract_str(stdout, 'Currently banned:\t', '\n')
        # important to convert the result to an integer for the comparison later on
        if f2b_currently_banned:
            f2b_currently_banned = int(f2b_currently_banned)
            count += f2b_currently_banned
        else:
            f2b_currently_banned = 0
        jail_state = lib.base.get_state(f2b_currently_banned, args.WARN, args.CRIT)
        state = lib.base.get_worst(state, jail_state)


        # build the message
        msg += (
            f'* {f2b_currently_banned} in jail "{jail}"'
            f'{lib.base.state2str(state, prefix=" ")}\n'
        )
        perfdata += lib.base.get_perfdata(
            jail,
            f2b_currently_banned,
            warn=args.WARN,
            crit=args.CRIT,
            _min=0,
        )

    # over and out
    msg = f'{count} {lib.txt.pluralize("IP", count)} banned\n{msg}'
    lib.base.oao(msg, state, perfdata, always_ok=args.ALWAYS_OK)


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