Measures latency to a host (Threaded) and plot it with matplotlib

ping-stat

import re
import sys
import click
import logging

import numpy as np
import matplotlib.pyplot as plt

from subprocess import check_output

from threading import Thread, Lock

logging.basicConfig(level=logging.INFO)


def ping_stat(host: str, count: int):
    logging.info(f'[*] Measuring latency for {host}')
    out = check_output(['ping', '-c', str(count), host])
    data = re.findall(r'time=([0-9.]+) ', str(out))

    return [float(v) for v in data]


def ping_stat_async(hosts: list, count: int) -> dict:
    threads = []
    result = {}

    lock = Lock()

    def return_wrapper(host: str, count: int):
        res = ping_stat(host,  count)
        lock.acquire()
        result[host] = res
        lock.release()

    for h in hosts:
        th = Thread(target=return_wrapper, args=(h, count))
        threads.append(th)
        th.start()

    for t in threads:
        t.join()

    return result

# TODO: Load these from file

dns_providers = {
    'CloudFlare': '1.1.1.1',
    'CloudFlare Secondary': '9.9.9.9',
    'Google': '8.8.8.8',
}


def plot_provider(provider, data, ax, color='r'):
    timeOffset = [np.sum(data[:i]) for i in range(len(data))]
    ax.plot(timeOffset, data, color, label=f'{provider}')
    ax.set(
        xlabel='time (ms)', ylabel='Latency (ms)',
        title='Latency Over Time'
    )


@click.command()
@click.option('--samples', default=5, help='Number of ping samples')
@click.option('--output', default='output.png', help='Output file')
def cli(samples, output):
    logging.info(f'[*] Measuring latency with {samples} samples')
    fig, ax = plt.subplots()

    results = ping_stat_async(dns_providers.values(), samples)
    print(results)

    colors = ['r', 'g', 'b', 'c', 'b', 'm', 'y', 'k']

    for host, data in results.items():
        plot_provider(host, data, ax, colors.pop())

    ax.grid()
    ax.legend()
    fig.savefig(output)

    return 0


if __name__ == '__main__':
    sys.exit(cli())