import sys
import time
import paramiko

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

def check_vip(ssh):
    stdin, stdout, stderr = ssh.exec_command("ip -4 addr show dev eth0 | grep '140.44.4.70' || true")
    return bool(stdout.read().decode().strip())

def curl_vip(ssh):
    stdin, stdout, stderr = ssh.exec_command("curl -s -o /dev/null -w '%{http_code}' -m 2 http://140.44.4.70:80 || true")
    return stdout.read().decode().strip()

def main():
    ssh1 = paramiko.SSHClient()
    ssh1.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh1.connect('140.44.4.71', username='snarf', key_filename=r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet', timeout=10)

    ssh2 = paramiko.SSHClient()
    ssh2.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh2.connect('140.44.4.72', username='snarf', key_filename=r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet', timeout=10)

    print("=== INITIAL STATE ===")
    v1 = check_vip(ssh1)
    v2 = check_vip(ssh2)
    code = curl_vip(ssh1)
    print(f"swm01 has VIP: {v1}, swm02 has VIP: {v2}, HTTP code on VIP: {code}")

    print("\n=== TRIGGERING FAILOVER: STOPPING KEEPALIVED ON swm01 ===")
    ssh1.exec_command("sudo systemctl stop keepalived")
    
    # Wait up to 5 seconds for swm02 to claim VIP
    claimed = False
    for i in range(5):
        time.sleep(1)
        if check_vip(ssh2):
            claimed = True
            break
            
    v1 = check_vip(ssh1)
    v2 = check_vip(ssh2)
    code = curl_vip(ssh2)
    print(f"Failover result -> swm01 has VIP: {v1}, swm02 has VIP: {v2}, HTTP code on VIP: {code}")

    print("\n=== RESTORING swm01: STARTING KEEPALIVED ON swm01 ===")
    ssh1.exec_command("sudo systemctl start keepalived")
    
    # Wait up to 5 seconds for swm01 to reclaim VIP
    reclaimed = False
    for i in range(5):
        time.sleep(1)
        if check_vip(ssh1):
            reclaimed = True
            break
            
    v1 = check_vip(ssh1)
    v2 = check_vip(ssh2)
    code = curl_vip(ssh1)
    print(f"Fallback result -> swm01 has VIP: {v1}, swm02 has VIP: {v2}, HTTP code on VIP: {code}")

    ssh1.close()
    ssh2.close()

if __name__ == '__main__':
    main()