#!/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 re
import sys

import lib.args
import lib.base
import lib.lftest
import lib.shell
from lib.globals import STATE_OK, STATE_UNKNOWN

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

DESCRIPTION = """Counts the sessions that `w` (Linux) or `query user` (Windows) lists,
by type: tty (console) and pts (SSH on Linux, RDP on Windows). On Windows, also counts
disconnected sessions (closed connections without logging out). Alerts when the number
of sessions of a type reaches its threshold; by default, every console session warns."""

DEFAULT_WARN_PTS = 20
DEFAULT_WARN_DISC = 1
DEFAULT_WARN_TTY = 1
DEFAULT_CRIT_PTS = None
DEFAULT_CRIT_DISC = None
DEFAULT_CRIT_TTY = None

# the WHAT column of the `manager` session, e.g. `/usr/lib/systemd/systemd --user`
SYSTEMD_USER_MANAGER = re.compile(r'(?:^|[\s/])systemd --user$')


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(
        '-c',
        '--critical',
        default=[DEFAULT_CRIT_TTY, DEFAULT_CRIT_PTS, DEFAULT_CRIT_DISC],
        dest='CRIT',
        help='Threshold for the number of tty and pts sessions, '
        'in the format "tty,pts". '
        'Both values are required, `None` leaves a threshold out. '
        'On Windows, you can additionally specify a threshold for disconnected '
        'sessions in the format "tty,pts,disc". '
        'Alerts when the number reaches the threshold. '
        'Example: `--critical=3,10`. '
        'Default: %(default)s',
        type=lib.args.csv,
    )

    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(
        '-w',
        '--warning',
        default=[DEFAULT_WARN_TTY, DEFAULT_WARN_PTS, DEFAULT_WARN_DISC],
        dest='WARN',
        help='Threshold for the number of tty and pts sessions, '
        'in the format "tty,pts". '
        'Both values are required, `None` leaves a threshold out. '
        'On Windows, you can additionally specify a threshold for disconnected '
        'sessions in the format "tty,pts,disc". '
        'Alerts when the number reaches the threshold. '
        'Example: `--warning=1,5`. '
        'Default: %(default)s',
        type=lib.args.csv,
    )

    args, _ = parser.parse_known_args()
    return args


def parse_linux_output(s):
    """Parse the output of `w` on Linux.

    Uses the header line to determine the TTY column position, so it works
    regardless of whether FROM is present, and regardless of column widths
    (which vary across distros and versions).

    procps-ng before 4.0.5, built against logind, lists every logind session,
    including the `manager` session systemd 256 and newer opens for the per-user
    `systemd --user`. That line has no TTY and is nobody logging in, so it is
    left out of the count; it would make one SSH login read `PTS: 2`. procps-ng
    4.0.5 (commit 734930e4) shows only sessions of class `user*`. The WHAT column
    is only that long when `COLUMNS` is unset, see `main()`. Verified against
    procps-ng 4.0.4 / systemd 257 on Rocky 10 (lists it), Debian 13 and Ubuntu
    26.04 (procps 4.0.4 reading utmp, do not list it) and procps-ng 4.0.6 /
    systemd 259 on Fedora 44 (does not list it).
    """
    # replace pipes in output, otherwise we will get problems with perfdata,
    # and ignore the first line of w's output
    s = s.strip().replace('|', '!').splitlines()[1:]
    if not s:
        return s, 0, 0

    header = s[0]
    tty_start = header.find('TTY')
    if tty_start < 0:
        return s, 0, 0

    # find end of TTY column: start of the next column header after TTY
    after_tty = header[tty_start + 3 :]
    tty_end = tty_start + 3 + (len(after_tty) - len(after_tty.lstrip()))

    count_tty, count_pts = 0, 0
    for line in s[1:]:
        tty_value = line[tty_start:tty_end].strip() if len(line) > tty_start else ''
        if tty_value.startswith('tty') or tty_value.startswith(':'):
            # tty = local terminal, ":0" = local X display
            count_tty += 1
        elif not tty_value.startswith('pts/') and SYSTEMD_USER_MANAGER.search(
            line.rstrip()
        ):
            # The per-user service manager, not a login. Its TTY column is empty,
            # but the cut can catch the first digit of LOGIN@, which `w` on Rocky 10
            # prints one position left of its header.
            continue
        else:
            # pts, empty (SSH without TTY) or anything else
            count_pts += 1

    return s, count_tty, count_pts


def parse_windows_output(s):
    """Parse the output of `query user` on Windows.

    Cuts the session name at the columns the header gives it, the way
    `parse_linux_output()` reads the TTY column. A disconnected session has no
    session name, so splitting the line at whitespace moves its ID into that
    column and the session is counted nowhere. The header is translated with the
    display language, its layout is not. `query user` prints nothing on stdout
    when nobody is logged in. Verified against Windows Server 2025.
    """
    # Only trailing whitespace goes: the header starts with the space the data
    # lines keep for the ">" that marks the calling session, and stripping it
    # would shift every header position by one.
    s = [line for line in s.rstrip().splitlines() if line.strip()]
    if not s:
        return s, 0, 0, 0

    columns = [m.start() for m in re.finditer(r'\S+', s[0])]
    if len(columns) < 3:
        return s, 0, 0, 0
    name_start, name_end = columns[1], columns[2]

    count_tty, count_pts, count_disc = 0, 0, 0
    for line in s[1:]:
        value = line[name_start:name_end].strip()
        if value == 'console':
            count_tty += 1
        elif 'rdp-' in value:
            count_pts += 1
        elif value == '':
            count_disc += 1

    return s, count_tty, count_pts, count_disc


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)

    try:
        WARN_TTY = lib.args.int_or_none(args.WARN[0])
        WARN_PTS = lib.args.int_or_none(args.WARN[1])
        CRIT_TTY = lib.args.int_or_none(args.CRIT[0])
        CRIT_PTS = lib.args.int_or_none(args.CRIT[1])
    except (IndexError, ValueError):
        lib.base.cu(
            '--warning and --critical take two integers, "tty,pts" (on Windows'
            ' optionally three, "tty,pts,disc"). Use `None` to leave a threshold out,'
            ' for example `--critical=None,50`',
            traceback=False,
        )

    # A Windows fixture is read like the output of a Windows host, thresholds and
    # message included, so the tests reach the disconnected sessions too.
    windows = (
        lib.base.WINDOWS
        if args.TEST is None
        else any('windows' in s for s in args.TEST)
    )
    WARN_DISC, CRIT_DISC = None, None
    if windows:
        # "tty,pts" is a valid threshold on Windows too, the third value is optional
        WARN_DISC = lib.args.int_or_none(
            args.WARN[2] if len(args.WARN) > 2 else DEFAULT_WARN_DISC
        )
        CRIT_DISC = lib.args.int_or_none(
            args.CRIT[2] if len(args.CRIT) > 2 else DEFAULT_CRIT_DISC
        )

    # fetch data
    if args.TEST is None:
        if windows:
            cmd = ['query', 'user']
            stdout, _, _ = lib.base.coe(lib.shell.shell_exec(cmd))
            # Could not find any documentation in the return codes of 'query user'
            # (%ERRORLEVEL% is always 1).
            lines, count_tty, count_pts, count_disc = parse_windows_output(stdout)
        else:
            cmd = ['/usr/bin/w']
            # Without a terminal, `w` cuts the WHAT column to `COLUMNS` if that is
            # set, else to 512 characters. Unset, `systemd --user` stays readable.
            stdout, stderr, retc = lib.base.coe(
                lib.shell.shell_exec(cmd, env={'COLUMNS': None})
            )
            if stderr or retc != 0:
                lib.base.cu(stderr)
            lines, count_tty, count_pts = parse_linux_output(stdout)
    else:
        # do not call the command, put in test data
        stdout, _, _ = lib.lftest.test(args.TEST)
        if windows:
            lines, count_tty, count_pts, count_disc = parse_windows_output(stdout)
        else:
            lines, count_tty, count_pts = parse_linux_output(stdout)

    # init some vars
    msg = ''
    perfdata = ''
    state = STATE_OK

    # analyze data
    # both parsers keep the header, so a header alone is nobody
    if len(lines) <= 1:
        msg = 'No one is logged in.'
    else:
        tty_state = lib.base.get_state(count_tty, WARN_TTY, CRIT_TTY)
        state = lib.base.get_worst(state, tty_state)

        # build the message
        msg += f'TTY: {count_tty}' + lib.base.state2str(tty_state, prefix=' ')

        pts_state = lib.base.get_state(count_pts, WARN_PTS, CRIT_PTS)
        state = lib.base.get_worst(state, pts_state)
        msg += f', PTS: {count_pts}' + lib.base.state2str(pts_state, prefix=' ')

        if windows:
            disc_state = lib.base.get_state(count_disc, WARN_DISC, CRIT_DISC)
            state = lib.base.get_worst(state, disc_state)
            msg += f', Disconnected: {count_disc}' + lib.base.state2str(
                disc_state, prefix=' '
            )

        msg += '\n\n' + '\n'.join(lines)

    # written whether or not anybody is logged in, like tty and pts, so the graph
    # has no gap
    if windows:
        perfdata += lib.base.get_perfdata(
            'disc',
            count_disc,
            uom=None,
            warn=WARN_DISC,
            crit=CRIT_DISC,
            _min=0,
            _max=None,
        )
    perfdata += lib.base.get_perfdata(
        'tty',
        count_tty,
        uom=None,
        warn=WARN_TTY,
        crit=CRIT_TTY,
        _min=0,
        _max=None,
    )
    perfdata += lib.base.get_perfdata(
        'pts',
        count_pts,
        uom=None,
        warn=WARN_PTS,
        crit=CRIT_PTS,
        _min=0,
        _max=None,
    )

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