import paramiko
import sys
import time

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

KEY_FILE = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'
TARGET_IP = '140.44.4.75'
TARGET_NAME = 'snshpd4lnxswm05'

print(f"=== TESTING ROLLING OS UPDATE ON SWARM NODE {TARGET_NAME} ({TARGET_IP}) ===")

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(TARGET_IP, username='snarf', key_filename=KEY_FILE, timeout=10)

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
"""

print("[*] Running apt dist-upgrade...")
stdin, stdout, stderr = ssh.exec_command(upgrade_cmd)
exit_code = stdout.channel.recv_exit_status()
print(f"[*] apt dist-upgrade exit code: {exit_code}")

# Check reboot
stdin, stdout, stderr = 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("[*] Gracefully rebooting node...")
    try:
        ssh.exec_command("sudo reboot")
    except Exception:
        pass
    ssh.close()
    
    print("[*] Waiting 30s for node to 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(TARGET_IP, username='snarf', key_filename=KEY_FILE, timeout=6)
            reconnected = True
            print(f"[*] Node reconnected on attempt {attempt+1}!")
            break
        except Exception:
            time.sleep(4)
            
    if not reconnected:
        print("[!] ERROR: Node failed to reconnect!")
        sys.exit(1)
        
    _, out, _ = ssh.exec_command("sudo systemctl is-active docker")
    print("  Docker service:", out.read().decode().strip())
    _, out, _ = ssh.exec_command("systemctl is-active systemd-timesyncd || systemctl is-active chrony")
    print("  TimeSync:", out.read().decode().strip())
    _, out, _ = ssh.exec_command("systemctl is-active zabbix-agent2")
    print("  ZabbixAgent2:", out.read().decode().strip())
    ssh.close()
else:
    ssh.close()

# Verify on Swarm Leader (140.44.4.73)
leader_ssh = paramiko.SSHClient()
leader_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
leader_ssh.connect('140.44.4.73', username='snarf', key_filename=KEY_FILE, timeout=10)
_, out, _ = leader_ssh.exec_command('sudo docker node ls')
print("\n=== SWARM NODES FROM LEADER ===")
print(out.read().decode())
leader_ssh.close()

print(f"[SUCCESS] Swarm node {TARGET_NAME} updated cleanly!")