#!/usr/bin/env python3
"""
check_crm_status

Icinga/Nagios plugin to monitor Pacemaker/Corosync clusters
(Pacemaker 2.x / 3.x – Debian Bookworm & Trixie compatible).

Features:
- Executes `crm_mon -1 --output-as=xml` with timeout
- Parses cluster state from XML output using defusedxml
- Supports simple resources, groups, clones, master/slave
- Detects quorum loss, offline/unhealthy nodes, failed/blocked/unmanaged resources
- Handles maintenance mode (cluster, nodes, resources)
- Compares current state with a reference JSON state unless disabled
- Outputs Nagios plugin format: short status line, state differences and detailed cluster summary
- Standard Nagios exit codes: OK / WARNING / CRITICAL / UNKNOWN
"""

import argparse
import json
import logging
import os
import subprocess  # nosec
import sys
import traceback
import xml.etree.ElementTree as ET  # nosec
from pprint import pformat
from typing import Any, Dict, List, Optional, Tuple

from defusedxml import ElementTree as DFT  # defusedxml safe parser

STATE_OK = 0
STATE_WARNING = 1
STATE_CRITICAL = 2
STATE_UNKNOWN = 3

DEFAULT_STATE_FILE = "/var/cache/ee/check_crm_status.state.json"
DEFAULT_CRM_MON_PATH = "crm_mon"
CRM_MON_CMD_ARGS = ("-1", "--output-as=xml")
CRM_MON_TIMEOUT = 10


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-i",
        "--init",
        action="store_true",
        help="Initialize or save current cluster state as reference",
    )
    parser.add_argument(
        "-D",
        "--disable-reference",
        action="store_true",
        help="Disable checking against reference state",
    )
    parser.add_argument("-s", "--state-file", default=DEFAULT_STATE_FILE, help="Reference state file")
    parser.add_argument(
        "-S",
        "--load-state-file",
        help="Load current cluster state from XML file (instead of running crm_mon)",
    )
    parser.add_argument(
        "-c",
        "--crm-mon-path",
        default=DEFAULT_CRM_MON_PATH,
        help="Path of crm_mon binary",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable INFO logging")
    parser.add_argument("-d", "--debug", action="store_true", help="Enable DEBUG logging")
    return parser.parse_args()


def run_crm_mon(crm_mon_path: str) -> str:
    """Execute crm_mon safely with timeout."""
    try:
        cmd = [crm_mon_path, *CRM_MON_CMD_ARGS]
        logging.info("Run '%s'", " ".join(cmd))
        result = subprocess.run(  # nosec
            cmd,
            capture_output=True,
            text=True,
            timeout=CRM_MON_TIMEOUT,
            check=True,
            shell=False,
        )
        return result.stdout
    except subprocess.TimeoutExpired as exc:
        raise RuntimeError("crm_mon execution timed out") from exc
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"crm_mon failed: {exc.stderr.strip()}") from exc


def parse_crm_mon_status(root: ET.Element) -> Dict[str, str]:
    """Parse the crm_mon status from XML root."""
    status_elem = root.find("status")
    if status_elem is None:
        raise RuntimeError("Can't find cluster summary in crm_mon output")
    status = dict(status_elem.attrib)
    logging.debug("Cluster status: %s", pformat(status))
    return status


def parse_cluster_summary(root: ET.Element) -> Dict[str, str]:
    """Parse the cluster summary from XML root."""
    summary: Dict[str, str] = {
        "current_dc": {},
        "cluster_options": {},
    }
    summary_elem = root.find("summary")
    if summary_elem is None:
        raise RuntimeError("Can't find cluster summary in crm_mon output")
    summary["current_dc"].update(summary_elem.find("current_dc").attrib)
    summary["cluster_options"].update(summary_elem.find("cluster_options").attrib)
    summary["with_quorum"] = summary["current_dc"].get("with_quorum", "false") == "true"
    summary["quorum_policy"] = summary["cluster_options"].get("no-quorum-policy", "stop")
    summary["maintenance_mode"] = summary["cluster_options"].get("maintenance-mode") == "true"
    logging.debug("Cluster summary: %s", pformat(summary))
    logging.info(
        "Cluster state: with quorum=%s, quorum policy=%s, maintenance mode=%s",
        summary["with_quorum"],
        summary["quorum_policy"],
        summary["maintenance_mode"],
    )
    return summary


def parse_nodes(root: ET.Element) -> Dict[str, Dict[str, str]]:
    """Parse nodes from XML root."""
    nodes: Dict[str, Dict[str, str]] = {}
    nodes_elem = root.find("nodes")
    if nodes_elem is None:
        raise RuntimeError("Can't find cluster nodes in crm_mon output")
    for node_elem in nodes_elem.findall(".//node"):
        node = dict(node_elem.attrib)
        name = node.get("name")
        nodes[name] = {
            "state": (
                "online"
                if node.get("online") == "true"
                else (
                    "standby"
                    if node.get("standby") == "true"
                    else (
                        "pending"
                        if node.get("pending") == "true"
                        else "shutdown" if node.get("shutdown") == "true" else "unknown"
                    )
                )
            ),
            "maintenance_mode": node.get("maintenance") == "true",
            **node,
        }
        logging.debug("Node %s: %s", name, nodes[name])
    logging.info("%d nodes found: %s", len(nodes), ", ".join(nodes))
    return nodes


def parse_resources(root: ET.Element) -> Dict[str, Any]:
    """Parse all resources, including clones and master/slave."""
    resources: Dict[str, Any] = {}

    resources_elem = root.find("resources")
    if resources_elem is None:
        raise RuntimeError("Can't find cluster resources in crm_mon output")

    cloned_resources = resources_elem.findall(".//clone")
    cloned_resources_id = {inst.attrib.get("id") for clone in cloned_resources for inst in clone.findall(".//resource")}
    logging.info(
        "%d cloned resources found (%s)",
        len(cloned_resources),
        ", ".join(cloned_resources_id),
    )

    for res_elem in resources_elem.findall(".//resource"):
        res = dict(res_elem.attrib)
        res_id = res.get("id")
        if res_id in cloned_resources_id:
            continue
        node_elem = res_elem.find("node")
        resources[res_id] = {
            "state": (
                "active"
                if res.get("active") == "true"
                else (
                    "failed"
                    if res.get("failed") == "true"
                    else "blocked" if res.get("blocked") == "true" else "inactive"
                )
            ),
            "maintenance_mode": res.get("maintenance") == "true",
            "node_name": (node_elem.attrib.get("name") if node_elem is not None else None),
            **res,
        }
        logging.debug("Resource %s: %s", res_id, pformat(resources[res_id]))
    logging.info("%d simple resources found (%s)", len(resources), ", ".join(resources))

    for clone_elem in cloned_resources:
        res_id = clone_elem.attrib.get("id")
        res = dict(clone_elem.attrib)
        resources[res_id] = {
            "maintenance_mode": res.get("maintenance") == "true",
            "instances": [],
            **res,
        }
        for res in clone_elem.findall(".//resource"):
            instance = dict(res.attrib)
            node_elem = res.find("node")
            resources[res_id]["instances"].append(
                {
                    "state": (
                        "active"
                        if instance.get("active") == "true"
                        else (
                            "failed"
                            if instance.get("failed") == "true"
                            else ("blocked" if instance.get("blocked") == "true" else "inactive")
                        )
                    ),
                    "maintenance_mode": instance.get("maintenance") == "true",
                    "node_name": (node_elem.attrib.get("name") if node_elem is not None else None),
                    **instance,
                }
            )
        logging.debug("Cloned resource %s: %s", res_id, pformat(resources[res_id]))

    return resources


def parse_failures(root: ET.Element) -> Dict[str, Any]:
    """Parse failures from XML root."""
    failures: List[Dict[str, Any]] = []
    for failure_elem in root.findall(".//failure"):
        failures.append(dict(failure_elem.attrib))
    logging.info("%d failures found", len(failures))
    return failures


def parse_crm_mon_xml(xml_data: str) -> Dict[str, Any]:
    """Parse XML output of crm_mon safely."""
    try:
        root = DFT.fromstring(xml_data)
    except DFT.ParseError as exc:
        raise RuntimeError(f"XML parsing error: {exc}") from exc
    return {
        "status": parse_crm_mon_status(root),
        "summary": parse_cluster_summary(root),
        "nodes": parse_nodes(root),
        "resources": parse_resources(root),
        "failures": parse_failures(root),
    }


def save_reference_state(state: Dict[str, Any], path: str) -> None:
    """Save the current state to reference file."""
    logging.info("Save current state in %s", path)
    logging.debug("State:\n%s", pformat(state))
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as file_handle:
        json.dump(state, file_handle, indent=2)
    logging.info("State file %s updated", path)


def load_reference_state(path: str) -> Optional[Dict[str, Any]]:
    """Load reference state from file if exists."""
    if not os.path.exists(path):
        return None
    logging.info("Loading current state from %s", path)
    with open(path, encoding="utf-8") as file_handle:
        state = json.load(file_handle)
    logging.debug("Reference state: %s", state)
    return state


def check_cluster_state(
    current: Dict[str, Any],
    reference: Optional[Dict[str, Any]],
    reference_enabled: bool = True,
) -> Tuple[int, List[str]]:
    """Check cluster health and return exit code and messages."""
    critical: List[str] = []
    warning: List[str] = []

    status = current["status"]
    summary = current["summary"]
    nodes = current["nodes"]
    resources = current["resources"]
    failures = current["failures"]

    if status["message"] != "OK":
        warning.append(f"crm_mon execution is {status['message']}")

    if not summary["with_quorum"]:
        (warning if summary["quorum_policy"] in ("ignore", "freeze") else critical).append(
            f"No quorum (policy={summary['quorum_policy']})"
        )

    if summary["maintenance_mode"]:
        warning.append("cluster is in maintenance mode")

    for name, node in nodes.items():
        if node["maintenance_mode"] and not summary["maintenance_mode"]:
            warning.append(f"node {name} in maintenance")
            continue
        if node.get("online") != "true":
            critical.append(f"node {name} is offline")
        elif node.get("standby") != "false":
            warning.append(f"node {name} is standby")
        elif node.get("health", "green") != "green":
            critical.append(f"node {name} health={node.get('health')}")

    for res_id, res in resources.items():
        if res["maintenance_mode"] and not summary["maintenance_mode"]:
            warning.append(f"resource {res_id} in maintenance")
            continue
        if res.get("managed") == "false" and not summary["maintenance_mode"]:
            warning.append(f"resource {res_id} unmanaged")
        elif res.get("failed") == "true":
            critical.append(f"resource {res_id} failed")
        elif res.get("blocked") == "true":
            critical.append(f"resource {res_id} blocked")
        if "instances" in res:
            # Close (or master/slave) resources
            instances = res["instances"]
            active = sum(1 for i in instances if i.get("active") == "true")
            expected = len(instances) if res.get("unique") == "false" else 1
            if active != expected:
                warning.append(f"{res_id} running on {active}/{expected} nodes")
            if res.get("multi_state") == "true":
                # Master/slave resources
                if not any(i.get("role") == "Promoted" for i in instances):
                    critical.append(f"{res_id} promoted on no node")
            else:
                for inst in instances:
                    if not inst["node_name"]:
                        continue
                    role = inst.get("role", "unknown")
                    target_role = inst.get("target_role", "unknown")
                    if role != target_role:
                        warning.append(f"{res_id} is {role} on {inst['node_name']} (target={target_role})")
        else:
            role = res.get("role", "unknown")
            target_role = res.get("target_role", "unknown")
            if role != target_role:
                warning.append(f"{res_id} is {role} (target={target_role})")

    if failures:
        warning.append(f"{len(failures)} failures detected")

    if reference_enabled:
        if reference is None:
            warning.append("reference state file missing")
        elif current != reference:
            warning.append("cluster state differs from reference")

    if critical:
        return STATE_CRITICAL, critical + warning
    if warning:
        return STATE_WARNING, warning
    return STATE_OK, []


def generate_cluster_summary(state: Dict[str, Any]) -> str:
    """Generate human-readable cluster summary."""
    lines = [""]
    summary = state["summary"]
    nodes = state["nodes"]
    resources = state["resources"]
    failures = state["failures"]
    if failures:
        lines.extend(
            [
                f"Failures: {len(failures)} failures detected",
                *[
                    f"- {f['op_key']} on {f['node']} (task: {f.get('task', 'unknown')}, date: {f['last-rc-change']}): "
                    f"{f.get('exitstatus') or 'unknown error'} (exit code: {f.get('exitcode', 'unknown')}"
                    f"{(', reason: ' + f['exitreason']) if f.get('exitreason') else ''})"
                    for f in failures
                ],
                "",
            ]
        )

    total_nodes = len(nodes)
    online_nodes = sum(1 for n in nodes.values() if n.get("online") == "true")
    maintenance_nodes = sum(1 for n in nodes.values() if n["maintenance_mode"] == "true")
    total_resources = len(resources)
    active_resources = sum(
        1
        for r in resources.values()
        if r.get("active") == "true" or any(i.get("active") == "true" for i in r.get("instances", []))
    )
    maintenance_resources = sum(1 for r in resources.values() if r["maintenance_mode"] == "true")
    lines += [
        f"Cluster: quorum={'yes' if summary['with_quorum'] else 'NO'}"
        f" (policy: {summary['quorum_policy']}) / "
        f"{'MAINTENANCE IN PROGRESS' if summary['maintenance_mode'] else 'maintenance: no'}",
        f"Nodes: {total_nodes} total, {online_nodes} online, {maintenance_nodes} in maintenance",
        "\n".join(
            [
                f"- {node_name} ({node['state']}" f"{' / MAINTENANCE IN PROGRESS' if node['maintenance_mode'] else ''})"
                for node_name, node in nodes.items()
            ]
        ),
        f"Resources: {total_resources} total, {active_resources} active, " f"{maintenance_resources} in maintenance",
        "\n".join(
            [
                f"- {res_id} "
                f"({res['state']}{(' on ' + res['node_name']) if res['node_name'] else ''}"
                f"{' / MAINTENANCE IN PROGRESS' if res['maintenance_mode'] else ''})"
                for res_id, res in resources.items()
                if "instances" not in res
            ]
        ),
        "\n".join(
            [
                f"- {res_id}{' (MAINTENANCE IN PROGRESS)' if res['maintenance_mode'] else ''}:"
                + "\n"
                + "\n".join(
                    [
                        f"  - {inst['node_name']} ({inst['state']} / {inst['role']}"
                        f"{' / MAINTENANCE IN PROGRESS' if inst['maintenance_mode'] else ''})"
                        for inst in res["instances"]
                        if inst["node_name"]
                    ]
                )
                for res_id, res in resources.items()
                if "instances" in res
            ]
        ),
    ]
    return "\n".join(lines)


def print_state_differences(reference, current, path=None):
    """Print differences between reference & current states."""
    if path is None:
        print("\nStates changes:")
    all_keys = set(reference.keys()).union(set(current.keys()))

    for key in sorted(all_keys):
        current_path = f"{path}.{key}" if path else key

        if key not in reference:
            print(f"+ ADDED   [{current_path}]: {pformat(current[key])}")
        elif key not in current:
            print(f"- REMOVED [{current_path}]: {pformat(reference[key])}")
        elif isinstance(reference[key], dict) and isinstance(current[key], dict):
            print_state_differences(reference[key], current[key], current_path)
        elif reference[key] != current[key]:
            print(f"~ CHANGED [{current_path}]: {pformat(reference[key])} ➔ {pformat(current[key])}")


def main() -> None:
    """Main entry point."""
    args = parse_args()
    logging.basicConfig(
        level=(logging.DEBUG if args.debug else logging.INFO if args.verbose else logging.WARNING),
        format="%(levelname)s: %(message)s",
    )

    if args.load_state_file:
        logging.info("Load state to check from %s", args.load_state_file)
        try:
            with open(args.load_state_file, encoding="utf-8") as file_desc:
                xml_output = file_desc.read()
        except OSError:
            print(f"UNKNOWN - Failed to load Corosync/Pacemaker cluster state from {args.load_state_file}")
            print(traceback.format_exc())
            sys.exit(STATE_UNKNOWN)
    else:
        try:
            xml_output = run_crm_mon(args.crm_mon_path)
        except RuntimeError:
            print("UNKNOWN - Failed to retrieve Corosync/Pacemaker cluster status")
            print(traceback.format_exc())
            sys.exit(STATE_UNKNOWN)
    logging.debug("Raw crm_mon XML output:\n%s", xml_output)

    try:
        current_state = parse_crm_mon_xml(xml_output)
        logging.debug("Parsed state:\n%s", json.dumps(current_state, indent=2))
    except RuntimeError:
        print("UNKNOWN - Failed to parse crm_mon XML output")
        print(traceback.format_exc())
        sys.exit(STATE_UNKNOWN)

    if args.init:
        try:
            save_reference_state(current_state, args.state_file)
            print(f"OK - Reference state saved to {args.state_file}")
            sys.exit(STATE_OK)
        except RuntimeError:
            print("UNKNOWN - Failed to save reference state")
            print(traceback.format_exc())
            sys.exit(STATE_UNKNOWN)

    reference_state = None
    if not args.disable_reference:
        try:
            reference_state = load_reference_state(args.state_file)
        except RuntimeError:
            print(f"UNKNOWN - Failed to load reference state from {args.state_file}")
            print(traceback.format_exc())
            sys.exit(STATE_UNKNOWN)

    exit_code, messages = check_cluster_state(current_state, reference_state, not args.disable_reference)
    state_label = ["OK", "WARNING", "CRITICAL", "UNKNOWN"][exit_code]

    if messages:
        print(f"{state_label} - {', '.join(messages[:5])}")
    else:
        print(f"{state_label} - Cluster is healthy")

    if not args.disable_reference and reference_state and reference_state != current_state:
        print_state_differences(reference_state, current_state)

    print(generate_cluster_summary(current_state))
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
