diff --git a/squinnondation/squinnondation.py b/squinnondation/squinnondation.py index d99873f..4ea5c5f 100644 --- a/squinnondation/squinnondation.py +++ b/squinnondation/squinnondation.py @@ -5,6 +5,7 @@ import socket from argparse import ArgumentParser from enum import Enum from ipaddress import IPv6Address +from threading import Thread from typing import Any, Optional, Tuple @@ -58,20 +59,7 @@ class Squinnondation: pkt.body_length = pkt.body.length + 2 squirrel.send_packet(hazelnut, pkt) - while True: - pkt, hazelnut = squirrel.receive_packet() - print(pkt.body.data.decode('UTF-8')) - pkt = Packet() - pkt.magic = 95 - pkt.version = 0 - pkt.body = DataTLV() - msg = f"Hello my dear hazelnut, I am {squirrel.nickname}!" - pkt.body.data = msg.encode("UTF-8") - pkt.body.sender_id = 42 - pkt.body.nonce = 18 - pkt.body.length = len(msg) + 1 + 1 + 8 + 4 - pkt.body_length = pkt.body.length + 2 - squirrel.send_packet(hazelnut, pkt) + Worm(squirrel).start() class TLV: @@ -393,7 +381,7 @@ class Squirrel(Hazelnut): def receive_packet(self) -> Tuple[Packet, Hazelnut]: """ Receive a packet from the socket and translate it into a Python object. - TODO: Translate the address into the correct 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]) @@ -403,3 +391,34 @@ class Squirrel(Hazelnut): Receive a packet from the socket. """ return self.socket.recvfrom(1024) + + +class Worm(Thread): + """ + The worm is the hazel listener. + It always waits for an incoming packet, then it treats it, and continues to wait. + It is in a dedicated thread. + """ + def __init__(self, squirrel: Squirrel, *args, **kwargs): + super().__init__(*args, **kwargs) + self.squirrel = squirrel + + def run(self) -> None: + while True: + try: + pkt, hazelnut = self.squirrel.receive_packet() + except ValueError as error: + print("An error occured while receiving a packet: ", error) + else: + print(pkt.body.data.decode('UTF-8')) + pkt = Packet() + pkt.magic = 95 + pkt.version = 0 + pkt.body = DataTLV() + msg = f"Hello my dear hazelnut, I am {self.squirrel.nickname}!" + pkt.body.data = msg.encode("UTF-8") + pkt.body.sender_id = 42 + pkt.body.nonce = 18 + pkt.body.length = len(msg) + 1 + 1 + 8 + 4 + pkt.body_length = pkt.body.length + 2 + self.squirrel.send_packet(hazelnut, pkt)