#!/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

import lib.args
import lib.base
import lib.container
import lib.db_sqlite
import lib.lftest
import lib.time
import lib.txt
from lib.globals import STATE_CRIT, STATE_OK, STATE_UNKNOWN, STATE_WARN

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

DESCRIPTION = """Reports CPU and memory usage for all running Docker containers. CPU usage is
normalized by dividing by the number of available host CPU cores, so 100% means all host
CPUs are fully utilized. Alerts when the CPU usage of a container has been outside its
threshold for a configurable number of consecutive check runs (default: 5), suppressing
short spikes, and immediately when its memory usage is outside its threshold. For
Podman, use the podman-stats check instead.
Requires root or sudo."""

DB_FILENAME = 'linuxfabrik-monitoring-plugins-docker-stats.db'
DEFAULT_COUNT = (
    5  # measurements; if check runs once per minute, this is a 5 minute period
)
DEFAULT_CRIT_CPU = '90'  # %
DEFAULT_CRIT_MEM = '95'  # %
DEFAULT_NO_MATCH_SEVERITY = 'ok'
DEFAULT_TIMEOUT = 8
DEFAULT_WARN_CPU = '80'  # %
DEFAULT_WARN_MEM = '90'  # %


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(
        '--count',
        help=lib.args.help('--count') + ' Default: %(default)s',
        dest='COUNT',
        type=int,
        default=DEFAULT_COUNT,
    )

    parser.add_argument(
        '--critical-cpu',
        help='CRIT threshold for CPU usage, in percent. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        default=DEFAULT_CRIT_CPU,
        dest='CRIT_CPU',
    )

    parser.add_argument(
        '--critical-mem',
        help='CRIT threshold for memory usage, in percent. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        default=DEFAULT_CRIT_MEM,
        dest='CRIT_MEM',
    )

    parser.add_argument(
        '--full-name',
        help='Use the full container name, for example `traefik_traefik.2.1idw12p2yqp`. '
        'Without this flag, the name is shortened after the replica number.',
        dest='FULL_NAME',
        action='store_true',
        default=False,
    )

    parser.add_argument(
        '--ignore',
        help='Ignore containers whose name matches this Python regular expression. '
        'Matched against the full container name, even when the displayed name is '
        'shortened (see --full-name). '
        'Case-sensitive by default; use `(?i)` for case-insensitive matching. '
        'Can be specified multiple times. '
        'Example: `--ignore="^k8s_"` to skip Kubernetes pod infrastructure containers. '
        'Example: `--ignore="(?i)test"` (case-insensitive) to skip any container with '
        '"test" in its name. '
        'Default: %(default)s',
        dest='IGNORE',
        action='append',
        default=None,
    )

    parser.add_argument(
        '--match',
        help='Only check containers whose name matches this Python regular expression. '
        'Matched against the full container name, even when the displayed name is '
        'shortened (see --full-name). '
        'Case-sensitive by default; use `(?i)` for case-insensitive matching. '
        'Can be specified multiple times. '
        + lib.args.MATCH_IGNORE_PRECEDENCE
        + ' Example: `--match="^traefik$"` to pin the check to one specific container. '
        'Example: `--match="(?i)^web"` (case-insensitive) to check every web container. '
        'Default: %(default)s',
        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(
        '--warning-cpu',
        help='WARN threshold for CPU usage, in percent. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        default=DEFAULT_WARN_CPU,
        dest='WARN_CPU',
    )

    parser.add_argument(
        '--warning-mem',
        help='WARN threshold for memory usage, in percent. '
        'Supports Nagios ranges. '
        'Default: %(default)s',
        default=DEFAULT_WARN_MEM,
        dest='WARN_MEM',
    )

    args, _ = parser.parse_known_args()
    return args


def get_cpu_state(conn, container, count, warn, crit):
    """Return the state the last COUNT CPU samples of a container agree on.

    A single sample outside the threshold is a burst, COUNT of them in a row are a
    container that stays busy. Until COUNT samples are recorded, the container cannot
    alert on its CPU usage.
    """
    rows = lib.base.coe(
        lib.db_sqlite.select(
            conn,
            'SELECT cpu_usage FROM cpu_usage WHERE container = :container',
            {'container': container},
        )
    )
    samples = [row['cpu_usage'] for row in rows]
    if len(samples) < count:
        return STATE_OK
    states = [
        lib.base.get_state(sample, warn, crit, _operator='range') for sample in samples
    ]
    if all(item == STATE_CRIT for item in states):
        return STATE_CRIT
    if all(item in (STATE_CRIT, STATE_WARN) for item in states):
        return STATE_WARN
    return STATE_OK


def keep_container(name, match_patterns, ignore_patterns):
    """Return True if `name` should be kept by the --match / --ignore filter pair,
    False if it should be dropped. Include first, then exclude: a name passes if it
    matches any `match_patterns` entry (or if `match_patterns` is empty) AND does not
    match any `ignore_patterns` entry. Same semantics as the lib.args canonical
    --match / --ignore convention.
    """
    if match_patterns and not any(p.search(name) for p in match_patterns):
        return False
    return not any(p.search(name) for p in ignore_patterns)


def parse_percent(value):
    """Return the number out of a percentage the engine formatted, or None if it did
    not report one. `docker stats` prints `--` for every value of a container whose
    statistics could not be collected, which happens when the container is removed
    while the command runs or when the daemon does not answer within two seconds.
    """
    try:
        return float(value.replace('%', '').strip())
    except (AttributeError, ValueError):
        return None


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 = []

    # reject a threshold that is not a valid Nagios range, instead of grading every
    # container against it
    for param, threshold in (
        ('--critical-cpu', args.CRIT_CPU),
        ('--critical-mem', args.CRIT_MEM),
        ('--warning-cpu', args.WARN_CPU),
        ('--warning-mem', args.WARN_MEM),
    ):
        success, result = lib.base.match_range(0, threshold)
        if not success:
            lib.base.cu(f'Invalid Nagios range for {param} ("{threshold}"): {result}.')

    # compile --match and --ignore patterns (case-sensitive by default, matching
    # the lib.args convention; the user can opt into case-insensitive matching
    # with the inline `(?i)` flag)
    match_patterns = [
        lib.base.coe(p) for p in lib.txt.compile_regex(args.MATCH, '--match')
    ]
    ignore_patterns = [
        lib.base.coe(p) for p in lib.txt.compile_regex(args.IGNORE, '--ignore')
    ]

    # fetch data
    if args.TEST is None:
        started = lib.time.now(as_type='float')
        # get the number of host CPUs
        success, result = lib.container.run(
            ['docker', 'info', '--format', '{{json .}}'],
            args.TIMEOUT,
            started,
        )
        if not success:
            # --always-ok suppresses what would alert somebody. An UNKNOWN (no
            # permission, client missing) says the check is deployed wrong and stays.
            lib.base.oao(
                *result, always_ok=args.ALWAYS_OK and result[1] != STATE_UNKNOWN
            )
        stdout, stderr, retc = result
        if retc != 0:
            error = lib.container.get_engine_error(stderr, stdout)
            lib.base.oao(*error, always_ok=args.ALWAYS_OK and error[1] != STATE_UNKNOWN)
        try:
            host_cpus = int(json.loads(stdout).get('NCPU') or 0)
        except Exception:
            host_cpus = 0
        if not host_cpus:
            lib.base.cu(
                'The daemon did not report the number of host CPUs.'
                ' If you are using Podman, use the podman-stats check instead.'
            )

        # get the container statistics for all running containers
        success, result = lib.container.run(
            ['docker', 'stats', '--no-stream', '--format', '{{json .}}'],
            args.TIMEOUT,
            started,
        )
        if not success:
            lib.base.oao(
                *result, always_ok=args.ALWAYS_OK and result[1] != STATE_UNKNOWN
            )
        stdout, stderr, retc = result
        if retc != 0:
            error = lib.container.get_engine_error(stderr, stdout)
            lib.base.oao(*error, always_ok=args.ALWAYS_OK and error[1] != STATE_UNKNOWN)
    else:
        # do not call the command, put in test data
        host_cpus = 1
        stdout, stderr, retc = lib.lftest.test(args.TEST)
        if retc != 0:
            error = lib.container.get_engine_error(stderr, stdout)
            lib.base.oao(*error, always_ok=args.ALWAYS_OK and error[1] != STATE_UNKNOWN)

    # init some vars
    msg = ''
    state = STATE_OK
    perfdata = ''
    table_values = []

    # parse newline-delimited JSON output
    containers = []
    for line in stdout.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            containers.append(json.loads(line))
        except (json.JSONDecodeError, ValueError):
            continue

    # sort containers by name, then apply the --match / --ignore name filters (dropping
    # entries without a usable name)
    containers.sort(key=lambda c: c.get('Name', ''))
    # a container whose statistics never arrived is reported without a name, which the
    # engine prints as the same placeholder it uses for every unknown value
    containers = [c for c in containers if c.get('Name', '') not in ('', '--')]
    # every name a running container is known under, whatever this run filters: another
    # service on the same host may filter differently or use --full-name, and must not
    # lose its history because this run looks elsewhere
    running_names = {c['Name'] for c in containers}
    running_names |= {lib.container.strip_task_id(name) for name in running_names}
    containers = [
        c
        for c in containers
        if keep_container(c.get('Name', ''), match_patterns, ignore_patterns)
    ]

    # analyze data
    # create the db table holding the CPU trend the --count evaluation reads
    definition = """
                container TEXT NOT NULL,
                cpu_usage REAL NOT NULL,
                timestamp REAL NOT NULL
        """
    # in test mode the trend lives in memory only, so a fixture neither reads the
    # history an earlier test run left behind nor leaves one for the next
    if args.TEST is None:
        conn = lib.base.coe(lib.db_sqlite.connect(filename=DB_FILENAME))
    else:
        conn = lib.base.coe(lib.db_sqlite.connect(in_memory=True))
    lib.base.coe(lib.db_sqlite.create_table(conn, definition, table='cpu_usage'))
    lib.base.coe(lib.db_sqlite.create_index(conn, 'container', table='cpu_usage'))
    # forget the history of containers that no longer exist, otherwise the state file
    # grows with every short-lived container (CI jobs, one-off runs)
    lib.base.coe(
        lib.db_sqlite.forget_sensors(
            conn, 'container', running_names, table='cpu_usage'
        )
    )

    # grade every container
    for container in containers:
        name = container.get('Name', '')
        # https://github.com/Linuxfabrik/monitoring-plugins/issues/586
        if not args.FULL_NAME:
            name = lib.container.strip_task_id(name)
        # a container name may carry characters that have no place in a metric name
        label = re.sub(r'\W+', '_', name)

        # the engine states the CPU percentage relative to a single core, so it goes
        # up to 100% per core; divide by the cores of the host to get its share
        cpu_percent = parse_percent(container.get('CPUPerc'))
        cpu_usage = None if cpu_percent is None else round(cpu_percent / host_cpus, 1)
        mem_percent = parse_percent(container.get('MemPerc'))
        mem_usage = None if mem_percent is None else round(mem_percent, 1)

        # per-container perfdata for long-term trending of individual workloads
        if cpu_usage is not None:
            perfdata += lib.base.get_perfdata(
                f'{label}_cpu_usage',
                cpu_usage,
                uom='%',
                warn=args.WARN_CPU,
                crit=args.CRIT_CPU,
                _min=0,
                _max=100,
            )
        if mem_usage is not None:
            perfdata += lib.base.get_perfdata(
                f'{label}_mem_usage',
                mem_usage,
                uom='%',
                warn=args.WARN_MEM,
                crit=args.CRIT_MEM,
                _min=0,
                _max=100,
            )

        cpu_state = mem_state = STATE_OK
        if cpu_usage is not None:
            # save trend data to local sqlite database, limited to "count" rows max.
            lib.base.coe(
                lib.db_sqlite.insert(
                    conn,
                    {
                        'container': name,
                        'cpu_usage': cpu_usage,
                        'timestamp': lib.time.now(as_type='float'),
                    },
                    table='cpu_usage',
                ),
            )
            # keep "count" rows per container; trimming the table to a total instead
            # lets a container that is sampled more often evict the others' history
            lib.base.coe(
                lib.db_sqlite.cut_per_sensor(
                    conn, sensorcol='container', _max=args.COUNT, table='cpu_usage'
                )
            )

            # alert when the cpu usage of the container stayed outside its threshold
            cpu_state = get_cpu_state(
                conn, name, args.COUNT, args.WARN_CPU, args.CRIT_CPU
            )
            if cpu_state != STATE_OK:
                # build the message
                msg += f'"{name}" cpu {cpu_usage}% {lib.base.state2str(cpu_state)}, '
            state = lib.base.get_worst(cpu_state, state)

        # alert when container mem_usage is exceeded
        if mem_usage is not None:
            mem_state = lib.base.get_state(
                mem_usage, args.WARN_MEM, args.CRIT_MEM, _operator='range'
            )
            if mem_state != STATE_OK:
                msg += f'"{name}" memory {mem_usage}% {lib.base.state2str(mem_state)}, '
            state = lib.base.get_worst(mem_state, state)

        table_values.append(
            {
                'name': name,
                'cpu_usage': '-'
                if cpu_usage is None
                else f'{cpu_usage}{lib.base.state2str(cpu_state, prefix=" ")}',
                'mem_usage': '-'
                if mem_usage is None
                else f'{mem_usage}{lib.base.state2str(mem_state, prefix=" ")}',
            }
        )

    # we don't need the database any more: save data and close connection
    lib.db_sqlite.commit(conn)
    lib.db_sqlite.close(conn)

    # nothing left after applying the --match / --ignore filters (or no running
    # containers); report the configured no-match severity
    if not table_values:
        lib.base.oao(
            'No containers to check.',
            lib.base.str2state(args.NO_MATCH_SEVERITY),
            always_ok=args.ALWAYS_OK,
        )

    # build the message
    perfdata += lib.base.get_perfdata(
        'containers_running',
        len(table_values),
        _min=0,
    )
    perfdata += lib.base.get_perfdata(
        'cpu',
        host_cpus,
        _min=0,
    )

    if state == STATE_OK:
        checked = len(table_values)
        msg = (
            f'Everything is ok. {checked} '
            f'{lib.txt.pluralize("container", checked)} checked.\n\n'
        )
    else:
        msg = msg[:-2] + '\n\n'

    # build table output
    if len(table_values) > 0:
        msg += lib.base.get_table(
            table_values,
            ['name', 'cpu_usage', 'mem_usage'],
            header=['Container', 'CPU %', 'Mem % '],
        )

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