#!/usr/lib64/linuxfabrik-monitoring-plugins/venv/bin/python
# -*- 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.db_sqlite
import lib.lftest
import lib.shell
import lib.txt
from lib.globals import STATE_CRIT, STATE_OK, STATE_UNKNOWN, STATE_WARN

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

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. CPU alerts only
trigger after the threshold has been exceeded for a configurable number of consecutive
check runs (default: 5), suppressing short spikes. Memory alerts trigger immediately.
Uses a local SQLite database for CPU trend tracking across runs. For Podman, use the
podman-stats check instead.
Requires root or sudo."""

DEFAULT_COUNT = (
    5  # measurements; if check runs once per minute, this is a 5 minute period
)
DEFAULT_NO_MATCH_SEVERITY = 'ok'
DEFAULT_WARN_CPU = 80  # %
DEFAULT_CRIT_CPU = 90  # %
DEFAULT_WARN_MEM = 90  # %
DEFAULT_CRIT_MEM = 95  # %


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. Default: >= %(default)s',
        default=DEFAULT_CRIT_CPU,
        dest='CRIT_CPU',
    )

    parser.add_argument(
        '--critical-mem',
        help='CRIT threshold for memory usage in percent. 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(
        '--warning-cpu',
        help='WARN threshold for CPU usage in percent. Default: >= %(default)s',
        default=DEFAULT_WARN_CPU,
        dest='WARN_CPU',
    )

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

    args, _ = parser.parse_known_args()
    return args


def get_engine_error(stderr, stdout=''):
    """Return `(message, state)` for a command that could not reach the container
    engine. A refused permission is a problem of how this check is deployed, not of
    the engine: the engine answers other callers just fine, this one is only not
    allowed to ask, so the check cannot say anything and reports UNKNOWN. Everything
    else, a socket that is not there or an engine that does not answer, is the
    outage this check exists to report.
    """
    text = f'{stderr}\n{stdout}'.strip()
    if 'permission denied' in text.lower():
        return (
            'No permission to talk to the container engine, so nothing can be said'
            ' about it. Run the check as root, or deploy the sudoers file that ships'
            f' with the plugins.\n{text}',
            STATE_UNKNOWN,
        )
    return (text, STATE_CRIT)


def get_cpu_from_db(conn, container, threshold):
    """Return the number of rows where cpu_usage >= threshold for a container."""
    result = lib.base.coe(
        lib.db_sqlite.select(
            conn,
            """
        SELECT count(*) as cnt
        FROM cpu
        WHERE container = :container and cpu_usage >= :threshold
        """,
            {'container': container, 'threshold': threshold},
            fetchone=True,
        )
    )
    return int(result['cnt'])


# A swarm task container is named `<service>.<slot or node id>.<task id>`, and the
# task id is 25 base36 characters (swarmkit identity.NewID). Only that suffix is cut
# off, so a container an admin named `backup.daily` keeps the name they gave it, and
# two such containers do not collapse into one row.
TASK_ID_REGEX = re.compile(r'\.[0-9a-z]{25}$')


def shorten(name):
    """
    >>> shorten('traefik_traefik.2.1idw12p2yqpxutlzkcwign4at')
    traefik_traefik.2
    >>> shorten('backup.daily')
    backup.daily
    """
    return TASK_ID_REGEX.sub('', name)


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

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

    # create the db tables
    definition = """
                container TEXT NOT NULL,
                cpu_usage REAL NOT NULL
        """
    conn = lib.base.coe(
        lib.db_sqlite.connect(
            filename='linuxfabrik-monitoring-plugins-docker-stats.db'
        ),
    )
    lib.base.coe(lib.db_sqlite.create_table(conn, definition, table='cpu'))
    lib.db_sqlite.create_index(conn, 'container', table='cpu')

    # fetch data
    if args.TEST is None:
        # get the number of host CPUs
        stdout, stderr, retc = lib.base.coe(
            lib.shell.shell_exec(['docker', 'info', '--format', '{{json .}}']),
        )
        if retc != 0:
            lib.db_sqlite.close(conn)
            lib.base.oao(*get_engine_error(stderr, stdout))
        try:
            host_cpus = int(json.loads(stdout).get('NCPU') or 0)
        except Exception:
            host_cpus = 0
        if not host_cpus:
            lib.db_sqlite.close(conn)
            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
        stdout, stderr, retc = lib.base.coe(
            lib.shell.shell_exec(
                ['docker', 'stats', '--no-stream', '--format', '{{json .}}']
            ),
        )
    else:
        # do not call the command, put in test data
        host_cpus = 1
        stdout, stderr, retc = lib.lftest.test(args.TEST)

    if retc != 0:
        lib.db_sqlite.close(conn)
        lib.base.oao(*get_engine_error(stderr, stdout))

    # 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), so the SQLite trend cut below is sized to the
    # number of tracked containers
    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', '')
        and c.get('Name') != '--'
        and keep_container(c.get('Name', ''), match_patterns, ignore_patterns)
    ]

    # analyze data
    for container in containers:
        name = container.get('Name', '')
        # https://github.com/Linuxfabrik/monitoring-plugins/issues/586
        if not args.FULL_NAME:
            name = shorten(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}, table='cpu'
                ),
            )
            lib.base.coe(
                lib.db_sqlite.cut(conn, _max=args.COUNT * len(containers), table='cpu')
            )

            # alert when container cpu_usage is exceeded
            # my container state is not ok, if in every of my historic rows the cpu
            # value is above the threshold
            if get_cpu_from_db(conn, name, args.CRIT_CPU) >= args.COUNT:
                cpu_state = STATE_CRIT
            elif get_cpu_from_db(conn, name, args.WARN_CPU) >= args.COUNT:
                cpu_state = STATE_WARN
            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)
            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 perfdata
    perfdata += lib.base.get_perfdata(
        'containers_running',
        len(table_values),
        _min=0,
    )
    perfdata += lib.base.get_perfdata(
        'cpu',
        host_cpus,
        _min=0,
    )

    # create output
    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'
    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()
