#!/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.human
import lib.kvm
import lib.lftest
import lib.txt
from lib.globals import STATE_OK, STATE_UNKNOWN

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

DESCRIPTION = """Reports what the storage pools of a libvirt host actually hold: how much space
their volumes have been promised, how much of it they occupy today, and how far the
promises exceed the storage underneath. Handing out more than there is is what thin
provisioning is for, and it is safe only for as long as the volumes stay unfilled, so
this is the number that says how much room is left for that to happen. Several pools
commonly share one filesystem, and the promises they make are added up and judged once
per filesystem rather than once per pool. Alerts if the promises exceed the storage by
more than the thresholds allow, which are unset by default because only the
administrator knows how far their storage may be oversubscribed. Supports extended
reporting via --lengthy. Runs without root or sudo."""

DEFAULT_CRIT = None
DEFAULT_NO_MATCH_SEVERITY = 'ok'
DEFAULT_TIMEOUT = lib.kvm.DEFAULT_TIMEOUT
DEFAULT_URL = lib.kvm.DEFAULT_URI
DEFAULT_WARN = None

# How many volumes `--lengthy` lists, the largest first. A pool holds as many volumes
# as the host has disks and old images, 40 in one directory of an ordinary
# workstation, and a table of all of them answers nothing that the ten largest do not.
# Whatever is left out is counted in the caption rather than silently dropped.
LENGTHY_VOLUMES = 10

# The volume types that hold storage. libvirt's own enumeration is `virStorageVolType`
# in include/libvirt/libvirt-storage.h, and all six of it are accounted for:
#
# - `file`, `block`, `network` (RBD and the like) and `ploop` hold data and are
#   counted.
# - `dir` and `netdir` are directories, not storage. A pool over a general-purpose
#   directory lists every subdirectory in it as a volume of capacity zero, which
#   inflates the count with entries that can never fill anything up: 16 of the 1191
#   "volumes" of a pool pointing at a Downloads folder were subdirectories. They are
#   left out, and the directory's contents are not the pool's either, because libvirt
#   does not descend into it.
VOLUME_TYPES = ('block', 'file', 'network', 'ploop')

# How far the free space two pools report may drift apart and still be the same store.
# Same reasoning and same figure as `kvm-storage-pool`, which does the grouping for
# the storage side: the pools are asked one after another, so a filesystem being
# written to has moved by the time the next one is asked.
STORE_DRIFT = 0.01

POOL_TARGET_REGEX = re.compile(r'<target>(.*?)</target>', re.DOTALL)
POOL_PATH_REGEX = re.compile(r'<path>(.*?)</path>', re.DOTALL)


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(
        '--brief',
        help=lib.args.help('--brief'),
        dest='BRIEF',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '-c',
        '--critical',
        help='CRIT threshold for the space a pool has promised its volumes, in '
        'percent of the storage underneath it. '
        'Above 100 the pool has promised more than it has. '
        'Supports Nagios ranges. '
        'Default: unset, the figure is reported but does not alert',
        dest='CRIT',
        default=DEFAULT_CRIT,
    )

    parser.add_argument(
        '--ignore',
        help=lib.args.help('--ignore-regex'),
        dest='IGNORE',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--lengthy',
        help=lib.args.help('--lengthy'),
        dest='LENGTHY',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '--match',
        help=lib.args.help('--match'),
        dest='MATCH',
        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(
        '--test',
        help=lib.args.help('--test'),
        dest='TEST',
        type=lib.args.csv,
    )

    parser.add_argument(
        '--timeout',
        help=lib.args.help('--timeout') + ' Default: %(default)s (seconds)',
        dest='TIMEOUT',
        type=int,
        default=DEFAULT_TIMEOUT,
    )

    parser.add_argument(
        '--url',
        help='libvirt connection URI, passed to `virsh --connect`. '
        'Use `qemu+ssh://user@host/system` to check a host that runs no local '
        'monitoring agent. '
        'Default: %(default)s',
        dest='URL',
        default=DEFAULT_URL,
    )

    parser.add_argument(
        '-w',
        '--warning',
        help='WARN threshold for the space a pool has promised its volumes, in '
        'percent of the storage underneath it. '
        'Above 100 the pool has promised more than it has. '
        'Supports Nagios ranges. '
        'Default: unset, the figure is reported but does not alert',
        dest='WARN',
        default=DEFAULT_WARN,
    )

    args, _ = parser.parse_known_args()
    return args


def get_pool_path(xml):
    """Return where a pool points, or an empty string when its definition says nothing.

    Same reader as `kvm-storage-pool`: libvirt writes the element inside `<target>`,
    and that is the only place a pool's XML carries a `<path>` at all, because a
    source names its device or directory in an attribute instead. A pool that is not
    backed by a path, an iSCSI or RBD pool for instance, has none and gets no made-up
    one.
    """
    target = POOL_TARGET_REGEX.search(xml or '')
    if not target:
        return ''
    path = POOL_PATH_REGEX.search(target.group(1))
    return path.group(1).strip() if path else ''


def get_pool_volumes(uri, timeout):
    """Collect the volumes of every pool the connection knows, with the sizes of the
    storage each pool sits on.

    Returns `(True, {pool: {'capacity': int, 'available': int, 'state': str,
    'volumes': list}})`, or `(False, errormessage)`.

    A pool that is not running is left out. libvirt cannot list what it has not
    opened, and the check that reports a pool for not running is `kvm-storage-pool`.
    """
    success, names = lib.kvm.get_pools(uri=uri, timeout=timeout)
    if not success:
        return False, names

    pools = {}
    for name in names:
        success, info = lib.kvm.get_pool_info(name, uri=uri, timeout=timeout)
        if not success:
            return False, info
        if info.get('state') != 'running':
            continue
        success, volumes = lib.kvm.get_volumes(name, uri=uri, timeout=timeout)
        if not success:
            return False, volumes
        # Where the pool points, which is only in its definition and is what names
        # the store below. A pool type without a path, RBD for one, has none.
        success, xml = lib.kvm.get_pool_xml(name, uri=uri, timeout=timeout)
        if not success:
            return False, xml
        pools[name] = {
            'available': info.get('available', 0),
            'capacity': info.get('capacity', 0),
            'path': get_pool_path(xml),
            'state': info.get('state', ''),
            'volumes': [volume for volume in volumes if volume['type'] in VOLUME_TYPES],
        }
    return True, pools


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.IGNORE is None:
        args.IGNORE = []
    if args.MATCH is None:
        args.MATCH = []

    # fetch data
    if args.TEST is None:
        pools = lib.base.coe(get_pool_volumes(args.URL, args.TIMEOUT))
    else:
        # One fixture names the running pools, and two per pool: what `pool-info`
        # answered for it and what `vol-list --details` did, both named after the
        # first with the pool appended.
        pools = {}
        for name in [
            line.strip()
            for line in lib.lftest.test_text(args.TEST).splitlines()
            if line.strip()
        ]:
            info = lib.kvm.parse_pool_info(
                lib.lftest.test_text(args.TEST, f'{args.TEST[0]}-{name}')
            )
            pools[name] = {
                'available': info.get('available', 0),
                'capacity': info.get('capacity', 0),
                'path': get_pool_path(
                    lib.lftest.test_text(args.TEST, f'{args.TEST[0]}-{name}-xml')
                ),
                'state': info.get('state', ''),
                'volumes': [
                    volume
                    for volume in lib.kvm.parse_volumes(
                        lib.lftest.test_text(args.TEST, f'{args.TEST[0]}-{name}-vols')
                    )
                    if volume['type'] in VOLUME_TYPES
                ],
            }

    if not pools:
        lib.base.oao(
            'No running storage pools found.', STATE_OK, always_ok=args.ALWAYS_OK
        )

    # init some vars
    allocated_total = 0
    msg = ''
    # What each pool reports about its storage, so the pools sitting on one and the
    # same filesystem can be found afterwards. libvirt never says what a pool sits on.
    measurements = []
    oversubscribed = []
    oversubscribed_state_worst = STATE_OK
    perfdata = ''
    store_data = []
    promised_total = 0
    state = STATE_OK
    table_data = []
    volume_data = []
    volumes_total = 0
    compiled_ignore = [
        lib.base.coe(item) for item in lib.txt.compile_regex(args.IGNORE, '--ignore')
    ]
    compiled_match = [
        lib.base.coe(item) for item in lib.txt.compile_regex(args.MATCH, '--match')
    ]

    # analyze data
    for name in sorted(pools):
        if compiled_match and not any(item.search(name) for item in compiled_match):
            continue
        if any(item.search(name) for item in compiled_ignore):
            continue

        pool = pools[name]
        volumes = pool['volumes']
        # What the volumes have been promised, and what they have taken so far. The
        # difference is the room the pool still owes them, and it is the figure that
        # decides whether thin provisioning is safe here: a pool may promise several
        # times its storage as long as the volumes stay unfilled.
        promised = sum(volume['capacity'] for volume in volumes)
        allocated = sum(volume['allocation'] for volume in volumes)
        promised_total += promised
        allocated_total += allocated
        volumes_total += len(volumes)

        # What a pool has promised is compared against its storage further down, not
        # here: several pools commonly sit on one filesystem and each of them reports
        # the whole of it, so a ratio per pool would measure each against a store it
        # shares with the others and no two of them could be added up.
        if pool['capacity']:
            measurements.append(
                {
                    'allocated': allocated,
                    'available': pool['available'],
                    'capacity': pool['capacity'],
                    'name': name,
                    'path': pool['path'],
                    'promised': promised,
                }
            )

        table_data.append(
            {
                'allocated': lib.human.bytes2human(allocated),
                'name': name,
                'promised': lib.human.bytes2human(promised),
                'volumes': len(volumes),
            }
        )
        for volume in volumes:
            volume_data.append(
                {
                    '_allocation': volume['allocation'],
                    'allocated': lib.human.bytes2human(volume['allocation']),
                    'name': volume['name'],
                    'pool': name,
                    'promised': lib.human.bytes2human(volume['capacity']),
                    'type': volume['type'] or '-',
                }
            )

        perfdata_name = re.sub(r'\W+', '_', name)
        perfdata += lib.base.get_perfdata(
            f'{perfdata_name}_volumes',
            len(volumes),
            uom=None,
            _min=0,
        )
        perfdata += lib.base.get_perfdata(
            f'{perfdata_name}_promised',
            promised,
            uom='B',
            _min=0,
        )
        perfdata += lib.base.get_perfdata(
            f'{perfdata_name}_allocated',
            allocated,
            uom='B',
            _min=0,
        )

    # nothing left to report on
    if not table_data:
        lib.base.oao(
            f'Nothing checked. {len(pools)} running storage pools are filtered out '
            f'by `--match` or `--ignore`.',
            lib.base.str2state(args.NO_MATCH_SEVERITY),
            always_ok=args.ALWAYS_OK,
        )

    # One store per filesystem, with the promises of every pool on it added up. This
    # is what the thresholds judge, and the only place that addition is sound: the
    # volumes of two pools are distinct files and do add up, the storage they sit on
    # does not. The store is named after the deepest path holding every pool on it,
    # the way `kvm-storage-pool` names it, so the two checks label one filesystem
    # alike.
    for members in lib.kvm.group_by_store(measurements, drift=STORE_DRIFT):
        promised = sum(item['promised'] for item in members)
        allocated = sum(item['allocated'] for item in members)
        capacity = members[0]['capacity']
        names = sorted(item['name'] for item in members)
        paths = sorted(item['path'] for item in members)
        store = (
            os.path.commonpath(paths)
            if all(path.startswith('/') for path in paths)
            else ', '.join(names)
        )
        # The store is named by a path or by the pools sitting on it, so it carries
        # slashes, commas and spaces. A perfdata label keeps neither of them.
        store_label = re.sub(r'\W+', '_', store).strip('_')
        subscription = round(promised / capacity * 100, 1)
        usage_state = lib.base.get_state(
            subscription, args.WARN, args.CRIT, _operator='range'
        )
        state = lib.base.get_worst(state, usage_state)
        if usage_state != STATE_OK:
            oversubscribed.append(f'{store} ({subscription}%)')
            oversubscribed_state_worst = lib.base.get_worst(
                oversubscribed_state_worst, usage_state
            )
        store_data.append(
            {
                '_state': usage_state,
                'allocated': lib.human.bytes2human(allocated),
                'pools': ', '.join(names),
                'promised': lib.human.bytes2human(promised),
                'size': lib.human.bytes2human(capacity),
                'store': store,
                # The verdict of the store shares the last cell with its figure:
                # IcingaWeb replaces the marker with an icon and would break the
                # table alignment anywhere else.
                'subscription': (
                    f'{subscription}%{lib.base.state2str(usage_state, prefix=" ")}'
                ),
            }
        )
        # Named after the store itself, the way `kvm-storage-pool` names the same
        # filesystem (`/-usage`), so the two checks graph one store under one name.
        perfdata += lib.base.get_perfdata(
            f'{store_label}_subscription',
            subscription,
            uom='%',
            warn=args.WARN,
            crit=args.CRIT,
            _min=0,
        )

    # build the message
    checked = len(table_data)
    msg += (
        f'{checked} pool{"s" if checked != 1 else ""}, '
        f'{volumes_total} volume{"s" if volumes_total != 1 else ""}, '
        f'{lib.human.bytes2human(promised_total)} promised, '
        f'{lib.human.bytes2human(allocated_total)} taken'
    )
    if oversubscribed:
        msg += (
            f'. Promised more than the storage holds: {", ".join(oversubscribed)}'
            f'{lib.base.state2str(oversubscribed_state_worst, prefix=" ")}'
        )
    perfdata += lib.base.get_perfdata(
        'volumes',
        volumes_total,
        uom=None,
        _min=0,
    )
    perfdata += lib.base.get_perfdata(
        'promised',
        promised_total,
        uom='B',
        _min=0,
    )
    perfdata += lib.base.get_perfdata(
        'allocated',
        allocated_total,
        uom='B',
        _min=0,
    )

    # build table output
    # Two tables, because they answer two questions. The first is about storage: how
    # far the promises on it reach, and which pools made them. The second is about the
    # pools themselves, which is where the volumes are counted.
    #
    # --brief only reshapes them. Every pool above has already emitted its performance
    # data and has already driven the overall state, so hiding a row here changes
    # nothing but what a reader has to scroll past.
    display_stores = (
        [row for row in store_data if row['_state'] != STATE_OK]
        if args.BRIEF
        else store_data
    )
    if display_stores:
        # Stripped, because `get_table()` ends in a newline of its own and the next
        # block opens with a blank line: left as it comes, every table is followed by
        # two.
        msg += (
            '\n\n'
            + lib.base.get_table(
                display_stores,
                ['store', 'pools', 'size', 'promised', 'allocated', 'subscription'],
                header=['Store', 'Pools', 'Size', 'Promised', 'Taken', 'Sub%'],
            ).rstrip()
        )
    # No verdict per pool: how far the promises reach is a property of the storage
    # under them and is judged once, above. `--brief` leaves this table out, because
    # a pool row carries nothing that could be within or outside a threshold.
    if table_data and not args.BRIEF:
        msg += (
            '\n\n'
            + lib.base.get_table(
                table_data,
                ['name', 'volumes', 'promised', 'allocated'],
                header=['Pool', 'Volumes', 'Promised', 'Taken'],
            ).rstrip()
        )

    # The largest volumes, which is what answers "what is eating this pool". All of
    # them would be a table nobody reads: one directory of an ordinary workstation
    # held 40. The caption counts the ones left out, so the cut is visible.
    if args.LENGTHY and volume_data:
        volume_data.sort(key=lambda row: row['_allocation'], reverse=True)
        hidden = len(volume_data) - LENGTHY_VOLUMES
        caption = f'The {min(len(volume_data), LENGTHY_VOLUMES)} largest volumes'
        if hidden > 0:
            caption += f', {hidden} more not listed'
        msg += (
            f'\n\n{caption}:\n\n'
            + lib.base.get_table(
                volume_data[:LENGTHY_VOLUMES],
                ['pool', 'name', 'type', 'promised', 'allocated'],
                header=['Pool', 'Volume', 'Type', 'Promised', 'Taken'],
            ).rstrip()
        )

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