Use hazelnut sets instead of dictionaries to avoid duplicate hazels

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()
 

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)