import paramiko
import sys
import time

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

key_path = r"C:\Users\Snarf\.ssh\id_ed25519_snarfnet"
client = paramiko.SSHClient()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

target_ip = "140.44.4.43"
print(f"Connecting to {target_ip} to issue systemctl reboot...")

client.connect(target_ip, username="root", key_filename=key_path, timeout=10)
try:
    # Use systemctl reboot --force or nohup /sbin/reboot
    stdin, stdout, stderr = client.exec_command("systemctl reboot")
    print("reboot stdout:", stdout.read().decode('utf-8'))
    print("reboot stderr:", stderr.read().decode('utf-8'))
except Exception as e:
    print("Exception on reboot command (expected):", e)
client.close()

print("Waiting for node to drop off network...")
for i in range(30):
    try:
        c = paramiko.SSHClient()
        c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        c.connect(target_ip, username="root", key_filename=key_path, timeout=2)
        c.close()
        sys.stdout.write("x")
        sys.stdout.flush()
        time.sleep(2)
    except Exception:
        print("\nNode is DOWN! Reboot has taken effect!")
        break

print("Waiting for node to boot back up...")
start = time.time()
up = False
while time.time() - start < 180:
    try:
        c = paramiko.SSHClient()
        c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        c.connect(target_ip, username="root", key_filename=key_path, timeout=3)
        stdin, stdout, stderr = c.exec_command("uptime; uname -r")
        out = stdout.read().decode('utf-8').strip()
        c.close()
        print(f"\nNode is UP!\n{out}")
        up = True
        break
    except Exception:
        sys.stdout.write(".")
        sys.stdout.flush()
        time.sleep(3)

if not up:
    print("\nTimeout waiting for node.")