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"""
    if target.family == pf.IPAddress.IPV4:
        ping = "ping"
    elif target.family == pf.IPAddress.IPV6:
        ping = "ping6"
    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])
                # Debug
                #time.sleep(random.randint(1, 5))
                print("Sending results")
                s.send(reply.SerializeToString())
        except socket.timeout:
            pass
            #print("Not 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)