cowmesh-network-test/cowmesh_router_iperf_test.py

"""
script to do a basic holistic test of the cowmesh network,
testing internet connection speed using iperf,
between all nodes in the mesh
"""
import datetime
import os
import re
import time
import asyncio, asyncssh, sys
import paramiko
import subprocess
import json
from cowmesh_constants import NODES, HOST_TO_IP

from cowmesh_helpers import cleanup_iperf_server

PROJECT_PATH = os.path.abspath(os.path.dirname(__file__))
SECRETS_PATH = os.path.join(PROJECT_PATH, "secrets.json")
with open(SECRETS_PATH, 'r') as f:
    SECRETS = json.loads(f.read())


class CowmeshRouterIperfTester:

    def __init__(self, log=None, debug=False, time=10):
        if log:
            self.log = log
        self.debug = debug
        self.time = time
        self.results = {}

    async def log(self, msg):
        print(msg)

    async def debug_log(self, msg):
        if self.debug:
            await self.log(msg)

    async def test_between_two_nodes(self, node_a, node_b):
        await self.log("++ running test from {} to {}".format(node_a, node_b))
        u_name = 'root'
        pswd = SECRETS["ROUTER_PASSWORD"]

        myconn = paramiko.SSHClient()
        myconn.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        myconn.connect(node_a, username =u_name, password=pswd)

        ip = HOST_TO_IP[node_b]

        remote_cmd = 'iperf -c {ip} -p 5001 -t {seconds}'.format(ip=ip, seconds=self.time)
        (stdin, stdout, stderr) = myconn.exec_command(remote_cmd)
        output = str(stdout.read())
        await self.debug_log("output: {}".format(output))
        await self.debug_log("errors: {}".format(stderr.read()))
        myconn.close()

        match = re.search("(\S+) Mbits", output)
        if match:
            to_return = match.group(1)
        else:
            match = re.search("(\S+) Kbits", output)
            if match:
                to_return = match.group(1)
                to_return = float(to_return) / 1000.0
            else:
                to_return = None

        return to_return

    async def start_iperf_servers(self):
        for node in NODES:
            try:
                print("++ starting iperf server on {}".format(node))
                u_name = 'root'
                pswd = SECRETS["ROUTER_PASSWORD"]

                myconn = paramiko.SSHClient()
                myconn.set_missing_host_key_policy(paramiko.AutoAddPolicy())

                myconn.connect(node, username=u_name, password=pswd)

                remote_cmd = 'iperf -s &'
                (stdin, stdout, stderr) = myconn.exec_command(remote_cmd)
                await self.debug_log("{}".format(stdout.read()))
                await self.debug_log("{}".format(type(myconn)))
                await self.debug_log("Options available to deal with the connections are many like\n{}".format(dir(myconn)))
                myconn.close()
            except:
                try:
                    await self.log("++ error starting iperf server on {}".format(node))
                except:
                    continue

    async def cleanup_iperf_servers(self):
        await self.log("shutting down iperf servers")
        for node in NODES:
            user = "root"
            password = SECRETS["ROUTER_PASSWORD"]
            ip = HOST_TO_IP[node]
            await cleanup_iperf_server(node=node, ip=ip, username=user, password=password, log=self.log)

    async def run_test(self):

        try:
            await self.start_iperf_servers()
        except:
            pass

        for node_a in NODES:
            for node_b in NODES:
                if node_a == node_b:
                    await self.debug_log("skip self")
                    continue

                try:
                    r = await self.test_between_two_nodes(node_a, node_b)
                    result_key = "{} -> {}".format(node_a, node_b)
                    self.results[result_key] = r
                except Exception as e:
                    try:
                        await self.log("++ error running test between {} and {}: {}".format(node_a, node_b, e))
                    except:
                        continue
        # try:
        #     await self.cleanup_iperf_servers()
        # except:
        #     await self.log("error shutting down iperf servers")

    async def output_results(self):
        results_str = ""
        now = datetime.datetime.now()
        date = now.date()
        time = now.time()
        results_str += "**** iperf results on {date:%m-%d-%Y} at {time:%H:%M}:\n\n".format(date=date, time=time)
        for test_name, result in self.results.items():
            results_str += "{}: {} mbps\n".format(test_name, result)
        await self.log(results_str)


if __name__ == "__main__":
    try:
        tester = CowmeshRouterIperfTester()

        async def main_fun():
            await tester.run_test()
            await tester.output_results()

        asyncio.get_event_loop().run_until_complete(main_fun())


    except (OSError, asyncssh.Error) as exc:
        sys.exit('SSH connection failed: ' + str(exc))