import paramiko
import sys
import time

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

MASTER_IP = '140.44.4.11'
WORKER_IP = '140.44.4.27'
NODE_NAME = 'sninfr2lnxwrk07'
KEY_FILE = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'

print(f"=== STEP 1: TEST ROLLING OS PATCH ON {NODE_NAME} ({WORKER_IP}) ===")

# 1. Connect to master to check pods on node
master_ssh = paramiko.SSHClient()
master_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
master_ssh.connect(MASTER_IP, username='root', key_filename=KEY_FILE, timeout=10)

print(f"[*] Checking active pods on {NODE_NAME}...")
_, stdout, _ = master_ssh.exec_command(f"kubectl get pods -A --field-selector spec.nodeName={NODE_NAME}")
print(stdout.read().decode())

print(f"[*] Cordoning {NODE_NAME}...")
_, stdout, stderr = master_ssh.exec_command(f"kubectl cordon {NODE_NAME}")
print("  " + stdout.read().decode().strip())

print(f"[*] Draining {NODE_NAME}...")
_, stdout, stderr = master_ssh.exec_command(f"kubectl drain {NODE_NAME} --ignore-daemonsets --delete-emptydir-data --force --timeout=180s")
print("  " + stdout.read().decode().strip())
err = stderr.read().decode().strip()
if err:
    print("  (drain notice):", err)

# 2. Connect to worker to perform apt-get update & dist-upgrade
print(f"[*] Connecting directly to worker {WORKER_IP} for package upgrades...")
worker_ssh = paramiko.SSHClient()
worker_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
worker_ssh.connect(WORKER_IP, username='root', key_filename=KEY_FILE, timeout=10)

upgrade_cmd = "DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get -o Dpkg::Options::='--force-confdef' -o Dpkg::Options::='--force-confold' dist-upgrade -y && apt-get autoremove -y"
print("[*] Running dist-upgrade (preserving configurations)...")
stdin, stdout, stderr = worker_ssh.exec_command(upgrade_cmd)
exit_status = stdout.channel.recv_exit_status()
print(f"[*] dist-upgrade completed with exit code: {exit_status}")

# 3. Check reboot requirement
_, stdout, _ = worker_ssh.exec_command("test -f /var/run/reboot-required && echo 'yes' || echo 'no'")
reboot_needed = stdout.read().decode().strip() == 'yes'
print(f"[*] Reboot required: {reboot_needed}")

if reboot_needed:
    print("[*] Rebooting worker node...")
    try:
        worker_ssh.exec_command("reboot")
    except Exception:
        pass
    worker_ssh.close()
    
    print("[*] Waiting 30 seconds for host to reboot...")
    time.sleep(30)
    
    # Poll SSH until online
    reconnected = False
    for attempt in range(25):
        try:
            worker_ssh = paramiko.SSHClient()
            worker_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            worker_ssh.connect(WORKER_IP, username='root', key_filename=KEY_FILE, timeout=6)
            reconnected = True
            print(f"[*] Worker reconnected successfully on attempt {attempt+1}!")
            break
        except Exception:
            time.sleep(5)
            
    if not reconnected:
        print("[!] ERROR: Worker node did not come back online within timeout!")
        sys.exit(1)
        
    # Verify services
    _, stdout, _ = worker_ssh.exec_command("systemctl is-active systemd-timesyncd || systemctl is-active chrony")
    print(f"  TimeSync: {stdout.read().decode().strip()}")
    _, stdout, _ = worker_ssh.exec_command("systemctl is-active zabbix-agent2")
    print(f"  ZabbixAgent2: {stdout.read().decode().strip()}")
    _, stdout, _ = worker_ssh.exec_command("systemctl is-active k3s-agent || systemctl is-active k3s")
    print(f"  K3sAgent: {stdout.read().decode().strip()}")
    worker_ssh.close()
else:
    worker_ssh.close()

# 4. Uncordon node
print(f"[*] Uncordoning {NODE_NAME}...")
_, stdout, _ = master_ssh.exec_command(f"kubectl uncordon {NODE_NAME}")
print("  " + stdout.read().decode().strip())

# 5. Wait for node to be Ready
print(f"[*] Verifying node status in Kubernetes...")
for _ in range(15):
    _, stdout, _ = master_ssh.exec_command(f"kubectl get node {NODE_NAME} -o jsonpath='{{.status.conditions[?(@.type==\"Ready\")].status}}'")
    status = stdout.read().decode().strip()
    if status == "True":
        print(f"[OK] Node {NODE_NAME} is fully Ready!")
        break
    time.sleep(4)

# 6. Check Longhorn volume robustness
print("[*] Settling Longhorn volume replicas...")
time.sleep(10)
_, stdout, _ = master_ssh.exec_command("kubectl get volumes.longhorn.io -n longhorn-system -o custom-columns=NAME:.metadata.name,STATE:.status.state,ROBUSTNESS:.status.robustness")
print(stdout.read().decode())

master_ssh.close()
print(f"[SUCCESS] Step 1 Test Worker OS Patch on {NODE_NAME} completed flawlessly!")