diff --git a/Shared/Disc0very.py b/Shared/Disc0very.py index 769aff9..ed3aff9 100644 --- a/Shared/Disc0very.py +++ b/Shared/Disc0very.py @@ -1,13 +1,15 @@ from socket import * from random import getrandbits +from random import choice from time import time from time import sleep import Utils import Logging from _thread import start_new_thread + HEY = "HEY".encode() -IDENT_SEED_LENGTH = 64 +NAY = "NAY".encode() WHAT = "WHAT".encode() @@ -17,80 +19,79 @@ class MissingInterfaceError(TypeError): class Disc0very: - def __init__(self, port, interface="eth0", max_peers=32): + def __init__(self, port, interface=None, max_peers=32): self.port = port - self.ident = getrandbits(max_peers * 2) self.sock = socket(AF_INET, SOCK_DGRAM) self.sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1) self.sock.setsockopt(SOL_SOCKET, SO_BROADCAST, 1) self.sock.bind(("", self.port)) + self.ip_address = Utils.get_ip_address(interface) self.discovering = False self.enlisting = False - - def discover(self, time_out=2): - if not self.enlisting: - self.discovering = True - self.sock.setblocking(False) - self.sock.sendto(HEY, ('255.255.255.255', self.port)) - end_time = time() + time_out - servers = [] - while time() < end_time: - try: - data = self.sock.recv(512) - if data != HEY: - servers.append(data) - except BlockingIOError: - sleep(0.2) - self.discovering = False - if len(servers) == 0: - return None - elif len(servers) > 1: - self.sock.sendto(WHAT, ('255.255.255.255', self.port)) - else: - return servers[0] + # Not thread-safe if ran in parallel with __enlist + def discover(self, time_out=1): + self.sock.setblocking(False) + end_time = time() + time_out + servers = [] + self.sock.sendto(HEY, ('255.255.255.255', self.port)) + while time() < end_time: + try: + data, address = self.sock.recvfrom(512) + address = address[0] + # If enlisted peer responds with ip address + if data not in (HEY, NAY) and address != self.ip_address and not data in servers: + servers.append(data) + # If another peer is currently discovering + elif data == HEY and address != self.ip_address: + self.sock.sendto(NAY, ('255.255.255.255', self.port)) + return self.discover(time_out=time_out) + # If another peer gave up + elif data == NAY and address != self.ip_address: + return None + except BlockingIOError: + sleep((choice(range(1, 10)) / 2) / 10) + if len(servers) == 0: + return None + elif len(servers) > 1: + self.sock.sendto(WHAT, ('255.255.255.255', self.port)) + else: + return servers[0] def enlist(self, interface, blocking=False): - if not self.discovering: - if blocking: - self.__enlist(interface) - else: - start_new_thread(self.__enlist, (interface, )) + if blocking: + self.__enlist(interface) + else: + start_new_thread(self.__enlist, (interface, )) + # Not thread-safe if ran in parallel with discover # Interface should always be provided when using a Wallaby # because wlan0 and wlan1 have an IP address assigned def __enlist(self, interface=None): - self.enlisting = True self.sock.setblocking(True) - ip_address = Utils.get_ip_address(interface) data = "" while True: try: - data = self.sock.recv(512) + data, address = self.sock.recvfrom(512) except BlockingIOError: sleep(0.1) if data == HEY: - self.sock.sendto(ip_address.encode(), ('255.255.255.255', self.port)) + self.sock.sendto(self.ip_address.encode(), ('255.255.255.255', self.port)) elif data == WHAT: Logging.error("Apparently more than one server is running. " "Investigating...") # Discover and if other server is found shutdown -def start(): - print("starting") - - if __name__ == "__main__": - from time import sleep disc0very = Disc0very(3077) server = disc0very.discover() if not server: - print("enlisted") + print("enlisting") disc0very.enlist(None, blocking=True) else: print(server) \ No newline at end of file