import paramiko
import json
import time

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

    print("=== STEP 1: Add node labels 'npm=true' to swm01 (.71) and swm02 (.72) ===")
    stdin, stdout, stderr = ssh.exec_command('sudo docker node update --label-add npm=true snshpd4lnxswm01 && sudo docker node update --label-add npm=true snshpd4lnxswm02')
    print("Labels update:", stdout.read().decode().strip(), stderr.read().decode().strip())

    # Check node labels
    stdin, stdout, stderr = ssh.exec_command("sudo docker node inspect snshpd4lnxswm01 --format '{{json .Spec.Labels}}'")
    print("Labels swm01:", stdout.read().decode().strip())
    stdin, stdout, stderr = ssh.exec_command("sudo docker node inspect snshpd4lnxswm02 --format '{{json .Spec.Labels}}'")
    print("Labels swm02:", stdout.read().decode().strip())

    ssh.close()

if __name__ == '__main__':
    main()