import paramiko
import sys

sys.stdout.reconfigure(encoding='utf-8')

key_path = r"C:\Users\Snarf\.ssh\id_ed25519_snarfnet"

def run_ssh(host, cmds, desc):
    print(f"\n=======================================================")
    print(f"[{desc}] Connecting to {host}...")
    print(f"=======================================================")
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(host, username="root", key_filename=key_path, timeout=15)
    for cmd in cmds:
        print(f"\n>>> Running on {host}: {cmd}")
        stdin, stdout, stderr = c.exec_command(cmd)
        out = stdout.read().decode('utf-8')
        err = stderr.read().decode('utf-8')
        if out: print(out.strip())
        if err: print("ERR:", err.strip())
    c.close()

# 1. Aerodrome Master (140.44.4.11)
aero_cmds = [
    # Take post-upgrade etcd snapshot
    "k3s etcd-snapshot save --name post-upgrade-v2.15.3-final",
    # List snapshots to verify
    "k3s etcd-snapshot ls | tail -n 5",
    # Clean up temporary rancher dirs in /tmp
    "rm -rf /tmp/rancher-* /tmp/rancher_upgrade* /tmp/chart* /tmp/extracted*",
    # Prune completed helm operations in cattle-system
    "k3s kubectl -n cattle-system delete pod -l helm.cattle.io/chart=rancher --field-selector status.phase=Succeeded 2>/dev/null || true",
    "k3s kubectl -n cattle-system delete pod -l app.kubernetes.io/managed-by=Helm --field-selector status.phase=Succeeded 2>/dev/null || true",
    # Check cattle-system pods
    "k3s kubectl -n cattle-system get pods"
]

run_ssh("140.44.4.11", aero_cmds, "Aerodrome Master")

# 2. Spaceport Master (140.44.4.41)
spaceport_cmds = [
    # Take post-upgrade etcd snapshot
    "k3s etcd-snapshot save --name post-upgrade-spaceport-final",
    # List snapshots to verify
    "k3s etcd-snapshot ls | tail -n 5",
    # Prune completed system agent upgrader and helm pods
    "k3s kubectl -n cattle-system delete pod -l app=system-agent-upgrader --field-selector status.phase=Succeeded 2>/dev/null || true",
    "k3s kubectl -n cattle-system delete pod --field-selector status.phase=Succeeded 2>/dev/null || true",
    # Check cattle-system pods
    "k3s kubectl -n cattle-system get pods"
]

run_ssh("140.44.4.41", spaceport_cmds, "Spaceport Master")