diff --git a/src/core/udp_server.py b/src/core/udp_server.py index 2b8a652..a2d5599 100644 --- a/src/core/udp_server.py +++ b/src/core/udp_server.py @@ -10,42 +10,36 @@ import traceback from core import utils -class UDPServer: +class UDPServer(asyncio.DatagramTransport): - def __init__(self, core, host, port): + def __init__(self, core, host=None, port=None): + super().__init__() self.log = utils.get_logger("UDPServer") self.loop = asyncio.get_event_loop() self.Core = core self.host = host self.port = port self.run = False + self.transport = None - async def handle_client(self, reader, writer): - while True: - try: - data = await reader.read(1) - if not data: - break - code = data.decode() - self.log.debug(f"Received {code!r} from {writer.get_extra_info('sockname')!r}") - # await self.handle_code(code, reader, writer) - # task = asyncio.create_task(self.handle_code(code, reader, writer)) - # await asyncio.wait([task], return_when=asyncio.FIRST_EXCEPTION) - if not writer.is_closing(): - writer.close() - self.log.debug("Disconnected.") - break - except Exception as e: - self.log.error("Error while connecting..") - self.log.error(f"Error: {e}") - traceback.print_exc() - break + def connection_made(self, transport): + self.transport = transport + + def datagram_received(self, data, addr): + message = data.decode() + print('Received %r from %s' % (message, addr)) + print('Send %r to %s' % (message, addr)) + self.transport.sendto(data, addr) async def start(self): self.log.debug("Starting UDP server.") self.run = True try: - pass + self.transport, _ = await self.loop.create_datagram_endpoint( + lambda: self, + local_addr=(self.host, self.port), + reuse_port=True + ) except OSError as e: self.log.error("Cannot bind port or other error") raise e @@ -58,3 +52,4 @@ class UDPServer: def stop(self): self.log.debug("Stopping UDP server") + self.transport.close() diff --git a/src/core/udp_server.pyi b/src/core/udp_server.pyi index 8b5d4e6..8666508 100644 --- a/src/core/udp_server.pyi +++ b/src/core/udp_server.pyi @@ -5,20 +5,22 @@ # Licence: FPA # (c) kuitoi.su 2023 import asyncio +from typing import Tuple from core import utils class UDPServer: - def __init__(self, core, host, port): + def __init__(self, core, host=None, port=None): self.log = utils.get_logger("UDPServer") self.loop = asyncio.get_event_loop() self.Core = core self.host = host self.port = port self.run = False - async def handle_client(self, srv_sock) -> None: ... + self.transport = None + def connection_made(self, transport: asyncio.DatagramTransport): ... + def datagram_received(self, data: bytes, addr: Tuple[str, int]): ... async def start(self) -> None: ... - async def stop(self) -> None: ... \ No newline at end of file