[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()

udpproxylogger.py

@@ -3,11 +3,10 @@ import asyncio
 import logging
 import traceback
 
-from bitarray import bitarray
+from dmrtools.udpproxy import AbstractUDPProxy
-
 from dmrtools import DMRPPacketFactory
 from dmrtools import hexdump
-from dmrtools.dmrproto import DMRPL2FullLC, DMRPL2VoiceBurst
+# from dmrtools.dmrproto import DMRPL2FullLC, DMRPL2VoiceBurst
 from dmrtools.dmrproto import DMRPPacketData
 from dmrtools.dmrproto import EmbLCAssembler, LCFactory
 
@@ -52,95 +51,33 @@ def get_packet_details(data):
         return f"Exception while decoding packet:\n{traceback.format_exc()}"
 
 
-def log_packet(direction, data, addr):
+def log_packet(direction, data):
     log_message = (
-        f"{direction} {addr[0]}:{addr[1]} | {len(data)} bytes\n"
+        f"{direction} | {len(data)} bytes\n"
         f"HD: {data.hex()}\n{hexdump(data)}\n"
         f"{get_packet_details(data)}"
     )
     logging.info(log_message)
 
 
-class UDPProxyClientProtocol(asyncio.DatagramProtocol):
+class UDPProxyLogger(AbstractUDPProxy):
-    def __init__(self, server_address, loop):
+    def on_forward(self, data: bytes, to_server: bool) -> bytes:
-        self.server_address = server_address
+        direction = "OUT=>" if to_server else "<==IN"
-        self.loop = loop
+        log_packet(direction, data)
-        self.transport = None
+        return data
-        self.server_transport = None
-        self.client_addr = None
-
-    def connection_made(self, transport):
-        self.transport = transport
-        logging.info(f"Listening for client datagrams")
-
-    def datagram_received(self, data, addr):
-        if not self.client_addr:
-            self.client_addr = addr
-            logging.info(f"Client connected from {self.client_addr}")
-
-        log_packet("OUT=>", data, self.server_address)
-
-        # Forward to server
-        self.server_transport.sendto(data)
-
-    def error_received(self, exc):
-        logging.error(f"Client protocol error: {exc}")
-
-    def connection_lost(self, exc):
-        logging.info("Client connection closed")
-
-class UDPProxyServerProtocol(asyncio.DatagramProtocol):
-    def __init__(self, client_protocol):
-        self.client_protocol = client_protocol
-        self.transport = None
-
-    def connection_made(self, transport):
-        self.transport = transport
-        self.client_protocol.server_transport = transport
-        peername = self.transport.get_extra_info('peername')
-        logging.info(f"Connected to server {peername}")
-
-    def datagram_received(self, data, addr):
-        if self.client_protocol.client_addr:
-            log_packet("<==IN", data, self.client_protocol.client_addr)
-
-            # Send data to client
-            self.client_protocol.transport.sendto(
-                data, self.client_protocol.client_addr)
-
-    def error_received(self, exc):
-        logging.error(f"Server protocol error: {exc}")
-
-    def connection_lost(self, exc):
-        logging.info("Server connection closed")
 
 
-async def start_udp_proxy(listen_ip, listen_port, server_ip, server_port):
+async def start_udp_proxy(server_ip, server_port, listen_ip, listen_port):
-    loop = asyncio.get_running_loop()
-
-    server_address = (server_ip, server_port)
-
-    # Create client listener
-    listen = await loop.create_datagram_endpoint(
-        lambda: UDPProxyClientProtocol(server_address, loop),
-        local_addr=(listen_ip, listen_port)
-    )
-    client_transport, client_protocol = listen
-
-    # Connect to server
-    connect = await loop.create_datagram_endpoint(
-        lambda: UDPProxyServerProtocol(client_protocol),
-        remote_addr=server_address
-    )
-    server_transport, server_protocol = connect
-
     logging.info(f"Proxy running: {listen_ip}:{listen_port} <=> {server_ip}:{server_port}")
 
+    proxy = UDPProxyLogger(server_ip, server_port, listen_ip, listen_port)
+    await proxy.start()
+
     try:
         await asyncio.Future()  # Run forever
     finally:
-        client_transport.close()
+        # await proxy.stop()
-        server_transport.close()
+        pass
 
 
 def parse_arguments():
@@ -158,8 +95,8 @@ def main():
     setup_logger(args.log_file)
 
     try:
-        asyncio.run(start_udp_proxy(args.listenip, args.listenport,
+        asyncio.run(start_udp_proxy(args.serverip, args.serverport,
-                                    args.serverip, args.serverport))
+                                    args.listenip, args.listenport))
     except KeyboardInterrupt:
         logging.info("Shutting down.")