#!/usr/bin/env python3
"""Checkmk special agent for Talos Linux.

Talks the native Talos API (mTLS gRPC on tcp/50000) with a read-only
``os:reader`` credential and emits Checkmk agent sections on stdout.

    agent_talos --talosconfig ~/etc/talos/talosconfig --endpoint 10.5.0.156

Exit codes:

    0   the node was reached and at least one section was produced (partial
        data is not failure -- individual failures are reported in-band,
        through the talos_agent_error section)
    1   the node could not be reached, or the configuration is unusable.
        Note that the certificate section is still emitted in that case, since
        it is derived from the talosconfig alone -- an expired certificate is
        exactly when everything else goes dark.
    2   bad invocation
"""

from __future__ import annotations

import argparse
import os
import sys

# Vendored pure-Python dependencies (h2, hpack, hyperframe) ship beside this
# file so the plugin does not depend on what the Checkmk site Python happens
# to include. A locally installed copy still wins, which keeps development
# environments simple.
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(_HERE, "_vendor"))
sys.path.insert(0, _HERE)

from _talos import __version__, collect, config as config_mod  # noqa: E402
from _talos.client import TalosClient  # noqa: E402
from _talos.transport import TransportError  # noqa: E402


#: Exit code for a bad invocation, matching what argparse itself uses. Bare
#: SystemExit("message") would print the text but exit 1, which is the code
#: reserved for "ran, but could not reach the node".
EXIT_USAGE = 2


def _usage_error(message: str) -> SystemExit:
    sys.stderr.write(f"agent_talos: {message}\n")
    return SystemExit(EXIT_USAGE)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="agent_talos",
        description="Checkmk special agent for Talos Linux.",
    )
    parser.add_argument(
        "--talosconfig", required=True, metavar="PATH",
        help="path to a talosconfig holding os:reader credentials, or '-' to "
             "read it from stdin (used when the config is uploaded through "
             "the Setup GUI, which keeps the private key out of the process "
             "table)")
    parser.add_argument(
        "--endpoint", default=None,
        help="address to connect to (usually the node itself). Omit it to use "
             "the endpoints recorded in the talosconfig -- see --node for when "
             "that is unambiguous")
    parser.add_argument(
        "--node", default=None,
        help="node to query, if different from --endpoint (apid proxying)")
    parser.add_argument("--port", type=int, default=50000)
    parser.add_argument("--timeout", type=float, default=10.0)
    parser.add_argument(
        "--context", default=None,
        help="talosconfig context; defaults to the file's current context")
    parser.add_argument(
        "--no-verify-hostname", action="store_true",
        help="skip TLS hostname verification (the CA is still enforced); "
             "needed when connecting via a name absent from the node's SANs")
    parser.add_argument(
        "--sections", default=None,
        help=f"comma-separated subset of: {','.join(collect.ALL_SECTIONS)}")
    parser.add_argument(
        "--no-optional", action="store_true",
        help=f"skip the chatty sections: {','.join(collect.OPTIONAL_SECTIONS)}")
    parser.add_argument("--debug", action="store_true",
                        help="let exceptions escape, for troubleshooting")
    parser.add_argument("--version", action="version",
                        version=f"agent_talos {__version__}")
    return parser


def resolve_endpoints(args: argparse.Namespace, cfg) -> list[str]:
    """Decide which API endpoints to try, in order.

    A talosconfig carries the cluster's endpoints, so a config uploaded once
    per cluster already knows how to reach it. Using them is only safe when
    the *target* is pinned, though:

    * ``--endpoint`` given -> use exactly that.
    * no endpoint but ``--node`` given -> try every endpoint from the config.
      apid proxies to the requested node, so whichever endpoint answers, the
      data comes from the intended machine.
    * no endpoint, no node, one endpoint in the config -> unambiguous, use it.
    * no endpoint, no node, several endpoints -> refuse. Connecting to an
      arbitrary control plane node and filing its CPU and disks under a
      different host is worse than not running at all.
    """
    if args.endpoint:
        return [args.endpoint]

    endpoints = list(cfg.endpoints)
    if not endpoints:
        raise _usage_error("no --endpoint given and the talosconfig lists none")
    if args.node or len(endpoints) == 1:
        return endpoints
    raise _usage_error(
        "the talosconfig lists several endpoints "
        f"({', '.join(endpoints)}) and neither --endpoint nor --node was "
        "given. Pass --endpoint to choose one, or --node to say which node "
        "the data should come from."
    )


def resolve_sections(args: argparse.Namespace) -> set[str] | None:
    if args.sections:
        requested = {s.strip() for s in args.sections.split(",") if s.strip()}
        unknown = requested - set(collect.ALL_SECTIONS)
        if unknown:
            raise _usage_error(
                f"unknown section(s): {', '.join(sorted(unknown))}")
        return requested
    if args.no_optional:
        return set(collect.ALL_SECTIONS) - set(collect.OPTIONAL_SECTIONS)
    return None


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    enabled = resolve_sections(args)

    # The config is needed even if the connection fails: the certificate
    # section is produced from it alone, and an expired certificate is exactly
    # the case where everything else goes dark.
    try:
        if args.talosconfig == "-":
            cfg = config_mod.load_stdin(context=args.context)
        else:
            cfg = config_mod.load(args.talosconfig, context=args.context)
    except config_mod.ConfigError as exc:
        if args.debug:
            raise
        sys.stderr.write(f"agent_talos: {exc}\n")
        return 1

    client = None
    connection_error = None
    for endpoint in resolve_endpoints(args, cfg):
        try:
            client = TalosClient.from_config(
                cfg,
                endpoint,
                port=args.port,
                timeout=args.timeout,
                node=args.node,
                verify_hostname=not args.no_verify_hostname,
            )
            connection_error = None
            break
        except (config_mod.ConfigError, TransportError) as exc:
            if args.debug:
                raise
            # Keep the first failure: with several endpoints it is usually the
            # most informative, and later ones are often the same story.
            connection_error = connection_error or exc

    try:
        result = collect.collect(
            collect.Sources(config=cfg, client=client), enabled, debug=args.debug)
    finally:
        if client is not None:
            client.close()

    sys.stdout.write(result.output)

    if connection_error is not None:
        sys.stderr.write(f"agent_talos: {connection_error}\n")
        return 1
    for name, message in sorted(result.failures.items()):
        sys.stderr.write(f"agent_talos: section {name}: {message}\n")

    return 0 if result.any_data else 1


if __name__ == "__main__":
    sys.exit(main())
