diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py index 00988ca..8ebec48 100644 --- a/squinnondation/hazel.py +++ b/squinnondation/hazel.py @@ -144,6 +144,23 @@ class Squirrel(Hazelnut): """ return self.socket.recvfrom(1024) + def start_threads(self) -> None: + """ + Start asynchronous threads. + """ + self.worm = Worm(self) + self.hazel_manager = HazelManager(self) + self.inondator = Inondator(self) + + # Kill subthreads when exitting the program + self.worm.setDaemon(True) + self.hazel_manager.setDaemon(True) + self.inondator.setDaemon(True) + + self.worm.start() + self.hazel_manager.start() + self.inondator.start() + def wait_for_key(self) -> None: """ Infinite loop where we are waiting for a key of the user. @@ -158,6 +175,10 @@ class Squirrel(Hazelnut): curses.LINES - 1, min(3 + len(self.nickname) + self.input_index, curses.COLS - 4)) except curses.error: continue + except KeyboardInterrupt: + # Exit the program and send GoAway to neighbours + self.leave() + return if key == curses.KEY_MOUSE: try: @@ -655,6 +676,24 @@ class Squirrel(Hazelnut): self.refresh_lock.release() + def leave(self) -> None: + """ + The program is exited. We send a GoAway to our neighbours, then close the program. + """ + self.refresh_lock.acquire() + + # Last inundation + self.main_inundation() + self.clean_inundation() + + # Broadcast a GoAway + gatlv = GoAwayTLV().construct(GoAwayType.EXIT, f"I am leaving! Good bye!") + pkt = Packet.construct(gatlv) + for hazelnut in self.activehazelnuts.values(): + self.send_packet(hazelnut[0], pkt) + + exit(0) + class Worm(Thread): """ diff --git a/squinnondation/messages.py b/squinnondation/messages.py index a5304f7..2cf0fb4 100644 --- a/squinnondation/messages.py +++ b/squinnondation/messages.py @@ -356,7 +356,7 @@ class GoAwayTLV(TLV): self.type = raw_data[0] self.length = raw_data[1] self.code = GoAwayType(raw_data[2]) - self.message = raw_data[3:self.length - 1].decode("UTF-8") + self.message = raw_data[3:2 + self.length].decode("UTF-8") def marshal(self) -> bytes: return self.type.to_bytes(1, sys.byteorder) + \ @@ -376,7 +376,7 @@ class GoAwayTLV(TLV): tlv.type = 6 tlv.code = ga_type tlv.message = message - tlv.length = 1 + len(tlv.message) + tlv.length = 1 + len(tlv.message.encode("UTF-8")) return tlv diff --git a/squinnondation/squinnondation.py b/squinnondation/squinnondation.py index b0d959d..e81d3c3 100644 --- a/squinnondation/squinnondation.py +++ b/squinnondation/squinnondation.py @@ -86,7 +86,5 @@ class Squinnondation: # hazelnut = Hazelnut(address='::1', port=8082) # squirrel.potentialhazelnuts['::1', 8082] = hazelnut - Worm(squirrel).start() - HazelManager(squirrel).start() - Inondator(squirrel).start() + squirrel.start_threads() squirrel.wait_for_key()