peerfinder/peerfinderv2/client.py

#!/usr/bin/env python

import socket
import select
import random
import subprocess
import time
import re

import peerfinder_pb2 as pf

SERVER = ("::1", 9999)
# In seconds
HEARTBEAT_INTERVAL = random.randint(5, 15)
RECONNECT_INTERVAL = random.randint(5, 15)
NB_PINGS = 3

ping_regexp = re.compile(r"time=(.*) ms")

def ping(target, nb_pings=NB_PINGS):
    """Returns the list of latencies towards the given target"""
    args = ["ping", "-n", "-c", str(nb_pings), str(target.address)]
    try:
        out = subprocess.check_output(args, shell=False)
    except subprocess.CalledProcessError:
        return []
    data = out.decode("utf8").split("\n")
    data = [ping_regexp.search(line) for line in data]
    return [float(m.group(1)) for m in data if m != None]

def run_measurements(s):
    while True:
        try:
            data = s.recv(1024)
            print(data)
            if len(data) == 0:
                s.close()
                return
            msg = pf.Message()
            msg.ParseFromString(data)
            if msg.type == pf.Message.Pong:
                print("Got pong")
            else:
                target = msg.target
                latencies = ping(target.target)
                reply = pf.Message()
                reply.type = pf.Message.Measurement
                reply.measurement.target_id = msg.target.target_id
                reply.measurement.probes_sent = NB_PINGS
                reply.measurement.probes_received = len(latencies)
                reply.measurement.latency_us.extend([int(1000 * l) for l in latencies])
                s.send(reply.SerializeToString())
        except socket.timeout:
            print("Sending heartbeat")
            msg = pf.Message()
            msg.type = pf.Message.Ping
            s.send(msg.SerializeToString())
        except ConnectionResetError:
            s.close()
            return


if __name__ == '__main__':
    while True:
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.settimeout(HEARTBEAT_INTERVAL)
        try:
            sock.connect(SERVER)
            print("Connected to {}".format(SERVER))
            run_measurements(sock)
        except ConnectionRefusedError:
            pass
        print("Disconnected, reconnecting in {} seconds".format(RECONNECT_INTERVAL))
        time.sleep(RECONNECT_INTERVAL)