#!/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 re  # pylint: disable=C0413
import sys  # pylint: disable=C0413

import lib.args  # pylint: disable=C0413
import lib.base  # pylint: disable=C0413
import lib.human  # 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__ = '2025061501'

DESCRIPTION = 'Checks the used disk space, for each partition.'

DEFAULT_WARN = '90%USED'
DEFAULT_CRIT = '95%USED'


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

    parser.add_argument(
        '-V', '--version',
        action='version',
        version='{0}: v{1} by {2}'.format('%(prog)s', __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='Critical threshold, '
             'of the form "<number>[unit][method]", where unit is '
             'one of `%%|K|M|G|T|P` and method is one of `USED|FREE`. If "unit" is omitted, '
             '`%%` is assumed. `K` means `kibibyte` etc. If "method" is omitted, '
             '`USED` is assumed. `USED` means "number ore more", `FREE` means "number or less". '
             'Examples: '
             '`95` = alert at 95%% usage or more. '
             '`9.5M` = alert when 9.5 MiB or more is used. '
             'Other self-explanatory examples are '
             '`95%%USED`, `5%%FREE`, `9.5GFREE`, `1400GUSED`. '
             'Default: %(default)s',
        dest='CRIT',
        type=lib.args.number_unit_method,
        default=DEFAULT_CRIT,
    )

    parser.add_argument(
        '--exclude-pattern',
        help='Any mountpoint matching this pattern (case-insensitive) will count as an exclude. '
             'The mountpoint is excluded if it contains the specified value. '
             'Example: "boot" excludes "/boot" as well as "/boot/efi". '
             'Can be specified multiple times. '
             'On Windows, use drive letters without backslash ("Y:" or "Y"). '
             'Includes are matched before excludes.',
        dest='EXCLUDE_PATTERN',
        action='append',
        default=[],
    )

    parser.add_argument(
        '--exclude-regex',
        help='Any mountpoint matching this python regex (case-insensitive) will count '
             'as an exclude. '
             'Can be specified multiple times. '
             'On Windows, use drive letters without backslash ("Y:" or "Y"). '
             'Includes are matched before excludes.',
        dest='EXCLUDE_REGEX',
        action='append',
        default=[],
    )

    parser.add_argument(
        '--fstype',
        help='By default, this plugin only checks physical devices (e.g. hard disks, '
             'CD-ROM drives, USB keys) and ignores all others (e.g. pseudo, memory, duplicate, '
             'inaccessible file systems). You can override this behaviour with this parameter by '
             'specifying which file system types should be checked explicitly. '
             'Can be specified multiple times. '
             'To get a list of file system types you can specify, run `disk-usage --list-fstype` '
             'first (as file system types are machine dependent).',
        dest='FSTYPE',
        action='append',
        default=[],
    )

    parser.add_argument(
        '--include-pattern',
        help='Any mountpoint matching this pattern (case-insensitive) will count as an include. '
             'The mountpoint is included if it contains the specified value. '
             'Example: "boot" includes "/boot" as well as "/boot/efi". '
             'Can be specified multiple times. '
             'On Windows, use drive letters without backslash ("Y:" or "Y"). '
             'Includes are matched before excludes.',
        dest='INCLUDE_PATTERN',
        action='append',
        default=[],
    )

    parser.add_argument(
        '--include-regex',
        help='Any mountpoint matching this python regex (case-insensitive) will count '
             'as an include. '
             'Can be specified multiple times. '
             'On Windows, use drive letters without backslash ("Y:" or "Y"). '
             'Includes are matched before excludes.',
        dest='INCLUDE_REGEX',
        action='append',
        default=[],
    )

    parser.add_argument(
        '--list-fstypes',
        help='Show which file system types are available and which are checked by default.',
        dest='LIST_FSTYPES',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '--perfdata-regex',
        help='Only print perfdata keys matching this python regex. '
             'For a list of perfdata keys, have a look at the README and run this plugin. '
             'Can be specified multiple times.',
        action='append',
        dest='PERFDATA_REGEX',
        default=[],
    )

    parser.add_argument(
        '-w', '--warning',
        help='Warning threshold, '
             'of the form "<number>[unit][method]", where unit is '
             'one of `%%|K|M|G|T|P` and method is one of `USED|FREE`. If "unit" is omitted, '
             '`%%` is assumed. `K` means `kibibyte` etc. If "method" is omitted, '
             '`USED` is assumed. `USED` means "number ore more", `FREE` means "number or less". '
             'Examples: '
             '`95` = alert at 95%% usage. '
             '`9.5M` = alert when 9.5 MiB is used. '
             'Other self-explanatory examples are '
             '`95%%USED`, `5%%FREE`, `9.5GFREE`, `1400GUSED`. '
             'Default: %(default)s',
        dest='WARN',
        type=lib.args.number_unit_method,
        default=DEFAULT_WARN,
    )

    return parser.parse_args()


def compile_regex(regex, what):
    """Return a compiled regex.
    """
    try:
        return [re.compile(item, re.IGNORECASE) for item in regex]
    except re.error as e:
        # don't want a full stacktrace here, therefore not using cu()
        lib.base.oao(
            'Your {} "{}" contains one or more errors: {}'.format(
                what,
                regex,
                e,
            ),
            STATE_UNKNOWN,
        )


def list_fstypes():
    """Print a nice table showing which file system types exist and which are checked by default.
    """
    # get all partitions, no matter which ones
    parts = psutil.disk_partitions(all=True)
    table_data = []
    for part in parts:
        table_data.append({
            'fstype': part.fstype,
            'mountpoint': part.mountpoint,
            'device': part.device,
            'checked': False,
        })

    # get which ones are checked by default
    try:
        parts = psutil.disk_partitions(all=False)
    except AttributeError:
        pass
    for i, item in enumerate(table_data):
        for part in parts:
            if part.mountpoint == item['mountpoint']:
                table_data[i]['checked'] = True
                continue

    # sort table by fstype, mountpoint
    keys = [
      'fstype',
      'mountpoint',
      'device',
      'checked'
    ]
    lib.base.oao(lib.base.get_table(table_data, keys, header=keys, sort_by_key='fstype'), STATE_OK)


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)
    # args.WARN[0] = number, args.WARN[1] = unit, args.WARN[2] = USED|FREE
    try:
        float(args.WARN[0])
        float(args.CRIT[0])
    except ValueError:
        lib.base.cu('Invalid parameter value.')

    # show partition information and exit
    if args.LIST_FSTYPES:
        list_fstypes()

    # init some vars
    state = STATE_OK
    perfdata = ''
    table_data = []
    compiled_include_regex = compile_regex(args.INCLUDE_REGEX, 'include-regex')
    compiled_exclude_regex = compile_regex(args.EXCLUDE_REGEX, 'exclude-regex')
    compiled_perfdata_regex = compile_regex(args.PERFDATA_REGEX, 'perfdata-regex')

    # analyze data
    if args.FSTYPE:
        # user wants to check file system types on his own
        parts = psutil.disk_partitions(all=True)
    else:
        # default behaviour - check physical devices only (e.g. hard disks, cd-rom drives, USB keys)
        # and ignore all others (e.g. pseudo, memory, duplicate, inaccessible filesystems)
        try:
            parts = psutil.disk_partitions(all=False)
        except AttributeError:
            lib.base.cu('Did not find physical devices (e.g. hard disks, cd-rom drives, USB keys).')

    for part in parts:
        # sdiskpart(device='/dev/vda2', mountpoint='/', fstype='ext4', opts='rw,relatime')
        # sdiskpart(
        #   device='/dev/sr0',
        #   mountpoint='/run/media/root/CentOS 7 x86_64',
        #   fstype='iso9660',
        #   opts='ro,nosuid,nodev,relatime,uid=0,gid=0,iocharset=utf8,mode=0400,dmode=0500'
        # )
        # ignore `/snap`, iso mountpoints and cdroms (UDF = universal disk format)
        if args.FSTYPE:
            # user wants to check file system types on his own
            if part.fstype not in args.FSTYPE:
                continue
        else:
            # default behaviour - ignore read-only and some other filesystems
            if part.fstype in ['CDFS', 'iso9660', 'squashfs', 'UDF'] or part.opts in ['cdrom']:
                continue

        # include (first) and exclude (second) specific partitions
        # hint: we can't do `if not part or part.mountpoint in args.IGNORE:` because it is
        # impossible to specify a "Y:\" on the command line ('Y:\' or 'Y:\\' all don't work).
        mountpoint = part.mountpoint.lower()
        if args.INCLUDE_PATTERN or args.INCLUDE_REGEX:
            if not any(include_pattern.lower() in mountpoint for include_pattern in args.INCLUDE_PATTERN) \
            and not any(item.search(mountpoint) for item in compiled_include_regex):
                continue
        if args.EXCLUDE_PATTERN or args.EXCLUDE_REGEX:
            if any(exclude_pattern.lower() in mountpoint for exclude_pattern in args.EXCLUDE_PATTERN) \
            or any(item.search(mountpoint) for item in compiled_exclude_regex):
                continue

        try:
            usage = psutil.disk_usage(part.mountpoint)
        except (PermissionError, FileNotFoundError, OSError):
            table_data.append({
                'mountpoint': '{}'.format(part.mountpoint),
                'type': '{}'.format(part.fstype),
                'used': 'N/A',
                'avail': 'N/A',
                'size': 'N/A',
                'percent': 'N/A',
            })
            continue

        # check % vs. USED|FREE and K|M|G|T|P vs. USED|FREE, first WARN, then CRIT
        disk_state = STATE_OK

        if args.WARN[1] == '%' and args.WARN[2] == 'USED':
            disk_state = lib.base.get_state(
                usage.percent,
                args.WARN[0],
                None,
                'ge',
            )
        elif args.WARN[1] == '%' and args.WARN[2] == 'FREE':
            disk_state = lib.base.get_state(
                100.0 - usage.percent,
                args.WARN[0],
                None,
                'le',
            )
        elif args.WARN[1] != '%' and args.WARN[2] == 'USED':
            disk_state = lib.base.get_state(
                usage.used,
                lib.human.human2bytes(''.join(args.WARN[:2])),
                None,
                'ge',
            )
        elif args.WARN[1] != '%' and args.WARN[2] == 'FREE':
            disk_state = lib.base.get_state(
                usage.free,
                lib.human.human2bytes(''.join(args.WARN[:2])),
                None,
                'le',
            )

        if args.CRIT[1] == '%' and args.CRIT[2] == 'USED':
            disk_state = lib.base.get_worst(
                disk_state,
                lib.base.get_state(
                    usage.percent,
                    None,
                    args.CRIT[0],
                    'ge',
                ),
            )
        elif args.CRIT[1] == '%' and args.CRIT[2] == 'FREE':
            disk_state = lib.base.get_worst(
                disk_state,
                lib.base.get_state(
                    100.0 - usage.percent,
                    None,
                    args.CRIT[0],
                    'le',
                ),
            )
        elif args.CRIT[1] != '%' and args.CRIT[2] == 'USED':
            disk_state = lib.base.get_worst(
                disk_state,
                lib.base.get_state(
                    usage.used,
                    None,
                    lib.human.human2bytes(''.join(args.CRIT[:2])),
                    'ge',
                ),
            )
        elif args.CRIT[1] != '%' and args.CRIT[2] == 'FREE':
            disk_state = lib.base.get_worst(
                disk_state,
                lib.base.get_state(
                    usage.free,
                    None,
                    lib.human.human2bytes(''.join(args.CRIT[:2])),
                    'le',
                ),
            )

        state = lib.base.get_worst(state, disk_state)

        perfdata_key = '{}-usage'.format(part.mountpoint)
        if not args.PERFDATA_REGEX \
        or any(item.search(perfdata_key) for item in compiled_perfdata_regex):
            perfdata += lib.base.get_perfdata(
                perfdata_key,
                usage.used,
                uom='B',
                warn=None,
                crit=None,
                _min=0,
                _max=usage.total,
            )
        perfdata_key = '{}-total'.format(part.mountpoint)
        if not args.PERFDATA_REGEX \
        or any(item.search(perfdata_key) for item in compiled_perfdata_regex):
            perfdata += lib.base.get_perfdata(
                perfdata_key,
                usage.total,
                uom='B',
                warn=None,
                crit=None,
                _min=0,
                _max=usage.total,
            )
        perfdata_key = '{}-percent'.format(part.mountpoint)
        if not args.PERFDATA_REGEX \
        or any(item.search(perfdata_key) for item in compiled_perfdata_regex):
            perfdata += lib.base.get_perfdata(
                perfdata_key,
                usage.percent,
                uom='%',
                warn=None,
                crit=None,
                _min=0,
                _max=100,
            )
        table_data.append({
            'mountpoint': '{}'.format(part.mountpoint),
            'type': '{}'.format(part.fstype),
            'used': lib.human.bytes2human(usage.used),
            'avail': lib.human.bytes2human(usage.free),
            'size': lib.human.bytes2human(usage.total),
            'percent': '{}%{}'.format(usage.percent, lib.base.state2str(disk_state, prefix=' ')),
        })

    # build the message
    if not table_data:
        msg = 'Nothing checked.'
    elif len(table_data) == 1:
        msg = '{} {} - total: {}, free: {}, used: {} (warn={} crit={})'.format(
            table_data[0]['mountpoint'],
            table_data[0]['percent'],
            table_data[0]['size'],
            table_data[0]['avail'],
            table_data[0]['used'],
            ''.join(args.WARN),
            ''.join(args.CRIT),
        )
    else:
        if state == STATE_CRIT:
            msg = 'There are critical errors.'
        elif state == STATE_WARN:
            msg = 'There are warnings.'
        else:
            msg = 'Everything is ok.'
        msg = '{} (warn={} crit={})\n\n{}'.format(
            msg,
            ''.join(args.WARN),
            ''.join(args.CRIT),
            lib.base.get_table(
                table_data,
                ['mountpoint', 'type', 'size', 'used', 'avail', 'percent'],
                ['Mountpoint', 'Type', 'Size', 'Used', 'Avail', 'Use%'],
                'percent',
            ),
        )

    # 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()
