diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py index 5bd436d..24fd6d2 100644 --- a/squinnondation/hazel.py +++ b/squinnondation/hazel.py @@ -115,12 +115,12 @@ class Squirrel(Hazelnut): return hazelnut @property - def active_hazelnuts(self) -> dict: - return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.active} + def active_hazelnuts(self) -> set: + return set(hazelnut for hazelnut in self.hazelnuts.values() if hazelnut.active) @property - def potential_hazelnuts(self) -> dict: - return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.potential} + def potential_hazelnuts(self) -> set: + return set(hazelnut for hazelnut in self.hazelnuts.values() if hazelnut.potential) def find_hazelnut(self, address: str, port: int) -> Hazelnut: """ @@ -284,7 +284,7 @@ class Squirrel(Hazelnut): self.add_message(msg) pkt = Packet.construct(DataTLV.construct(msg, self)) - for hazelnut in list(self.active_hazelnuts.values()): + for hazelnut in self.active_hazelnuts: self.send_packet(hazelnut, pkt) def handle_mouse_click(self, y: int, x: int, attr: int) -> None: @@ -369,7 +369,7 @@ class Squirrel(Hazelnut): Takes the active hazels dictionnary and returns a list of [hazel, date+random, 0] """ res = dict() - hazels = list(self.active_hazelnuts.items()) + hazels = self.active_hazelnuts for key, hazel in hazels: if hazel.symmetric: next_send = uniform(1, 2) @@ -622,8 +622,8 @@ class Squirrel(Hazelnut): not have enough symmetric neighbours. """ res = [] - lp = len(self.potential_hazelnuts) - val = list(self.potential_hazelnuts.values()) + val = list(self.potential_hazelnuts) + lp = len(val) for i in range(min(lp, max(0, self.minNS - self.nbNS))): a = randint(0, lp - 1) @@ -634,7 +634,7 @@ class Squirrel(Hazelnut): """ Sends a long HelloTLV to all active neighbours. """ - for hazelnut in self.active_hazelnuts.values(): + for hazelnut in self.active_hazelnuts: htlv = HelloTLV().construct(16, self, hazelnut) pkt = Packet().construct(htlv) self.send_packet(hazelnut, pkt) @@ -644,7 +644,7 @@ class Squirrel(Hazelnut): All neighbours that have not sent a HelloTLV in the last 2 minutes are considered not active. """ - val = list(self.active_hazelnuts.values()) # create a copy because the dict size will change + val = list(self.active_hazelnuts) # create a copy because the dict size will change for hazelnut in val: if time.time() - hazelnut.last_hello_time > 2 * 60: @@ -681,17 +681,17 @@ class Squirrel(Hazelnut): nb_ns = 0 # could send the same to all neighbour, but it means that neighbour # A could receive a message with itself in it -> if the others do not pay attention, trouble - for key, hazelnut in self.active_hazelnuts.items(): + for hazelnut in self.active_hazelnuts: if time.time() - hazelnut.last_long_hello_time <= 2 * 60: nb_ns += 1 - self.hazelnuts[key].symmetric = True + hazelnut.symmetric = True ntlv = NeighbourTLV().construct(*hazelnut.main_address) pkt = Packet().construct(ntlv) - for destination in self.active_hazelnuts.values(): + for destination in self.active_hazelnuts: if destination.id != hazelnut.id: self.send_packet(destination, pkt) else: - self.active_hazelnuts[key].symmetric = False + hazelnut.symmetric = False self.nbNS = nb_ns def leave(self) -> None: @@ -705,7 +705,7 @@ class Squirrel(Hazelnut): # Broadcast a GoAway gatlv = GoAwayTLV().construct(GoAwayType.EXIT, "I am leaving! Good bye!") pkt = Packet.construct(gatlv) - for hazelnut in self.active_hazelnuts.values(): + for hazelnut in self.active_hazelnuts: self.send_packet(hazelnut, pkt) exit(0) @@ -762,7 +762,7 @@ class HazelManager(Thread): # Second part: send long HelloTLVs to neighbours every 30 seconds if time.time() - self.last_check > 30: - self.squirrel.add_system_message(f"I have {len(list(self.squirrel.active_hazelnuts))} friends") + self.squirrel.add_system_message(f"I have {len(self.squirrel.active_hazelnuts)} friends") self.squirrel.send_hello() self.last_check = time.time() diff --git a/squinnondation/messages.py b/squinnondation/messages.py index 7ce7d1a..5145560 100644 --- a/squinnondation/messages.py +++ b/squinnondation/messages.py @@ -233,8 +233,7 @@ class NeighbourTLV(TLV): # This case should never happen (and in our protocol it is not possible), # but we include this test as a security measure. return - if not (str(self.ip_address), self.port) in squirrel.active_hazelnuts \ - and not (str(self.ip_address), self.port) in squirrel.potential_hazelnuts: + if not (str(self.ip_address), self.port) in squirrel.hazelnuts: hazelnut = squirrel.new_hazel(str(self.ip_address), self.port) hazelnut.potential = True squirrel.update_hazelnut_table(hazelnut)