diff --git a/squinnondation/hazel.py b/squinnondation/hazel.py index 4a1da53..958fa2d 100644 --- a/squinnondation/hazel.py +++ b/squinnondation/hazel.py @@ -193,9 +193,7 @@ class Squirrel(Hazelnut): 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.") + return Packet.construct(), hazelnut try: pkt = Packet.unmarshal(data) except ValueError as error: @@ -203,6 +201,10 @@ class Squirrel(Hazelnut): hazelnut.errors += 1 self.send_packet(hazelnut, Packet.construct(WarningTLV.construct( f"An error occured while reading your packet: {error}"))) + if hazelnut.banned: + self.send_packet(hazelnut, Packet.construct(WarningTLV.construct( + "You got banned since you sent too much errored packets."))) + raise ValueError("Client is banned since there were too many errors.", error) raise error else: return pkt, hazelnut @@ -761,7 +763,13 @@ class Worm(Thread): pkt, hazelnut = self.squirrel.receive_packet() except ValueError as error: self.squirrel.add_system_message("An error occurred while receiving a packet: {}".format(error)) + self.squirrel.refresh_history() + self.squirrel.refresh_input() else: + if hazelnut.banned: + # Ignore banned hazelnuts + continue + for tlv in pkt.body: tlv.handle(self.squirrel, hazelnut) self.squirrel.refresh_history()