diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py index e258d90..a09f1c0 100644 --- a/squinnondation/hazel.py +++ b/squinnondation/hazel.py @@ -24,6 +24,8 @@ class Hazelnut: self.last_hello_time = 0 self.last_long_hello_time = 0 self.symmetric = False + self.active = False + self.potential = False try: # Resolve DNS as an IPv6 @@ -77,8 +79,7 @@ class Squirrel(Hazelnut): curses.init_pair(i + 1, i, curses.COLOR_BLACK) # dictionnaries of neighbours - self.potentialhazelnuts = dict() - self.activehazelnuts = dict() + self.hazelnuts = dict() self.nbNS = 0 self.minNS = 3 # minimal number of symmetric neighbours a squirrel needs to have. @@ -96,22 +97,21 @@ class Squirrel(Hazelnut): hazelnut = Hazelnut(address=address, port=port) return hazelnut - def is_active(self, hazel: Hazelnut) -> bool: - return (hazel.address, hazel.port) in self.activehazelnuts + @property + def activehazelnuts(self) -> dict: + return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.active} - def is_potential(self, hazel: Hazelnut) -> bool: - return (hazel.address, hazel.port) in self.potentialhazelnuts - - def remove_from_potential(self, address: str, port: int) -> None: - self.potentialhazelnuts.pop((address, port), None) + @property + def potentialhazelnuts(self) -> dict: + return {(address, port): hazelnut for (address, port), hazelnut in self.hazelnuts.items() if hazelnut.potential} def find_hazelnut(self, address: str, port: int) -> Hazelnut: """ Translate an address into a hazelnut. If this hazelnut does not exist, creates a new hazelnut. """ - if (address, port) in self.activehazelnuts: - return self.activehazelnuts[(address, port)] + if (address, port) in self.hazelnuts: + return self.hazelnuts[(address, port)] hazelnut = Hazelnut(address=address, port=port) return hazelnut @@ -402,9 +402,8 @@ class Squirrel(Hazelnut): if self.recent_messages[key][2][key2][2] >= 5: # the neighbour is not reactive enough gatlv = GoAwayTLV().construct(GoAwayType.TIMEOUT, f"{self.id} No acknowledge") pkt = Packet().construct(gatlv) - self.send_packet(self.recent_messages[key][2][key2][0], pkt) - self.activehazelnuts.pop(key2) - self.potentialhazelnuts[key2] = self.recent_messages[key][2][key2][0] + hazelnut = self.recent_messages[key][2][key2][0] + self.send_packet(hazelnut, pkt) self.recent_messages[key][2].pop(key2) self.refresh_lock.release() @@ -648,11 +647,15 @@ class Squirrel(Hazelnut): gatlv = GoAwayTLV().construct(GoAwayType.TIMEOUT, "you did not talk to me") pkt = Packet().construct(gatlv) self.send_packet(hazelnut, pkt) - self.activehazelnuts.pop((hazelnut.address, hazelnut.port)) - self.potentialhazelnuts[(hazelnut.address, hazelnut.port)] = hazelnut + hazelnut.active = False + hazelnut.potential = True + self.update_hazelnut_table(hazelnut) self.refresh_lock.release() + def update_hazelnut_table(self, hazelnut: Hazelnut) -> None: + self.hazelnuts[(hazelnut.address, hazelnut.port)] = hazelnut + def send_neighbours(self) -> None: """ Update the number of symmetric neighbours and @@ -666,7 +669,7 @@ class Squirrel(Hazelnut): for key, hazelnut in self.activehazelnuts.items(): if time.time() - hazelnut.last_long_hello_time <= 2 * 60: nb_ns += 1 - self.activehazelnuts[key].symmetric = True + self.hazelnuts[key].symmetric = True ntlv = NeighbourTLV().construct(hazelnut.address, hazelnut.port) pkt = Packet().construct(ntlv) for destination in self.activehazelnuts.values(): @@ -748,7 +751,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.activehazelnuts.values()))} friends") + self.squirrel.add_system_message(f"I have {len(list(self.squirrel.activehazelnuts))} friends") self.squirrel.send_hello() self.last_check = time.time() diff --git a/squinnondation/messages.py b/squinnondation/messages.py index 661e6d2..531f3c9 100644 --- a/squinnondation/messages.py +++ b/squinnondation/messages.py @@ -159,7 +159,7 @@ class HelloTLV(TLV): def handle(self, squirrel: Any, sender: Any) -> None: time_h = time.time() - if not squirrel.is_active(sender): + if not sender.active: sender.id = self.source_id # The sender we are given misses an id time_hl = time.time() else: @@ -167,14 +167,13 @@ class HelloTLV(TLV): if self.is_long and self.dest_id == squirrel.id: time_hl = time.time() - # Make sure the sender is not in the potential hazelnuts - squirrel.remove_from_potential(sender.address, sender.port) - # Add entry to/actualize the active hazelnuts dictionnary sender.last_hello_time = time_h sender.last_long_hello_time = time_hl sender.symmetric = True - squirrel.activehazelnuts[(sender.address, sender.port)] = sender + sender.potential = False + sender.active = True + squirrel.update_hazelnut_table(sender) squirrel.nbNS += 1 squirrel.add_system_message(f"{self.source_id} sent me a Hello " + ("long" if self.is_long else "short")) @@ -374,9 +373,10 @@ class GoAwayTLV(TLV): self.message.encode("UTF-8")[:self.length - 1] def handle(self, squirrel: Any, sender: Any) -> None: - if squirrel.is_active(sender): - squirrel.activehazelnuts.pop((sender.address, sender.port)) - squirrel.potentialhazelnuts[(sender.address, sender.port)] = sender + if sender.active: + sender.active = False + sender.potential = True + squirrel.update_hazelnut_table(sender) squirrel.add_system_message("Some told me that he went away : " + self.message) @staticmethod