#!/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.net  # pylint: disable=C0413
from lib.globals import (STATE_CRIT, STATE_OK,  # pylint: disable=C0413
                          STATE_UNKNOWN, STATE_WARN)

try:
    import psutil  # pylint: disable=C0413
except ImportError:
    print('Python module "psutil" is not installed.')
    sys.exit(STATE_UNKNOWN)


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

DESCRIPTION = '''Counts system-wide socket connections like tcp, tcp6, udp or udp6. If you have
                 too many connections like TCP_CLOSE and therefore get errors like 
                 "too many files open", reconfigure and/or restart the application that is
                 receiving or processing those connections.'''

DEFAULT_CONN_STATUS = 'all'
DEFAULT_CONN_TYPE = 'all'
DEFAULT_CRIT = None
DEFAULT_WARN = None


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(
        '--conn-status',
        help='Filter the status of the connections (repeating). Default: %(default)s',
        default=None,   # due to https://bugs.python.org/issue16399, see in main() below
        dest='CONN_STATUS',
        action='append',
        choices=[
            'all',
            'close',
            'close_wait',
            'closing',
            'established',
            'fin_wait1',
            'fin_wait2',
            'last_ack',
            'listen',
            'none',
            'syn_recv',
            'syn_sent',
            'time_wait',
        ],
    )

    parser.add_argument(
        '--conn-type',
        help='Filter the family/type of the connections (repeating). Default: %(default)s',
        default=None,   # due to https://bugs.python.org/issue16399, see in main() below
        dest='CONN_TYPE',
        action='append',
        choices=[
            'all',
            'tcp',
            'tcp6',
            'udp',
            'udp6',
        ],
    )

    parser.add_argument(
        '-c', '--critical',
        help='Threshold for the number of connections. Type: None or Range. Default: %(default)s',
        default=DEFAULT_CRIT,
        dest='CRIT',
    )

    parser.add_argument(
        '-w', '--warning',
        help='Threshold for the number of connections. Type: None or Range. Default: %(default)s',
        default=DEFAULT_WARN,
        dest='WARN',
    )

    return parser.parse_args()


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)

    # due to https://bugs.python.org/issue16399, set the default value here
    if args.CONN_STATUS is None:
        args.CONN_STATUS = DEFAULT_CONN_STATUS
    if args.CONN_TYPE is None:
        args.CONN_TYPE = DEFAULT_CONN_TYPE

    # fetch data
    stats = {}
    for c in psutil.net_connections(kind='inet'):
        if 'all' not in args.CONN_STATUS \
            and c.status.lower() not in args.CONN_STATUS:
            continue
        if 'all' not in args.CONN_TYPE \
            and lib.net.PROTO_MAP[(c.family, c.type)] not in args.CONN_TYPE:
            continue
        key = '{}_{}'.format(lib.net.PROTO_MAP[(c.family, c.type)], c.status)
        if key in stats:
            stats[key] += 1
        else:
            stats[key] = 1

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

    # build the message
    for item in lib.base.sort(stats, reverse=True):
        proto, value = item
        if value:
            conn_state = STATE_OK
            if not lib.base.coe(lib.base.match_range(value, args.CRIT)):
                conn_state = STATE_CRIT
            elif not lib.base.coe(lib.base.match_range(value, args.WARN)):
                conn_state = STATE_WARN
            state = lib.base.get_worst(conn_state, state)
            msg += '{}: {}{}, '.format(
                proto.replace('_', ' '),
                value,
                lib.base.state2str(conn_state, prefix=' '),
            )
        perfdata += lib.base.get_perfdata(proto, value, None, None, None, None, None)
    if not msg:
        msg = 'No connections of type "{}" in status "{}" found.  '.format(
            ','.join(args.CONN_TYPE),
            ','.join(args.CONN_STATUS),
        )

    # over and out
    lib.base.oao(msg[:-2], STATE_OK, perfdata)


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