From 5bb4748056823d8cba1076dc85a149bc369dd49c Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <ynerant@crans.org>
Date: Tue, 5 Jan 2021 20:54:14 +0100
Subject: [PATCH] Warn users when they send poop messages and ban them if they
 are too many

---
 squinnondation/hazel.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py
index ed79fd6..2a4a5ab 100644
--- a/squinnondation/hazel.py
+++ b/squinnondation/hazel.py
@@ -11,7 +11,7 @@ import re
 import socket
 import time
 
-from .messages import Packet, DataTLV, HelloTLV, GoAwayTLV, GoAwayType, NeighbourTLV
+from .messages import Packet, DataTLV, HelloTLV, GoAwayTLV, GoAwayType, NeighbourTLV, WarningTLV
 
 
 class Hazelnut:
@@ -25,6 +25,7 @@ class Hazelnut:
         self.last_long_hello_time = 0
         self.symmetric = False
         self.active = False
+        self.errors = 0
 
         try:
             # Resolve DNS as an IPv6
@@ -48,8 +49,19 @@ class Hazelnut:
 
     @property
     def main_address(self) -> Tuple[str, int]:
+        """
+        A client can have multiple addresses.
+        We contact it only on one of them.
+        """
         return self.addresses[0]
 
+    @property
+    def banned(self) -> bool:
+        """
+        If a client send more than 5 invalid packets, we don't trust it anymore.
+        """
+        return self.errors >= 5
+
     def __repr__(self):
         return self.nickname or str(self.id) or str(self.main_address)
 
@@ -165,7 +177,22 @@ class Squirrel(Hazelnut):
         Warning: the process is blocking, it should be ran inside a dedicated thread.
         """
         data, addr = self.receive_raw_data()
-        return Packet.unmarshal(data), self.find_hazelnut(addr[0], addr[1])
+        hazelnut = self.find_hazelnut(addr[0], addr[1])
+        if hazelnut.banned:
+            # The client already sent errored packets
+            self.send_packet(hazelnut, Packet.construct(WarningTLV.construct(
+                "You got banned since you sent too much errored packets.")))
+            raise ValueError("Client is banned.")
+        try:
+            pkt = Packet.unmarshal(data)
+        except ValueError as error:
+            # The packet contains an error. We memorize it and warn the other user.
+            hazelnut.errors += 1
+            self.send_packet(hazelnut, Packet.construct(WarningTLV.construct(
+                f"An error occured while reading your packet: {error}")))
+            raise error
+        else:
+            return pkt, hazelnut
 
     def receive_raw_data(self) -> Tuple[bytes, Any]:
         """