#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-only
# Copyright (C) 2026 Sudhir Chauhan

"""Independent Checkmk special agent for KubeVirt.

This agent intentionally talks only to the KubeVirt API resources and to a
Prometheus endpoint containing KubeVirt metrics. It does not import or call
agent_kube and does not collect generic Kubernetes pods/nodes/deployments.
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import sys
from urllib.parse import urlsplit
from collections import defaultdict
from collections.abc import Iterable, Mapping, Sequence
from typing import Any

import requests
import urllib3

from cmk.password_store.v1_unstable import parser_add_secret_option, resolve_secret_option

LOG = logging.getLogger("agent_kubevirt")

# Kubernetes annotation values may contain multiline JSON or secrets. Only
# label-compatible values are eligible for Checkmk host labels.
LABEL_VALUE = re.compile(r"(?:[A-Za-z0-9](?:[-A-Za-z0-9_.]{0,61}[A-Za-z0-9])?)?")


def valid_annotation_value(value: object) -> bool:
    return isinstance(value, str) and LABEL_VALUE.fullmatch(value) is not None


def validate_prometheus_token_transport(endpoint: str, verify_certificate: bool) -> None:
    """Reject insecure destinations for a separately configured Prometheus credential."""
    parsed = urlsplit(endpoint)
    loopback = parsed.hostname in {"localhost", "127.0.0.1", "::1"}
    if parsed.scheme == "https" and (verify_certificate or loopback):
        return
    if parsed.scheme == "http" and loopback:
        return
    raise ValueError("Prometheus bearer authentication requires verified HTTPS (or a loopback endpoint)")

KUBEVIRT_QUERIES = {
    "cpu": 'sum(rate(kubevirt_vmi_cpu_usage_seconds_total[5m])) by (exported_namespace, name)',
    "memory": 'sum(kubevirt_vmi_memory_resident_bytes) by (exported_namespace, name)',
    "net_rx": 'sum(rate(kubevirt_vmi_network_receive_bytes_total[5m])) by (exported_namespace, name, interface)',
    "net_tx": 'sum(rate(kubevirt_vmi_network_transmit_bytes_total[5m])) by (exported_namespace, name, interface)',
    "disk_read": 'sum(rate(kubevirt_vmi_storage_read_traffic_bytes_total[5m])) by (exported_namespace, name, drive)',
    "disk_write": 'sum(rate(kubevirt_vmi_storage_write_traffic_bytes_total[5m])) by (exported_namespace, name, drive)',
    "version": 'kubevirt_info',
    "component_up": 'up{job=~".*kubevirt.*|kubevirt.*"}',
    "controller_leading": 'max(kubevirt_virt_controller_leading_status)',
    "controller_ready": 'max(kubevirt_virt_controller_ready_status)',
    "operator_leading": 'max(kubevirt_virt_operator_leading_status)',
    "operator_ready": 'max(kubevirt_virt_operator_ready_status)',
    "workqueue_depth": 'sum(kubevirt_workqueue_depth) by (name, queue, workqueue)',
    "workqueue_adds": 'sum(rate(kubevirt_workqueue_adds_total[5m])) by (name, queue, workqueue)',
    "workqueue_retries": 'sum(rate(kubevirt_workqueue_retries_total[5m])) by (name, queue, workqueue)',
    "workqueue_unfinished": 'sum(kubevirt_workqueue_unfinished_work_seconds) by (name, queue, workqueue)',
    # Current upstream names in the lab output:
    "migrations_pending": 'sum(kubevirt_vmi_migrations_in_pending_phase)',
    "migrations_running": 'sum(kubevirt_vmi_migrations_in_running_phase)',
    "migrations_scheduling": 'sum(kubevirt_vmi_migrations_in_scheduling_phase)',
    "migrations_unset": 'sum(kubevirt_vmi_migrations_in_unset_phase)',
    "vmi_phase_count": 'sum(kubevirt_vmi_phase_count) by (phase)',
    "vm_running": 'kubevirt_vm_running_status_last_transition_timestamp_seconds',
    "vm_starting": 'kubevirt_vm_starting_status_last_transition_timestamp_seconds',
    "vm_non_running": 'kubevirt_vm_non_running_status_last_transition_timestamp_seconds',
    "vm_migrating": 'kubevirt_vm_migrating_status_last_transition_timestamp_seconds',
    "vm_error": 'kubevirt_vm_error_status_last_transition_timestamp_seconds',
}


def parse_args(argv: list[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Monitor KubeVirt virtual machines and KubeVirt platform metrics")
    p.add_argument("--debug", action="store_true")
    p.add_argument("-v", "--verbose", action="count", default=0)
    p.add_argument("--cluster", required=True)
    p.add_argument("--kubernetes-cluster-hostname", required=True)
    parser_add_secret_option(p, long="--token", help="Kubernetes bearer token", required=True)
    p.add_argument("--api-server-endpoint", required=True)
    p.add_argument("--api-server-proxy", default="FROM_ENVIRONMENT")
    p.add_argument("--verify-cert-api", action="store_true")
    p.add_argument("--k8s-api-connect-timeout", type=int, default=10)
    p.add_argument("--k8s-api-read-timeout", type=int, default=12)
    ns = p.add_mutually_exclusive_group()
    ns.add_argument("--namespace-include-patterns", "-n", action="append", default=[])
    ns.add_argument("--namespace-exclude-patterns", action="append", default=[])
    p.add_argument("--include-annotations-as-host-labels", action="store_true")
    p.add_argument("--include-matching-annotations-as-host-labels")
    p.add_argument("--prometheus-endpoint")
    parser_add_secret_option(p, long="--prometheus-token", help="Optional, separate Prometheus bearer token", required=False)
    p.add_argument("--usage-proxy", default="FROM_ENVIRONMENT")
    p.add_argument("--usage-verify-cert", action="store_true")
    p.add_argument("--usage-connect-timeout", type=int, default=10)
    p.add_argument("--usage-read-timeout", type=int, default=12)
    return p.parse_args(argv)


def setup_logging(verbosity: int) -> None:
    if verbosity >= 3:
        level = logging.DEBUG
    elif verbosity == 2:
        level = logging.INFO
    elif verbosity == 1:
        level = logging.WARNING
    else:
        logging.disable(logging.CRITICAL)
        level = logging.CRITICAL
    logging.basicConfig(level=level, format="%(asctime)s %(levelname)s %(message)s")


def _resolve_token(args: argparse.Namespace) -> str:
    # parser_add_secret_option creates both --token and --token-id.
    # resolve_secret_option always returns a Secret wrapper; reveal() is needed
    # before using it in the Authorization header.
    return resolve_secret_option(args, "token").reveal()


def _resolve_prometheus_token(args: argparse.Namespace) -> str | None:
    if not (getattr(args, "prometheus_token", None) or getattr(args, "prometheus_token_id", None)):
        return None
    return resolve_secret_option(args, "prometheus_token").reveal()


def _proxy(value: str) -> Mapping[str, str] | None:
    if value == "FROM_ENVIRONMENT":
        return None
    if value == "NO_PROXY":
        return {"http": "", "https": ""}
    return {"http": value, "https": value}


def _get_json(session: requests.Session, url: str, **kwargs: Any) -> Mapping[str, Any]:
    # Reject redirects rather than forwarding credentials or silently moving to
    # a different endpoint.
    response = session.get(url, allow_redirects=False, **kwargs)
    if 300 <= response.status_code < 400:
        raise RuntimeError(f"HTTP redirect refused for monitoring endpoint: {url}")
    response.raise_for_status()
    data = response.json()
    if not isinstance(data, Mapping):
        raise RuntimeError(f"Unexpected JSON payload from {url}")
    return data


def _list_items(session: requests.Session, endpoint: str, path: str, **kwargs: Any) -> list[Mapping[str, Any]]:
    data = _get_json(session, endpoint.rstrip("/") + path, **kwargs)
    return [item for item in data.get("items", []) if isinstance(item, Mapping)]


def _volume_source(volume: Mapping[str, Any]) -> str | None:
    for key in volume.keys():
        if key != "name":
            return str(key)
    return None


def _parse_quantity(value: Any) -> float | None:
    if value is None:
        return None
    if isinstance(value, (int, float)):
        return float(value)
    s = str(value)
    units = [("Ki", 1024), ("Mi", 1024**2), ("Gi", 1024**3), ("Ti", 1024**4), ("m", 0.001)]
    for suffix, factor in units:
        if s.endswith(suffix):
            return float(s[:-len(suffix)]) * factor
    return float(s)


def _vm_section_data(vm: Mapping[str, Any], vmi: Mapping[str, Any] | None, cluster: str, host: str, import_annotations: bool, annotation_pattern: str | None) -> dict[str, Mapping[str, Any]]:
    meta = vm.get("metadata", {}) or {}
    spec = vm.get("spec", {}) or {}
    status = vm.get("status", {}) or {}
    template_spec = ((spec.get("template") or {}).get("spec") or {})
    domain = template_spec.get("domain") or {}
    devices = domain.get("devices") or {}
    resources = domain.get("resources") or {}
    requests = resources.get("requests") or {}
    cpu = domain.get("cpu") or {}
    vmi_status = (vmi or {}).get("status", {}) or {}

    raw_annotations = meta.get("annotations") or {}
    if import_annotations:
        annotations = {k: v for k, v in raw_annotations.items() if valid_annotation_value(v)}
    elif annotation_pattern:
        regex = re.compile(annotation_pattern)
        annotations = {k: v for k, v in raw_annotations.items() if regex.search(k) and valid_annotation_value(v)}
    else:
        annotations = {}

    disks = []
    for disk in devices.get("disks", []) or []:
        if isinstance(disk, Mapping):
            disks.append({"name": disk.get("name"), "bus": (disk.get("disk") or {}).get("bus") or disk.get("bus")})
    volumes = []
    for volume in template_spec.get("volumes", []) or []:
        if isinstance(volume, Mapping):
            volumes.append({"name": volume.get("name"), "source": _volume_source(volume)})
    interfaces = []
    for iface in vmi_status.get("interfaces", []) or []:
        if isinstance(iface, Mapping):
            interfaces.append({"name": iface.get("name"), "ip_address": iface.get("ipAddress"), "mac": iface.get("mac"), "info_source": iface.get("infoSource")})

    return {
        "kubevirt_vm_info_v1": {
            "namespace": meta.get("namespace"),
            "name": meta.get("name"),
            "uid": meta.get("uid"),
            "labels": meta.get("labels") or {},
            "annotations": annotations,
            "cluster": cluster,
            "kubernetes_cluster_hostname": host,
            "run_strategy": spec.get("runStrategy"),
            "disks": disks,
            "volumes": volumes,
        },
        "kubevirt_vm_status_v1": {
            "printable_status": status.get("printableStatus"),
            "run_strategy": spec.get("runStrategy"),
            "desired_running": spec.get("running"),
            "ready": status.get("ready"),
            "created": status.get("created"),
            "phase": vmi_status.get("phase"),
            "conditions": status.get("conditions") or [],
        },
        "kubevirt_vm_placement_v1": {"node": vmi_status.get("nodeName")},
        "kubevirt_vm_interfaces_v1": {"interfaces": interfaces},
        "kubevirt_vm_resources_v1": {"cpu_cores": _parse_quantity(cpu.get("cores")), "memory_request": _parse_quantity(requests.get("memory"))},
    }


def _write_pb(piggyback: str, sections: Mapping[str, Mapping[str, Any]]) -> None:
    sys.stdout.write(f"<<<<{piggyback}>>>>\n")
    for section_name, data in sections.items():
        sys.stdout.write(f"<<<{section_name}:sep(0)>>>\n")
        sys.stdout.write(json.dumps(data, separators=(",", ":")) + "\n")
    sys.stdout.write("<<<<>>>>\n")


def _monitor_namespace(name: str, include: Sequence[str], exclude: Sequence[str]) -> bool:
    if include:
        return any(re.search(pattern, name) for pattern in include)
    if exclude:
        return not any(re.search(pattern, name) for pattern in exclude)
    return True


def collect_api(args: argparse.Namespace, token: str) -> None:
    if not args.verify_cert_api:
        LOG.warning("Disabling SSL certificate verification for Kubernetes API endpoint.")
    kwargs = {
        "verify": args.verify_cert_api,
        "timeout": (args.k8s_api_connect_timeout, args.k8s_api_read_timeout),
        "proxies": _proxy(args.api_server_proxy),
    }
    with requests.Session() as session:
        session.headers.update({"Authorization": f"Bearer {token}"})
        vms = _list_items(session, args.api_server_endpoint, "/apis/kubevirt.io/v1/virtualmachines", **kwargs)
        vmis = _list_items(session, args.api_server_endpoint, "/apis/kubevirt.io/v1/virtualmachineinstances", **kwargs)
    vmis_by_key = {((vmi.get("metadata") or {}).get("namespace"), (vmi.get("metadata") or {}).get("name")): vmi for vmi in vmis}
    for vm in vms:
        meta = vm.get("metadata", {}) or {}
        namespace = str(meta.get("namespace", "default"))
        name = str(meta.get("name"))
        if not _monitor_namespace(namespace, args.namespace_include_patterns, args.namespace_exclude_patterns):
            continue
        piggyback = f"virtualmachine_{args.cluster}_{namespace}_{name}"
        sections = _vm_section_data(vm, vmis_by_key.get((namespace, name)), args.cluster, args.kubernetes_cluster_hostname, args.include_annotations_as_host_labels, args.include_matching_annotations_as_host_labels)
        _write_pb(piggyback, sections)


def prom_query(session: requests.Session, base_url: str, promql: str, **kwargs: Any) -> list[Mapping[str, Any]]:
    data = _get_json(session, base_url.rstrip("/") + "/api/v1/query", params={"query": promql}, **kwargs)
    if data.get("status") != "success":
        # A failed query is a collection error, not a successful empty vector.
        # Never include server-provided error strings (or credentials) in output.
        raise RuntimeError("Prometheus API returned a failed query status")
    payload = data.get("data")
    if not isinstance(payload, Mapping) or not isinstance(payload.get("result"), list):
        raise RuntimeError("Prometheus API returned an invalid query result")
    result = payload["result"]
    return [sample for sample in result if isinstance(sample, Mapping)]


def _value(sample: Mapping[str, Any]) -> float:
    value = sample.get("value", [0, 0])
    try:
        return float(value[1])
    except Exception:
        return 0.0


def _metric(sample: Mapping[str, Any]) -> Mapping[str, Any]:
    metric = sample.get("metric") or {}
    return metric if isinstance(metric, Mapping) else {}


def _vm_key(metric: Mapping[str, Any]) -> tuple[str, str]:
    return (str(metric.get("exported_namespace") or metric.get("namespace") or "default"), str(metric.get("name") or metric.get("domain") or metric.get("vm") or "unknown"))


def collect_prometheus(args: argparse.Namespace, prometheus_token: str | None = None) -> None:
    if not args.prometheus_endpoint:
        return
    if prometheus_token:
        validate_prometheus_token_transport(args.prometheus_endpoint, args.usage_verify_cert)
    if not args.usage_verify_cert:
        LOG.warning("Disabling SSL certificate verification for Prometheus endpoint.")
    kwargs = {
        "verify": args.usage_verify_cert,
        "timeout": (args.usage_connect_timeout, args.usage_read_timeout),
        "proxies": _proxy(args.usage_proxy),
    }
    with requests.Session() as session:
        if prometheus_token:
            session.headers.update({"Authorization": f"Bearer {prometheus_token}"})
        results = {name: prom_query(session, args.prometheus_endpoint, query, **kwargs) for name, query in KUBEVIRT_QUERIES.items()}

    per_vm: dict[tuple[str, str], dict[str, Any]] = defaultdict(lambda: {"net": defaultdict(dict), "disk": defaultdict(dict)})
    for sample in results["cpu"]:
        per_vm[_vm_key(_metric(sample))]["cpu"] = _value(sample)
    for sample in results["memory"]:
        per_vm[_vm_key(_metric(sample))]["memory"] = _value(sample)
    for key, field in [("net_rx", "receive_bytes_per_second"), ("net_tx", "transmit_bytes_per_second")]:
        for sample in results[key]:
            metric = _metric(sample)
            per_vm[_vm_key(metric)]["net"][str(metric.get("interface", "default"))][field] = _value(sample)
    for key, field in [("disk_read", "read_bytes_per_second"), ("disk_write", "write_bytes_per_second")]:
        for sample in results[key]:
            metric = _metric(sample)
            per_vm[_vm_key(metric)]["disk"][str(metric.get("drive", "unknown"))][field] = _value(sample)

    for (namespace, name), values in per_vm.items():
        sections = {}
        if "cpu" in values:
            sections["kubevirt_vm_cpu_v1"] = {"usage": values["cpu"]}
        if "memory" in values:
            sections["kubevirt_vm_memory_v1"] = {"resident_bytes": values["memory"]}
        if values["net"]:
            sections["kubevirt_vm_network_v1"] = {"interfaces": [{"interface": iface, "receive_bytes_per_second": data.get("receive_bytes_per_second", 0.0), "transmit_bytes_per_second": data.get("transmit_bytes_per_second", 0.0)} for iface, data in values["net"].items()]}
        if values["disk"]:
            sections["kubevirt_vm_disk_v1"] = {"drives": [{"drive": drive, "read_bytes_per_second": data.get("read_bytes_per_second", 0.0), "write_bytes_per_second": data.get("write_bytes_per_second", 0.0)} for drive, data in values["disk"].items()]}
        if sections:
            _write_pb(f"virtualmachine_{args.cluster}_{namespace}_{name}", sections)

    # Platform sections are written to the main host.
    platform = {}
    if results["version"]:
        platform["kubevirt_version_v1"] = {"labels": dict(_metric(results["version"][0]))}
    platform["kubevirt_component_health_v1"] = {"targets": [{"component": str(_metric(s).get("pod") or _metric(s).get("container") or _metric(s).get("job") or "unknown").split("-")[0], "pod": _metric(s).get("pod"), "namespace": _metric(s).get("namespace"), "up": _value(s)} for s in results["component_up"]]}
    platform["kubevirt_controller_status_v1"] = {"controller_leading": _value(results["controller_leading"][0]) if results["controller_leading"] else None, "controller_ready": _value(results["controller_ready"][0]) if results["controller_ready"] else None, "operator_leading": _value(results["operator_leading"][0]) if results["operator_leading"] else None, "operator_ready": _value(results["operator_ready"][0]) if results["operator_ready"] else None}
    queues: dict[str, dict[str, float | str]] = defaultdict(lambda: {"depth": 0.0, "adds_per_second": 0.0, "retries_per_second": 0.0, "unfinished_work_seconds": 0.0})
    for qname, field in [("workqueue_depth", "depth"), ("workqueue_adds", "adds_per_second"), ("workqueue_retries", "retries_per_second"), ("workqueue_unfinished", "unfinished_work_seconds")]:
        for sample in results[qname]:
            metric = _metric(sample)
            name = str(metric.get("name") or metric.get("queue") or metric.get("workqueue") or "unknown")
            queues[name]["name"] = name
            queues[name][field] = _value(sample)
    platform["kubevirt_workqueues_v1"] = {"queues": list(queues.values())}
    platform["kubevirt_migrations_v1"] = {"pending": _value(results["migrations_pending"][0]) if results["migrations_pending"] else 0.0, "running": _value(results["migrations_running"][0]) if results["migrations_running"] else 0.0, "scheduling": _value(results["migrations_scheduling"][0]) if results["migrations_scheduling"] else 0.0, "unset": _value(results["migrations_unset"][0]) if results["migrations_unset"] else 0.0}
    platform["kubevirt_vmi_phase_count_v1"] = {"phases": {str(_metric(s).get("phase", "unknown")): _value(s) for s in results["vmi_phase_count"]}}
    lifecycle = []
    for key, status in [("vm_running", "running"), ("vm_starting", "starting"), ("vm_non_running", "non_running"), ("vm_migrating", "migrating"), ("vm_error", "error")]:
        for sample in results[key]:
            metric = _metric(sample)
            lifecycle.append({"namespace": metric.get("exported_namespace") or metric.get("namespace"), "name": metric.get("name"), "status": status, "timestamp": _value(sample)})
    platform["kubevirt_vm_lifecycle_v1"] = {"timestamps": lifecycle}
    _write_pb("", platform)


def _prometheus_failure_summary(exc: Exception) -> str:
    """Return a useful error without exposing URLs, tokens or server payloads.

    requests exceptions can contain credentials embedded in endpoint URLs.
    Never display their raw exception messages in an agent section or log.
    """
    if isinstance(exc, requests.HTTPError):
        response = exc.response
        status = response.status_code if response is not None else None
        if status in (401, 403):
            return f"HTTP {status}: Prometheus authentication/authorization failed"
        if status is not None:
            return f"HTTP {status}: Prometheus request failed"
        return "Prometheus request failed (HTTP error)"
    if isinstance(exc, requests.exceptions.SSLError):
        return "Prometheus TLS certificate validation failed"
    if isinstance(exc, requests.exceptions.Timeout):
        return "Prometheus connection timed out"
    if isinstance(exc, requests.exceptions.ConnectionError):
        return "Prometheus connection failed"
    if isinstance(exc, requests.exceptions.RequestException):
        return "Prometheus HTTP request failed"
    if isinstance(exc, ValueError) and str(exc).startswith("Prometheus bearer authentication requires"):
        return "Prometheus bearer token requires verified HTTPS or a loopback endpoint"
    if isinstance(exc, RuntimeError) and str(exc).startswith("HTTP redirect refused"):
        return "Prometheus HTTP redirect refused"
    if isinstance(exc, RuntimeError) and str(exc).startswith("Prometheus API returned"):
        return str(exc)
    return f"Prometheus collection failed ({type(exc).__name__})"


def _emit_prometheus_health(ok: bool, message: str) -> None:
    # Empty piggyback name means source host, even after preceding VM sections.
    _write_pb("", {"kubevirt_prometheus_connection_v1": {"ok": ok, "message": message}})


def main(argv: list[str] | None = None) -> int:
    if argv is None:
        argv = sys.argv[1:]
    args = parse_args(argv)
    setup_logging(args.verbose)
    # Never log argparse's Namespace here: it may include explicit bearer tokens.
    if not args.verify_cert_api or (args.prometheus_endpoint and not args.usage_verify_cert):
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    try:
        token = _resolve_token(args)
        collect_api(args, token)
    except Exception as exc:
        # Kubernetes API collection is mandatory; preserving its failure signal
        # avoids silently reporting empty VM inventories as healthy.
        if args.debug:
            raise
        sys.stderr.write(str(exc) + "\n")
        return 1

    # Prometheus is optional. A broken metrics endpoint, invalid Prometheus
    # password-store entry or unsafe transport must not suppress API VM data.
    if args.prometheus_endpoint:
        try:
            prometheus_token = _resolve_prometheus_token(args)
            collect_prometheus(args, prometheus_token)
        except Exception as exc:
            message = _prometheus_failure_summary(exc)
            sys.stderr.write(f"Prometheus: {message}\n")
            _emit_prometheus_health(False, message)
        else:
            _emit_prometheus_health(True, "Prometheus queries completed successfully")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
