From a985dac4b0c5c85c7d8660bef5b23898f9a28cad Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <ynerant@crans.org>
Date: Fri, 1 Jan 2021 19:55:30 +0100
Subject: [PATCH] Ctrl+C properly stop the program

---
 squinnondation/hazel.py          | 39 ++++++++++++++++++++++++++++++++
 squinnondation/messages.py       |  4 ++--
 squinnondation/squinnondation.py |  4 +---
 3 files changed, 42 insertions(+), 5 deletions(-)

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