[proxy] Switch to AbstractUDPProxy

dmrtools/udpproxy.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from abc import ABC, abstractmethod
+
+
+class AbstractUDPProxy(ABC):
+    def __init__(self, server_host: str, server_port: int,
+                 listen_host: str, listen_port: int) -> None:
+        self.listen_host: str = listen_host
+        self.listen_port: int = listen_port
+        self.server_host: str = server_host
+        self.server_port: int = server_port
+
+        self.transport_client: asyncio.DatagramTransport|None = None
+        self.transport_server: asyncio.DatagramTransport|None = None
+
+        self.client_address: tuple|None = None
+
+    def on_forward(self, data: bytes, to_server: bool) -> bytes:
+        """
+        Intercept or modify packet data before forwarding.
+        to_server: True means direction from client to server, False otherwise
+        """
+        return data
+
+    async def start(self) -> None:
+        loop = asyncio.get_running_loop()
+
+        # Start the listening endpoint for the client
+        self.transport_client, _ = await loop.create_datagram_endpoint(
+            lambda: self.ClientProtocol(self),
+            local_addr=(self.listen_host, self.listen_port)
+        )
+
+        logging.info(
+            f"Listening for on {self.listen_host}:{self.listen_port}")
+
+    def reset(self) -> None:
+        if self.transport_client:
+            self.transport_client.close()
+        self.transport_client = None
+
+        if self.transport_server:
+            self.transport_server.close()
+        self.transport_server = None
+
+        self.client_address = None
+
+        logging.debug("Proxy reset to LISTENING state")
+
+    async def connect_to_dest(self):
+        loop = asyncio.get_running_loop()
+
+        self.transport_server, _ = await loop.create_datagram_endpoint(
+            lambda: self.ServerProtocol(self),
+            remote_addr=(self.server_host, self.server_port)
+        )
+
+        logging.info("Proxying "
+              f"{self.client_address[0]}:{self.client_address[1]}"
+              f" <=> {self.server_host}:{self.server_port}")
+
+    class ClientProtocol(asyncio.DatagramProtocol):
+        def __init__(self, proxy: AbstractUDPProxy) -> None:
+            self.proxy: AbstractUDPProxy = proxy
+
+        def datagram_received(self, data: bytes, addr: tuple) -> None:
+            if self.proxy.client_address is None:
+                self.proxy.client_address = addr
+                logging.info(f"Client {addr[0]}:{addr[1]} connected")
+                asyncio.create_task(self.proxy.connect_to_dest())
+
+            if addr != self.proxy.client_address:
+                logging.error("Ignoring packet from unknown client"
+                      f" {addr[0]}:{addr[1]}")
+                return
+
+            if self.proxy.transport_server:
+                try:
+                    data = self.proxy.on_forward(data, to_server=True)
+                    self.proxy.transport_server.sendto(data)
+                except Exception as e:
+                    logging.error(f"Error forwarding to server: {e}")
+                    self.proxy.reset()
+
+        def error_received(self, exc: Exception|None) -> None:
+            logging.error(f"Client socket error: {exc}")
+            self.proxy.reset()
+
+        def connection_lost(self, exc: Exception|None) -> None:
+            logging.error(f"Client connection lost: {exc}")
+            self.proxy.reset()
+
+    class ServerProtocol(asyncio.DatagramProtocol):
+        def __init__(self, proxy: AbstractUDPProxy) -> None:
+            self.proxy: AbstractUDPProxy = proxy
+
+        def datagram_received(self, data: bytes, addr: tuple) -> None:
+            if self.proxy.transport_client and self.proxy.client_address:
+                try:
+                    data = self.proxy.on_forward(data, to_server=False)
+                    self.proxy.transport_client.sendto(
+                        data, self.proxy.client_address)
+                except Exception as e:
+                    logging.error(f"Error forwarding to client: {e}")
+                    self.proxy.reset()
+
+        def error_received(self, exc: Exception|None) -> None:
+            logging.error(f"Server socket error: {exc}")
+            self.proxy.reset()
+
+        def connection_lost(self, exc: Exception|None) -> None:
+            logging.error(f"Server connection lost: {exc}")
+            self.proxy.reset()