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

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

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

# fail2ban-client deserializes the reply it reads from its Unix socket, so a
# socket under an attacker-writable path would let a local user execute code in
# this root process. The system runtime directories are the only places a custom
# socket may live, because an unprivileged user cannot create files there.
ALLOWED_SOCKET_ROOTS = ('/run', '/var/run')

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_NO_MATCH_SEVERITY = 'ok'
DEFAULT_WARN = '2500'


def parse_args():
    """Parse command line arguments using argparse."""
    parser = argparse.ArgumentParser(
        description=DESCRIPTION,
        epilog=lib.args.epilog(__file__),
        formatter_class=lib.args.HelpFormatter,
    )

    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(
        '--no-match-severity',
        help=lib.args.help('--no-match-severity') + ' Default: %(default)s',
        dest='NO_MATCH_SEVERITY',
        choices=['ok', 'warn', 'crit', 'unknown'],
        default=DEFAULT_NO_MATCH_SEVERITY,
    )

    parser.add_argument(
        '--no-perfdata',
        help=lib.args.help('--no-perfdata'),
        dest='NO_PERFDATA',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '--socket',
        help='Path to the fail2ban server Unix socket. '
        'Passed to `fail2ban-client --socket`. '
        'Must resolve within /run or /var/run. '
        '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.
    # This plugin runs as root and fail2ban-client deserializes the reply it reads from
    # the socket, so a socket under an attacker-writable path is a root code-execution
    # primitive. Confine --socket to the system runtime directories, which an
    # unprivileged user cannot write into.
    client = ['fail2ban-client']
    if args.SOCKET:
        if not lib.disk.is_within(os.path.realpath(args.SOCKET), ALLOWED_SOCKET_ROOTS):
            lib.base.cu(
                f'Refusing to use socket "{args.SOCKET}": it resolves outside '
                f'{" and ".join(ALLOWED_SOCKET_ROOTS)}.'
            )
        client += ['--socket', args.SOCKET]

    # fetch data
    # fail2ban-client ping
    if args.TEST is None:
        stdout, _stderr, retc = lib.base.coe(lib.shell.shell_exec([*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([*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([*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.',
            lib.base.str2state(args.NO_MATCH_SEVERITY),
            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, no_perfdata=args.NO_PERFDATA
    )


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