#!/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_CRIT, STATE_OK, STATE_UNKNOWN, STATE_WARN

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

DESCRIPTION = """Checks the storage pools a libvirt host keeps its virtual machines on. Reports the
storage behind them and the pools themselves in two tables: one per store, with how full
it is and which pools sit on it, and one per pool. Alerts if a store is filling up, if a
pool lost part of what it serves from, if it cannot be reached at all, and if a pool that
is configured to start together with the host is not active. Several pools commonly share
one filesystem, which is reported and alerted on once rather than once per pool. Supports
extended reporting via --lengthy. Runs without root or sudo."""

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

# How each of libvirt's storage pool states is judged. The list of states is libvirt's
# (`virStoragePoolState` in `include/libvirt/libvirt-storage.h`, mirrored by
# `lib.kvm.POOL_STATES`); which of them is worth an alert is not, because libvirt says
# nothing about that. The judgement below rests on the one-line description libvirt
# gives each value, and every one of the five is accounted for:
#
# - `running`, "Running normally": OK.
# - `building`, "Initializing pool, not available": OK. A pool passes through it while
#   it is being created, and a check that happens to catch that moment has found
#   somebody at work rather than a fault.
# - `inactive`, "Not running": OK on its own. A pool nobody asked to start is a normal
#   thing to have. The case that is not normal, a pool set to start with the host and
#   yet not started, is found separately below.
# - `degraded`, "Running degraded": WARN. The pool is serving and has lost part of what
#   it serves from, so it is one more failure away from not serving.
# - `inaccessible`, "Running, but not accessible": CRIT. The pool is up and what it
#   sits on cannot be reached at all, so the machines on it fail their next I/O.
#
# Anything libvirt adds to the enumeration later falls through to WARN: an unrecognised
# state is worth a look, and saying so beats passing it over in silence.
CRIT_STATES = ('inaccessible',)
OK_STATES = ('building', 'inactive', 'running')

# How far the free space two pools report may drift apart and still be the same store.
# The pools are asked one after another, so a filesystem being written to has moved by
# the time the next one is asked; 0.18% of capacity was measured between two pools of
# one filesystem on an idle workstation. A percent leaves room for a busy host and is
# still far below what two genuinely separate filesystems would differ by. The
# grouping itself is `lib.kvm.group_by_store()`, which `kvm-volume` needs as well.
STORE_DRIFT = 0.01

# Where a pool points, out of its definition. libvirt writes the element on a line of
# its own inside `<target>`, and `<target>` is the only place a pool's XML carries a
# `<path>` element at all: a source names its device or directory in an attribute
# instead. The search is scoped to the target anyway, so a pool type that changes that
# later cannot quietly hand over the wrong path. Verified against libvirt 12.0.0.
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 share of a pool that is taken, in percent. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        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 share of a pool that is taken, in percent. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        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.

    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_severity(pool_state):
    """Return the state to report for one of libvirt's storage pool states."""
    if pool_state in CRIT_STATES:
        return STATE_CRIT
    if pool_state in OK_STATES:
        return STATE_OK
    return STATE_WARN


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
    # Asked for per pool rather than read off the `pool-list` table, which is the only
    # way to the two states worth reacting to: without `--details` that table prints
    # `active` for a pool that is running, degraded and inaccessible alike, and with
    # `--details` it rounds the sizes to two decimals with a unit. There are a handful
    # of pools on a host, so the extra calls cost little.
    if args.TEST is None:
        names = lib.base.coe(lib.kvm.get_pools(uri=args.URL, timeout=args.TIMEOUT))
        pools = {}
        for name in names:
            pools[name] = lib.base.coe(
                lib.kvm.get_pool_info(name, uri=args.URL, timeout=args.TIMEOUT)
            )
            pools[name]['path'] = get_pool_path(
                lib.base.coe(
                    lib.kvm.get_pool_xml(name, uri=args.URL, timeout=args.TIMEOUT)
                )
            )
    else:
        # One fixture names the pools, two per pool: what `pool-info` answered for it
        # and what its definition looks like, both named after the first with the pool
        # appended.
        names = [
            line.strip()
            for line in lib.lftest.test_text(args.TEST).splitlines()
            if line.strip()
        ]
        pools = {}
        for name in names:
            pools[name] = lib.kvm.parse_pool_info(
                lib.lftest.test_text(args.TEST, f'{args.TEST[0]}-{name}')
            )
            pools[name]['path'] = get_pool_path(
                lib.lftest.test_text(args.TEST, f'{args.TEST[0]}-{name}-xml')
            )

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

    # init some vars
    counts = {}
    filling_up = []
    filling_up_state_worst = STATE_OK
    msg = ''
    not_active = []
    perfdata = ''
    state = STATE_OK
    # What each pool reports, so pools looking at one and the same store can be found
    # afterwards. libvirt never says what a pool sits on, so it has to be recognised
    # from the figures; see `group_by_store()` for how, and why they cannot simply be
    # compared for equality.
    measurements = []
    table_data = []
    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

        info = pools[name]
        pool_state = info.get('state', 'unknown')
        counts[pool_state] = counts.get(pool_state, 0) + 1
        item_state = get_severity(pool_state)

        # These three describe the storage the pool sits on, not what the pool itself
        # holds. For a directory-backed pool libvirt fills them from a `statvfs()` of
        # the pool's path and works the allocation out as capacity minus free space
        # (`storage_util.c`), so they cover everything on that filesystem, this pool's
        # share of it included. Several pools on one filesystem therefore report it
        # identically, which is why they are named as a group below and why the
        # columns say `Store`. The figures are still the right ones to alert on: a
        # pool runs out when the storage under it does.
        capacity = info.get('capacity', 0)
        allocation = info.get('allocation', 0)
        available = info.get('available', 0)
        # A pool that reports no sizes is not measuring any storage, so it does not
        # belong to a store and is left out here. It is still reported below, as a pool
        # in whatever state it is in.
        if capacity:
            measurements.append(
                {
                    'allocation': allocation,
                    'available': available,
                    'capacity': capacity,
                    'name': name,
                    'path': info.get('path', ''),
                }
            )

        # A pool set to start with the host and yet not started is the one case a plain
        # state count cannot show: `inactive` is a perfectly normal state for a pool
        # nobody asked to start.
        #
        # `inactive` is the only state that means the pool was not started. libvirt
        # spells the rest out in its own enumeration: `degraded` is "Running
        # degraded", `inaccessible` is "Running, but not accessible" and `building` is
        # "Initializing pool, not available" (`libvirt-storage.h`). Reading any of
        # those as "did not start" would report a pool that is up, and busy losing its
        # storage, as one nobody switched on.
        autostart = str(info.get('autostart', '')).lower() == 'yes'
        if autostart and pool_state == 'inactive':
            not_active.append(name)
            item_state = lib.base.get_worst(item_state, STATE_WARN)

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

        table_data.append(
            {
                '_state': item_state,
                'autostart': 'yes' if autostart else 'no',
                'name': name,
                'path': info.get('path') or '-',
                'persistent': str(info.get('persistent', '-')),
                # The state word and the verdict of the row share the last cell:
                # IcingaWeb replaces the marker with an icon and would break the table
                # alignment anywhere else. The verdict is about the pool itself; how
                # full the storage under it is has a table and a verdict of its own.
                'pool_state': pool_state,
                'state': f'{pool_state} '
                f'{lib.base.state2str(item_state, empty_ok=False)}',
            }
        )

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

    # still analysing: the pools are grouped into the stores they sit on, and the
    # figures of each store are worked out before the message can name them
    stores = lib.kvm.group_by_store(measurements, drift=STORE_DRIFT)

    # One set of numbers per store, named after the store rather than after the pools
    # on it. Pools sharing storage report the same figures to the byte, so a metric per
    # pool would draw one line several times over and let a dashboard add one
    # filesystem up as if it were several.
    #
    # The name is the deepest path that contains every pool on the store, which is how
    # `check_disk_usage` labels a filesystem too: four pools spread over `/home` and
    # `/var` of one root filesystem come out as `/`, and a store holding a single pool
    # comes out as that pool's own path. It is worked out from the paths as text and
    # never from the filesystem, so it stays right for a hypervisor checked over ssh.
    # A pool type without a path, RBD for one, falls back to the pool names.
    #
    # It follows that the name moves when the set of pools on the store does. Stop the
    # two pools under `/home` and the store that was `/` is reported as `/var`, and the
    # metrics are renamed with it. That is the price of libvirt never saying what a
    # pool sits on; it is in the README so nobody builds a dashboard on the assumption
    # that the name is fixed.
    store_data = []
    for members in stores:
        # The fullest member's figures stand for the store: they are the same numbers
        # a moment apart, and the least free space is the one worth acting on.
        fullest = max(members, key=lambda item: item['allocation'])
        capacity = fullest['capacity']
        allocation = fullest['allocation']
        available = fullest['available']
        names = sorted(item['name'] for item in members)
        paths = sorted(item['path'] for item in members)
        if all(path.startswith('/') for path in paths):
            store = os.path.commonpath(paths)
        else:
            store = ', '.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('_')
        usage = round(allocation / capacity * 100, 1)
        usage_state = lib.base.get_state(usage, args.WARN, args.CRIT, _operator='range')
        # How full a store is is judged once, here, and not once per pool on it: two
        # pools of one store report figures a moment apart and could otherwise land on
        # either side of a threshold.
        state = lib.base.get_worst(state, usage_state)
        if usage_state != STATE_OK:
            filling_up.append(f'{", ".join(names)} ({usage}%)')
            filling_up_state_worst = lib.base.get_worst(
                filling_up_state_worst, usage_state
            )
        store_data.append(
            {
                '_state': usage_state,
                'available': lib.human.bytes2human(available),
                'pools': ', '.join(names),
                'size': lib.human.bytes2human(capacity),
                'store': store,
                # The verdict of the store, in the last column for IcingaWeb's sake.
                'usage': f'{usage}%{lib.base.state2str(usage_state, prefix=" ")}',
                'used': lib.human.bytes2human(allocation),
                # Sorted on the number rather than on its rendering.
                'raw_usage': usage,
            }
        )
        perfdata += lib.base.get_perfdata(
            f'{store_label}_usage',
            usage,
            uom='%',
            warn=args.WARN,
            crit=args.CRIT,
            _min=0,
            _max=100,
        )
        perfdata += lib.base.get_perfdata(
            f'{store_label}_allocation',
            allocation,
            uom='B',
            _min=0,
            _max=capacity,
        )
        perfdata += lib.base.get_perfdata(
            f'{store_label}_available',
            available,
            uom='B',
            _min=0,
            _max=capacity,
        )
        perfdata += lib.base.get_perfdata(
            f'{store_label}_capacity',
            capacity,
            uom='B',
            _min=0,
        )

    # build the message
    # The verdict and the thresholds first, then whatever a notification has to name
    # without opening the table: which pools are involved, and why.
    if state == STATE_CRIT:
        msg += 'There are critical errors'
    elif state == STATE_WARN:
        msg += 'There are warnings'
    else:
        msg += 'Everything is ok'
    msg += f' (warn={args.WARN} crit={args.CRIT}).'

    # Every sentence carries its own verdict. The line opens with the worst of them,
    # which on a host where a store is filling up and a pool has become unreachable
    # is the pool, and without a marker per sentence nothing says which of the two
    # the CRITICAL belongs to.
    if filling_up:
        msg += (
            f' Filling up: {"; ".join(sorted(filling_up))}'
            f'{lib.base.state2str(filling_up_state_worst, prefix=" ")}.'
        )
    for pool_state in sorted(counts):
        severity = get_severity(pool_state)
        if severity == STATE_OK:
            continue
        involved = sorted(
            row['name'] for row in table_data if row['pool_state'] == pool_state
        )
        msg += (
            f' {pool_state.capitalize()}: {", ".join(involved)}'
            f'{lib.base.state2str(severity, prefix=" ")}.'
        )
    if not_active:
        msg += (
            f' Set to start with the host but not running: {", ".join(not_active)}'
            f'{lib.base.state2str(STATE_WARN, prefix=" ")}.'
        )

    # build table output
    # Two tables, because they answer two questions. The first is about storage: how
    # full it is, and which pools are on it. The second is about the pools themselves.
    # A pool that is not running has no sizes and so appears only in the second.
    #
    # --brief only reshapes them. Every store and 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. It applies to both
    # tables, and a table it empties is left out rather than printed as a header with
    # nothing under it.
    display_stores = (
        [row for row in store_data if row['_state'] != STATE_OK]
        if args.BRIEF
        else store_data
    )
    display_pools = (
        [row for row in table_data if row['_state'] != STATE_OK]
        if args.BRIEF
        else table_data
    )
    if display_stores:
        msg += (
            '\n\n'
            + lib.base.get_table(
                display_stores,
                ['store', 'pools', 'size', 'used', 'available', 'usage'],
                header=['Store', 'Pools', 'Size', 'Used', 'Avail', 'Use%'],
                sort_by_key='raw_usage',
                sort_order_reverse=True,
            ).rstrip()
        )

    if display_pools:
        cols = ['name', 'path', 'autostart', 'state']
        header = ['Pool', 'Path', 'Autostart', 'State']
        if args.LENGTHY:
            cols = ['name', 'path', 'autostart', 'persistent', 'state']
            header = ['Pool', 'Path', 'Autostart', 'Persistent', 'State']
        msg += '\n\n' + lib.base.get_table(display_pools, cols, header=header)

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