From 0a4ebd0c921b0a2f4e01d8921b42d128488d8836 Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <ynerant@crans.org>
Date: Tue, 5 Jan 2021 20:38:19 +0100
Subject: [PATCH] An unknown client can only send HelloTLV and GoAwayTLV for
 security

---
 squinnondation/hazel.py    |  6 +++---
 squinnondation/messages.py | 40 +++++++++++++++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py
index 24fd6d2..ed79fd6 100644
--- a/squinnondation/hazel.py
+++ b/squinnondation/hazel.py
@@ -51,7 +51,7 @@ class Hazelnut:
         return self.addresses[0]
 
     def __repr__(self):
-        return f"{self.id or '<unknown id>'}: {self.nickname or '<unknown nickname>'}, {self.addresses}"
+        return self.nickname or str(self.id) or str(self.main_address)
 
     def __str__(self):
         return repr(self)
@@ -370,10 +370,10 @@ class Squirrel(Hazelnut):
         """
         res = dict()
         hazels = self.active_hazelnuts
-        for key, hazel in hazels:
+        for hazel in hazels:
             if hazel.symmetric:
                 next_send = uniform(1, 2)
-                res[key] = [hazel, time.time() + next_send, 0]
+                res[hazel.main_address] = [hazel, time.time() + next_send, 0]
         return res
 
     def remove_from_inundation(self, hazel: Hazelnut, sender_id: int, nonce: int) -> None:
diff --git a/squinnondation/messages.py b/squinnondation/messages.py
index 5145560..f8948c9 100644
--- a/squinnondation/messages.py
+++ b/squinnondation/messages.py
@@ -73,7 +73,12 @@ class Pad1TLV(TLV):
         return self.type.to_bytes(1, sys.byteorder)
 
     def handle(self, squirrel: Any, sender: Any) -> None:
-        # TODO Add some easter eggs
+        if not sender.active or not sender.id:
+            # It doesn't say hello, we don't listen to it
+            squirrel.send_packet(sender, Packet.construct(WarningTLV.construct(
+                "You are not my neighbour, I don't listen to your Pad1TLV. Please say me Hello before.")))
+            return
+
         squirrel.add_system_message("I received a Pad1TLV, how disapointing.")
 
     def __len__(self) -> int:
@@ -118,7 +123,12 @@ class PadNTLV(TLV):
             + self.mbz[:self.length]
 
     def handle(self, squirrel: Any, sender: Any) -> None:
-        # TODO Add some easter eggs
+        if not sender.active or not sender.id:
+            # It doesn't say hello, we don't listen to it
+            squirrel.send_packet(sender, Packet.construct(WarningTLV.construct(
+                "You are not my neighbour, I don't listen to your PadNTLV. Please say me Hello before.")))
+            return
+
         squirrel.add_system_message(f"I received {self.length} zeros.")
 
     @staticmethod
@@ -229,6 +239,12 @@ class NeighbourTLV(TLV):
             self.port.to_bytes(2, sys.byteorder)
 
     def handle(self, squirrel: Any, sender: Any) -> None:
+        if not sender.active or not sender.id:
+            # It doesn't say hello, we don't listen to it
+            squirrel.send_packet(sender, Packet.construct(WarningTLV.construct(
+                "You are not my neighbour, I don't listen to your NeighbourTLV. Please say me Hello before.")))
+            return
+
         if (self.ip_address, self.port) in squirrel.addresses:
             # This case should never happen (and in our protocol it is not possible),
             # but we include this test as a security measure.
@@ -279,6 +295,12 @@ class DataTLV(TLV):
         """
         A message has been sent. We log it.
         """
+        if not sender.active or not sender.id:
+            # It doesn't say hello, we don't listen to it
+            squirrel.send_packet(sender, Packet.construct(WarningTLV.construct(
+                "You are not my neighbour, I don't listen to your DataTLV. Please say me Hello before.")))
+            return
+
         msg = self.data.decode('UTF-8')
 
         # Acknowledge the packet
@@ -342,7 +364,13 @@ class AckTLV(TLV):
         """
         When an AckTLV is received, we know that we do not have to inundate that neighbour anymore.
         """
-        squirrel.add_system_message("I received an AckTLV")
+        if not sender.active or not sender.id:
+            # It doesn't say hello, we don't listen to it
+            squirrel.send_packet(sender, Packet.construct(WarningTLV.construct(
+                "You are not my neighbour, I don't listen to your AckTLV. Please say me Hello before.")))
+            return
+
+        squirrel.add_system_message(f"I received an AckTLV from {sender}")
         squirrel.remove_from_inundation(sender, self.sender_id, self.nonce)
 
     @staticmethod
@@ -381,6 +409,12 @@ class GoAwayTLV(TLV):
             self.message.encode("UTF-8")[:self.length - 1]
 
     def handle(self, squirrel: Any, sender: Any) -> None:
+        if not sender.active or not sender.id:
+            # It doesn't say hello, we don't listen to it
+            squirrel.send_packet(sender, Packet.construct(WarningTLV.construct(
+                "You are not my neighbour, I don't listen to your GoAwayTLV. Please say me Hello before.")))
+            return
+
         if sender.active:
             sender.active = False
             squirrel.update_hazelnut_table(sender)