import paramiko
import time
import json

def ssh_connect_with_key(hostname, port, username, key_path):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname, port, username, key_filename=key_path)
    return ssh

def start_iperf_server(ssh):
    stdin, stdout, stderr = ssh.exec_command('iperf3 -s -1')
    return stdout, stderr

def start_iperf_client(ssh, server_ip, duration):
    command = f'iperf3 -c {server_ip} -t {duration} -J'
    stdin, stdout, stderr = ssh.exec_command(command)
    return stdout, stderr

def main():
    server_hostname = '192.168.0.10'
    server_port = 22
    server_username = 'root'
    server_key_path = '/home/apu/.ssh/id_ed25519'
    
    client_hostname = '192.168.0.11'
    client_port = 22
    client_username = 'root'
    client_key_path = '/home/apu/.ssh/id_ed25519'
    
    duration = 10  # duration of the iperf3 test in seconds

    # Connect to the server
    server_ssh = ssh_connect_with_key(server_hostname, server_port, server_username, server_key_path)
    print(f"Connected to server {server_hostname}")

    # Start iperf3 server
    server_stdout, server_stderr = start_iperf_server(server_ssh)
    print("Started iperf3 server")

    # Connect to the client
    client_ssh = ssh_connect_with_key(client_hostname, client_port, client_username, client_key_path)
    print(f"Connected to client {client_hostname}")

    # Start iperf3 client
    client_stdout, client_stderr = start_iperf_client(client_ssh, server_hostname, duration)
    print("Started iperf3 client")

    # Read iperf3 client output
    client_output = client_stdout.read().decode('utf-8')
    print(f"Client output: {client_output}")

    # Parse the result
    result = json.loads(client_output)
    throughput = result['end']['sum_received']['bits_per_second']
    
    # Save the results in the defined structure
    log_structure = f"{server_hostname} | {client_hostname} | 1 | {throughput / 1e9:.2f} Gbps\n"
    
    # Write the log to a file
    with open('iperf3_results.log', 'a') as log_file:
        log_file.write(log_structure)
    
    print(f"Results written to iperf3_results.log: {log_structure}")

    # Close SSH connections
    server_ssh.close()
    client_ssh.close()

if __name__ == "__main__":
    main()