#!/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 re
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__ = '2026052701'

DESCRIPTION = """Checks the number of currently banned IP addresses across all fail2ban jails.
Reports the total ban count and a per-jail breakdown. Jails can be excluded from the
check with a regular expression. 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. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        dest='CRIT',
        default=DEFAULT_CRIT,
    )

    parser.add_argument(
        '--ignore',
        help='Exclude jails whose name matches this Python regular expression. '
        'Case-sensitive by default; use `(?i)` for case-insensitive matching. '
        'Can be specified multiple times. '
        'Example: `--ignore="^recidive$"`. '
        'Default: %(default)s',
        dest='IGNORE',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--socket',
        help='Path to the fail2ban server Unix socket. '
        'Passed to `fail2ban-client --socket`. '
        'Default: %(default)s',
        dest='SOCKET',
        default=None,
    )

    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. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        dest='WARN',
        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('')

    # compile the ignore regexes early so an invalid pattern fails fast
    ignore_patterns = []
    if args.IGNORE:
        try:
            ignore_patterns = [re.compile(p) for p in args.IGNORE]
        except re.error as e:
            lib.base.cu(f'Invalid regular expression: {e}')

    # build the fail2ban-client command prefix, optionally pointing at a custom socket
    client = 'fail2ban-client'
    if args.SOCKET:
        client += f' --socket {args.SOCKET}'

    # fetch data
    # fail2ban-client ping
    if args.TEST is None:
        stdout, _stderr, retc = lib.base.coe(lib.shell.shell_exec(f'{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(f'{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
    perfdata = ''
    state = STATE_OK
    count = 0
    table_data = []

    # analyze data
    #     for each jail_name:
    #         get status of jail_name
    for jail in jail_list:
        # skip jails matching one of the ignore regexes
        if any(pattern.search(jail) for pattern in ignore_patterns):
            continue

        if args.TEST is None:
            stdout, _stderr, retc = lib.base.coe(
                lib.shell.shell_exec(f'{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, _operator='range'
        )
        state = lib.base.get_worst(state, jail_state)

        table_data.append(
            {
                'jail': jail,
                'banned': f2b_currently_banned,
                'state': lib.base.state2str(jail_state, empty_ok=False),
            }
        )
        perfdata += lib.base.get_perfdata(
            jail,
            f2b_currently_banned,
            warn=args.WARN,
            crit=args.CRIT,
            _min=0,
        )

    # nothing left after applying the ignore filter
    if not table_data:
        lib.base.oao('Nothing checked.', STATE_OK, always_ok=args.ALWAYS_OK)

    # build the message
    msg = f'{count} {lib.txt.pluralize("IP", count)} banned\n\n'
    msg += lib.base.get_table(
        table_data,
        ['jail', 'banned', 'state'],
        header=['Jail', 'Banned', 'Status'],
        sort_by_key='banned',
        sort_order_reverse=True,
    )

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


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