import paramiko
import sys
import time

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

KEY_FILE = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'
IP = '14.4.4.44'

print(f"=== UPGRADING NUTSERVER ({IP}) ===")

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(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
"""

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("[*] Gracefully rebooting nutserver...")
    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"[*] nutserver reconnected on attempt {attempt+1}!")
            break
        except Exception:
            time.sleep(4)
            
    if not reconnected:
        print("[!] ERROR: nutserver failed to reconnect!")
        sys.exit(1)
        
    _, out, _ = ssh.exec_command("systemctl is-active chrony || systemctl is-active chronyd")
    print("  Chrony NTP:", out.read().decode().strip())
    _, out, _ = ssh.exec_command("systemctl is-active nut-server")
    print("  NUT Server:", out.read().decode().strip())
    _, out, _ = ssh.exec_command("systemctl is-active zabbix-agent2")
    print("  ZabbixAgent2:", out.read().decode().strip())
    ssh.close()
else:
    ssh.close()

print("[SUCCESS] nutserver OS upgraded cleanly!")