#!/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 json
import re
import sys
from collections import namedtuple

import lib.args
import lib.base
import lib.human
import lib.lftest
from lib.globals import STATE_CRIT, STATE_OK, STATE_UNKNOWN, STATE_WARN

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

# psutil 5.x exposes sdiskpart/sdiskusage under psutil._common,
# psutil 6+ under psutil._ntuples. Fall back to plain namedtuples
# with the fields we touch if neither is available.
try:
    from psutil._ntuples import sdiskpart, sdiskusage
except ImportError:
    try:
        from psutil._common import sdiskpart, sdiskusage
    except ImportError:
        sdiskpart = namedtuple(
            'sdiskpart', ['device', 'mountpoint', 'fstype', 'opts']
        )
        sdiskusage = namedtuple(
            'sdiskusage', ['total', 'used', 'free', 'percent']
        )


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

DESCRIPTION = """Checks used or free disk space for each mounted partition. By default, only physical
devices are checked (hard disks, USB drives), ignoring pseudo and memory filesystems.
Supports filtering by mountpoint pattern or filesystem type. Thresholds can be set as
percentages or absolute values, and can target either used or free space. On systems
with many filesystems (hundreds of mounts), --brief hides rows that are within the
thresholds so the table only shows the filesystems in WARN/CRIT state. Note that on
Unix systems, 5% of disk space is typically reserved for root and not reflected in the
available space shown to regular users.
Alerts when usage exceeds the configured thresholds."""

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=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(
        '--brief',
        help='Hide table rows for filesystems within the thresholds and show only '
        'those in WARN/CRIT state. Inverse of `--lengthy` (which adds columns); '
        '`--brief` filters rows. The two are orthogonal and can be combined. '
        'Perfdata and alerting are unaffected: all filesystems still emit perfdata '
        'and still drive the overall check state. '
        'Default: %(default)s',
        dest='BRIEF',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '-c',
        '--critical',
        help='CRIT threshold in the form `<number>[unit][method]`. '
        'Unit is one of `%%|K|M|G|T|P` (default: `%%`). `K` means kibibyte etc. '
        'Method is one of `USED|FREE` (default: `USED`). '
        '`USED` means "number or more", `FREE` means "number or less". '
        'Examples: `95` = 95%% used. `9.5M` = 9.5 MiB used. `5%%FREE`. `1400GUSED`. '
        'Default: %(default)s',
        dest='CRIT',
        type=lib.args.number_unit_method,
        default=DEFAULT_CRIT,
    )

    parser.add_argument(
        '--exclude-pattern',
        help='Exclude any mountpoint containing this substring (case-insensitive). '
        'Example: `boot` excludes `/boot` and `/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=None,
    )

    parser.add_argument(
        '--exclude-regex',
        help='Exclude any mountpoint matching this Python regex (case-insensitive). '
        '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=None,
    )

    parser.add_argument(
        '--fstype',
        help='Override the default behaviour (check physical devices only) and check these '
        'file system types instead. '
        'Can be specified multiple times. '
        'Run `disk-usage --list-fstypes` first to see available types (they are machine dependent).',
        dest='FSTYPE',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--include-pattern',
        help='Only include mountpoints containing this substring (case-insensitive). '
        'Example: `boot` includes `/boot` and `/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=None,
    )

    parser.add_argument(
        '--include-regex',
        help='Only include mountpoints matching this Python regex (case-insensitive). '
        '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=None,
    )

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

    parser.add_argument(
        '--perfdata-regex',
        help='Only emit perfdata keys matching this Python regex. '
        'For a list of perfdata keys, see the README or run this plugin. '
        'Can be specified multiple times.',
        action='append',
        dest='PERFDATA_REGEX',
        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 in the form `<number>[unit][method]`. '
        'Unit is one of `%%|K|M|G|T|P` (default: `%%`). `K` means kibibyte etc. '
        'Method is one of `USED|FREE` (default: `USED`). '
        '`USED` means "number or more", `FREE` means "number or less". '
        'Examples: `95` = 95%% used. `9.5M` = 9.5 MiB used. `5%%FREE`. `1400GUSED`. '
        'Default: %(default)s',
        dest='WARN',
        type=lib.args.number_unit_method,
        default=DEFAULT_WARN,
    )

    args, _ = parser.parse_known_args()
    return args


def _load_disk_usage_fixture(raw_json):
    """Convert a test fixture into the two data structures the plugin
    expects from `psutil.disk_partitions()` and
    `psutil.disk_usage(mountpoint)`:

    - `partitions`: a list of `sdiskpart(device, mountpoint, fstype, opts)`
      namedtuples
    - `usage`: a dict keyed by mountpoint, each value an
      `sdiskusage(total, used, free, percent)` namedtuple

    Fixture shape:

        {
          "partitions": [
            {"device": "/dev/vda1", "mountpoint": "/", "fstype": "ext4",
             "opts": "rw,relatime"},
            ...
          ],
          "usage": {
            "/": {"total": <bytes>, "used": <bytes>, "free": <bytes>,
                  "percent": <0..100>},
            ...
          }
        }
    """
    data = json.loads(raw_json)
    partitions = [
        sdiskpart(
            p['device'],
            p['mountpoint'],
            p.get('fstype', ''),
            p.get('opts', ''),
        )
        for p in data.get('partitions', [])
    ]
    usage = {
        mp: sdiskusage(u['total'], u['used'], u['free'], u['percent'])
        for mp, u in data.get('usage', {}).items()
    }
    return partitions, usage


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(
            f'Your {what} "{regex}" contains one or more errors: {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. This is where the magic happens."""

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

    # set default values for append parameters that were not specified
    if args.EXCLUDE_PATTERN is None:
        args.EXCLUDE_PATTERN = []
    if args.EXCLUDE_REGEX is None:
        args.EXCLUDE_REGEX = []
    if args.FSTYPE is None:
        args.FSTYPE = []
    if args.INCLUDE_PATTERN is None:
        args.INCLUDE_PATTERN = []
    if args.INCLUDE_REGEX is None:
        args.INCLUDE_REGEX = []
    if args.PERFDATA_REGEX is None:
        args.PERFDATA_REGEX = []

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

    # fetch data
    fixture_usage = None
    if args.TEST is None:
        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).'
                )
    else:
        stdout, _, _ = lib.lftest.test(args.TEST)
        parts, fixture_usage = _load_disk_usage_fixture(stdout)

    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:
            if fixture_usage is not None:
                if part.mountpoint not in fixture_usage:
                    raise FileNotFoundError(part.mountpoint)
                usage = fixture_usage[part.mountpoint]
            else:
                usage = psutil.disk_usage(part.mountpoint)
        except (PermissionError, FileNotFoundError, OSError):
            table_data.append(
                {
                    'mountpoint': f'{part.mountpoint}',
                    'type': f'{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 = f'{part.mountpoint}-usage'
        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 = f'{part.mountpoint}-total'
        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 = f'{part.mountpoint}-percent'
        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': f'{part.mountpoint}',
                'type': f'{part.fstype}',
                'used': lib.human.bytes2human(usage.used),
                'avail': lib.human.bytes2human(usage.free),
                'size': lib.human.bytes2human(usage.total),
                'percent': f'{usage.percent}%{lib.base.state2str(disk_state, prefix=" ")}',
                # internal per-row state for the --brief filter below;
                # not rendered because get_table() uses an explicit
                # column whitelist.
                '_state': disk_state,
            }
        )

    # Filter table rows for --brief display. --brief is the inverse of
    # --lengthy: --lengthy adds columns, --brief filters rows. They
    # are orthogonal and both can be set at the same time. Perfdata
    # and alerting stay untouched above this point; --brief only
    # reshapes the human-readable output.
    if args.BRIEF:
        display_rows = [
            row for row in table_data if row.get('_state', STATE_OK) != STATE_OK
        ]
    else:
        display_rows = table_data

    # build the message
    thresholds = f'warn={"".join(args.WARN)} crit={"".join(args.CRIT)}'
    if not table_data:
        msg = 'Nothing checked.'
    elif args.BRIEF and not display_rows:
        # --brief and nothing worth showing: one-line summary only.
        msg = f'Everything is ok. ({thresholds})'
    elif len(table_data) == 1 and not args.BRIEF:
        # single-partition system (no --brief): full single-line
        # summary with the values inline.
        msg = (
            f'{table_data[0]["mountpoint"]}'
            f' {table_data[0]["percent"]}'
            f' - total: {table_data[0]["size"]}'
            f', free: {table_data[0]["avail"]}'
            f', used: {table_data[0]["used"]}'
            f' ({thresholds})'
        )
    else:
        if state == STATE_CRIT:
            header = 'There are critical errors.'
        elif state == STATE_WARN:
            header = 'There are warnings.'
        else:
            header = 'Everything is ok.'
        table = lib.base.get_table(
            display_rows,
            ['mountpoint', 'type', 'size', 'used', 'avail', 'percent'],
            ['Mountpoint', 'Type', 'Size', 'Used', 'Avail', 'Use%'],
            'percent',
        )
        msg = f'{header} ({thresholds})\n\n{table}'

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


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