#!/usr/bin/env python3
"""
fence_pve_ssh - Proxmox VE SSH fence agent for Pacemaker.

Uses SSH to reach pve1.sweet.home and run 'qm stop/start <vmid>'.
Designed for test-lab HA cluster only — not for production.

Configuration (as pacemaker stonith resource attributes):
  pve_host     Proxmox host to SSH to (default: pve1.sweet.home)
  pve_user     SSH user (default: wayne)
  key_file     SSH private key path (default: /etc/fence-pve-ssh-key)
  vmid_node1   VMID for ha-test-node1 (e.g. 200)
  vmid_node2   VMID for ha-test-node2 (e.g. 201)
  plug         Node name to act on (set by pacemaker: ha-test-node1 or ha-test-node2)
  action       Action: off|on|reboot|status|list|metadata
"""

import argparse
import subprocess
import sys
import os


METADATA = """<?xml version="1.0" ?>
<resource-agent name="fence_pve_ssh" shortdesc="Proxmox VE SSH fence agent (test lab)">
  <longdesc>Fences a VM on a Proxmox VE host by SSHing to the PVE host and
  running qm stop/start.  For test use only.</longdesc>
  <vendor-url>https://proxmox.com</vendor-url>
  <parameters>
    <parameter name="action" required="1" unique="0">
      <getopt mixed="-a, --action=[action]"/>
      <content type="string" default="reboot"/>
      <shortdesc lang="en">Fencing action: off|on|reboot|status|list</shortdesc>
    </parameter>
    <parameter name="plug" required="0" unique="0">
      <getopt mixed="-n, --plug=[nodename]"/>
      <content type="string"/>
      <shortdesc lang="en">Cluster node name to fence</shortdesc>
    </parameter>
    <parameter name="pve_host" required="0" unique="0">
      <getopt mixed="--pve-host=[host]"/>
      <content type="string" default="pve1.sweet.home"/>
      <shortdesc lang="en">Proxmox VE host to SSH to</shortdesc>
    </parameter>
    <parameter name="pve_user" required="0" unique="0">
      <getopt mixed="--pve-user=[user]"/>
      <content type="string" default="wayne"/>
      <shortdesc lang="en">SSH user on the Proxmox host</shortdesc>
    </parameter>
    <parameter name="key_file" required="0" unique="0">
      <getopt mixed="--key-file=[path]"/>
      <content type="string" default="/etc/fence-pve-ssh-key"/>
      <shortdesc lang="en">SSH private key file path</shortdesc>
    </parameter>
    <parameter name="vmid_node1" required="1" unique="0">
      <getopt mixed="--vmid-node1=[vmid]"/>
      <content type="string"/>
      <shortdesc lang="en">VMID for ha-test-node1</shortdesc>
    </parameter>
    <parameter name="vmid_node2" required="1" unique="0">
      <getopt mixed="--vmid-node2=[vmid]"/>
      <content type="string"/>
      <shortdesc lang="en">VMID for ha-test-node2</shortdesc>
    </parameter>
  </parameters>
  <actions>
    <action name="off" timeout="60s"/>
    <action name="on" timeout="60s"/>
    <action name="reboot" timeout="60s"/>
    <action name="status" timeout="30s"/>
    <action name="list" timeout="10s"/>
    <action name="metadata" timeout="5s"/>
  </actions>
</resource-agent>
"""


def parse_args():
    p = argparse.ArgumentParser(add_help=False)
    p.add_argument("-a", "--action", default="reboot")
    p.add_argument("-n", "--plug")
    p.add_argument("--pve-host", default="pve1.sweet.home")
    p.add_argument("--pve-user", default="wayne")
    p.add_argument("--key-file", default="/etc/fence-pve-ssh-key")
    p.add_argument("--vmid-node1")
    p.add_argument("--vmid-node2")
    # Allow remaining unknown args (pacemaker may pass extra ones)
    return p.parse_known_args()[0]


def ssh(pve_host, pve_user, key_file, cmd):
    result = subprocess.run(
        [
            "ssh",
            "-i", key_file,
            "-o", "StrictHostKeyChecking=no",
            "-o", "BatchMode=yes",
            "-o", "ConnectTimeout=10",
            f"{pve_user}@{pve_host}",
            cmd,
        ],
        capture_output=True,
        text=True,
        timeout=30,
    )
    return result


def get_vmid(args):
    node = args.plug
    if not node:
        print("ERROR: --plug not specified", file=sys.stderr)
        sys.exit(1)
    mapping = {
        "ha-test-node1": args.vmid_node1,
        "ha-test-node2": args.vmid_node2,
    }
    vmid = mapping.get(node)
    if not vmid:
        print(f"ERROR: unknown node '{node}'", file=sys.stderr)
        sys.exit(1)
    return vmid


def main():
    args = parse_args()
    action = args.action.lower()

    if action == "metadata":
        print(METADATA)
        sys.exit(0)

    if action == "list":
        if args.vmid_node1:
            print("ha-test-node1")
        if args.vmid_node2:
            print("ha-test-node2")
        sys.exit(0)

    vmid = get_vmid(args)

    if not os.path.exists(args.key_file):
        print(f"ERROR: SSH key not found at {args.key_file}", file=sys.stderr)
        sys.exit(1)

    if action in ("off", "reboot"):
        print(f"Stopping VM {vmid} ({args.plug}) on {args.pve_host}...")
        r = ssh(args.pve_host, args.pve_user, args.key_file,
                f"sudo /usr/sbin/qm stop {vmid}")
        if r.returncode != 0:
            print(f"ERROR stopping VM: {r.stderr}", file=sys.stderr)
            sys.exit(1)
        print(f"VM {vmid} stopped")

    if action in ("on", "reboot"):
        print(f"Starting VM {vmid} ({args.plug}) on {args.pve_host}...")
        r = ssh(args.pve_host, args.pve_user, args.key_file,
                f"sudo /usr/sbin/qm start {vmid}")
        if r.returncode != 0:
            print(f"ERROR starting VM: {r.stderr}", file=sys.stderr)
            sys.exit(1)
        print(f"VM {vmid} started")

    if action == "status":
        r = ssh(args.pve_host, args.pve_user, args.key_file,
                f"sudo /usr/sbin/qm status {vmid}")
        if r.returncode != 0:
            print(f"ERROR querying VM status: {r.stderr}", file=sys.stderr)
            sys.exit(1)
        # qm status returns "status: running" or "status: stopped"
        status_line = r.stdout.strip()
        print(status_line)
        if "stopped" in status_line:
            sys.exit(2)   # pacemaker interprets exit 2 as "off"
        sys.exit(0)       # running = exit 0


if __name__ == "__main__":
    main()