import paramiko
import sys
import time

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

KEY_FILE = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'
swarm_nodes = [
    ('snshpd4lnxswm04', '140.44.4.74'),
    ('snshpd4lnxswm03', '140.44.4.73'),
    ('snshpd4lnxswm02', '140.44.4.72'),
    ('snshpd4lnxswm01', '140.44.4.71')
]

upgrade_cmd = """
echo "debconf debconf/frontend select Noninteractive" | sudo debconf-set-selections
echo "linux-base linux-base/removing-running-kernel boolean false" | sudo debconf-set-selections
export DEBIAN_FRONTEND=noninteractive
sudo apt-get update -qq
sudo apt-get -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" dist-upgrade -y
sudo apt-get autoremove -y
"""

for name, ip in swarm_nodes:
    print(f"\n=======================================================")
    print(f"[*] Upgrading Docker Swarm node {name} ({ip})...")
    print(f"=======================================================")
    
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(ip, username='snarf', key_filename=KEY_FILE, timeout=10)
    
    stdin, stdout, stderr = ssh.exec_command(upgrade_cmd)
    exit_code = stdout.channel.recv_exit_status()
    print(f"[*] apt dist-upgrade finished with exit code {exit_code}")
    
    # Check reboot
    _, out, _ = ssh.exec_command("test -f /var/run/reboot-required && echo 'yes' || echo 'no'")
    reboot_needed = out.read().decode().strip() == 'yes'
    print(f"[*] Reboot required: {reboot_needed}")
    
    if reboot_needed:
        print(f"[*] Gracefully rebooting {name}...")
        try:
            ssh.exec_command("sudo reboot")
        except Exception:
            pass
        ssh.close()
        
        print("[*] Waiting 30s for reboot...")
        time.sleep(30)
        
        reconnected = False
        for attempt in range(25):
            try:
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                ssh.connect(ip, username='snarf', key_filename=KEY_FILE, timeout=6)
                reconnected = True
                print(f"[*] {name} reconnected successfully on attempt {attempt+1}!")
                break
            except Exception:
                time.sleep(4)
                
        if not reconnected:
            print(f"[!] ERROR: {name} failed to reconnect!")
            sys.exit(1)
            
        _, out, _ = ssh.exec_command("sudo systemctl is-active docker")
        print(f"  Docker: {out.read().decode().strip()}")
        ssh.close()
    else:
        ssh.close()
        
    # Settle Swarm consensus
    print(f"[*] Settling Swarm Raft consensus (15s)...")
    time.sleep(15)

# Verify from swm01
final_ssh = paramiko.SSHClient()
final_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
final_ssh.connect('140.44.4.71', username='snarf', key_filename=KEY_FILE, timeout=10)
_, out, _ = final_ssh.exec_command('sudo docker node ls')
print("\n=== FINAL DOCKER SWARM NODES ===")
print(out.read().decode())

_, out, _ = final_ssh.exec_command('sudo docker service ls')
print("=== DOCKER SERVICES ===")
print(out.read().decode())

_, out, _ = final_ssh.exec_command('ip addr show | grep 140.44.4')
print("=== KEEPALIved VIP 140.44.4.70 ===")
print(out.read().decode())
final_ssh.close()

print("[SUCCESS] All Docker Swarm nodes upgraded to latest Docker Engine and packages!")