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.42"
node_name = "sninfr0lnxctl02"
print(f"Connecting to {node_name} ({target_ip}) to issue reboot -f...")

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

print(f"Waiting for {node_name} 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(f"\n{node_name} is DOWN! Reboot has taken effect!")
        break

print(f"Waiting for {node_name} 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; modprobe veth; lsmod | grep veth")
        out = stdout.read().decode('utf-8').strip()
        c.close()
        print(f"\n{node_name} is UP!\n{out}")
        up = True
        break
    except Exception:
        sys.stdout.write(".")
        sys.stdout.flush()
        time.sleep(3)

if not up:
    print(f"\nTimeout waiting for {node_name}.")