[proxy] multi-client proxy

dmrtools/udpproxy.py

@@ -3,23 +3,22 @@ from __future__ import annotations
 import asyncio
 import logging
 
-from abc import ABC, abstractmethod
+from typing import Optional
 
 
-class AbstractUDPProxy(ABC):
+class UDPProxy:
     def __init__(self, server_host: str, server_port: int,
                  listen_host: str, listen_port: int) -> None:
-        self.listen_host: str = listen_host
+        self.listen_host = listen_host
-        self.listen_port: int = listen_port
+        self.listen_port = listen_port
-        self.server_host: str = server_host
+        self.server_host = server_host
-        self.server_port: int = server_port
+        self.server_port = server_port
 
-        self.transport_client: asyncio.DatagramTransport|None = None
+        self.listener_transport: asyncio.DatagramTransport|None = None
-        self.transport_server: asyncio.DatagramTransport|None = None
+        self.sessions: dict[tuple[str, int], UDPProxy.Session] = {}
 
-        self.client_address: tuple|None = None
+    def on_forward(self, data: bytes, to_server: bool,
-
+                   client_addr: tuple[str, int]) -> bytes:
-    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
@@ -29,89 +28,107 @@ class AbstractUDPProxy(ABC):
     async def start(self) -> None:
         loop = asyncio.get_running_loop()
 
-        # Start the listening endpoint for the client
+        self.listener_transport, _ = await loop.create_datagram_endpoint(
-        self.transport_client, _ = await loop.create_datagram_endpoint(
+            lambda: self.ListenerProtocol(self),
-            lambda: self.ClientProtocol(self),
             local_addr=(self.listen_host, self.listen_port)
         )
 
-        logging.info(
+        logging.info(f"Listening for clients on {self.listen_host}:{self.listen_port}")
-            f"Listening for on {self.listen_host}:{self.listen_port}")
 
-    def reset(self) -> None:
+    async def _create_session(self, client_addr: tuple[str, int]) -> Session:
-        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(
+        transport, _ = await loop.create_datagram_endpoint(
-            lambda: self.ServerProtocol(self),
+            lambda: self.ServerProtocol(self, client_addr),
             remote_addr=(self.server_host, self.server_port)
         )
 
-        logging.info("Proxying "
+        session = UDPProxy.Session(proxy=self,
-              f"{self.client_address[0]}:{self.client_address[1]}"
+                                           client_addr=client_addr,
+                                           server_transport=transport)
+
+        self.sessions[client_addr] = session
+
+        logging.info(f"Proxying {client_addr[0]}:{client_addr[1]}"
                      f" <=> {self.server_host}:{self.server_port}")
 
-    class ClientProtocol(asyncio.DatagramProtocol):
+        return session
-        def __init__(self, proxy: AbstractUDPProxy) -> None:
-            self.proxy: AbstractUDPProxy = proxy
 
-        def datagram_received(self, data: bytes, addr: tuple) -> None:
+    def remove_session(self, client_addr: tuple[str, int]) -> None:
-            if self.proxy.client_address is None:
+        session = self.sessions.pop(client_addr, None)
-                self.proxy.client_address = addr
+        if session:
-                logging.info(f"Client {addr[0]}:{addr[1]} connected")
+            session.server_transport.close()
-                asyncio.create_task(self.proxy.connect_to_dest())
+            logging.info(f"Closed session for {client_addr}")
 
-            if addr != self.proxy.client_address:
+    class Session:
-                logging.error("Ignoring packet from unknown client"
+        def __init__(self, proxy: UDPProxy, client_addr: tuple[str, int],
-                      f" {addr[0]}:{addr[1]}")
+                     server_transport: asyncio.DatagramTransport) -> None:
-                return
+            self.proxy = proxy
+            self.client_addr = client_addr
+            self.server_transport = server_transport
 
-            if self.proxy.transport_server:
+        def handle_from_client(self, data: bytes) -> None:
             try:
-                    data = self.proxy.on_forward(data, to_server=True)
+                data = self.proxy.on_forward(data, to_server=True,
-                    self.proxy.transport_server.sendto(data)
+                                             client_addr=self.client_addr)
+
+                self.server_transport.sendto(data)
             except Exception as e:
-                    logging.error(f"Error forwarding to server: {e}")
+                logging.error(
-                    self.proxy.reset()
+                    "Error forwarding to server for "
+                    f"{self.client_addr[0]}:{self.client_addr[1]}: {e}")
+                self.proxy.remove_session(self.client_addr)
 
-        def error_received(self, exc: Exception|None) -> None:
+        def handle_from_server(self, data: bytes) -> None:
-            logging.error(f"Client socket error: {exc}")
+            try:
-            self.proxy.reset()
+                data = self.proxy.on_forward(data, to_server=False,
+                                             client_addr=self.client_addr)
+                if self.proxy.listener_transport:
+                    self.proxy.listener_transport.sendto(data, self.client_addr)
+            except Exception as e:
+                logging.error(
+                    "Error forwarding to client "
+                    f"{self.client_addr[0]}:{self.client_addr[1]}: {e}")
+                self.proxy.remove_session(self.client_addr)
 
-        def connection_lost(self, exc: Exception|None) -> None:
+    class ListenerProtocol(asyncio.DatagramProtocol):
-            logging.error(f"Client connection lost: {exc}")
+        def __init__(self, proxy: UDPProxy) -> None:
-            self.proxy.reset()
+            self.proxy = proxy
+
+        def datagram_received(self, data: bytes,
+                              addr: tuple[str, int]) -> None:
+            session = self.proxy.sessions.get(addr)
+            if not session:
+                asyncio.create_task(self._handle_new_client(data, addr))
+                return
+            session.handle_from_client(data)
+
+        async def _handle_new_client(self, data: bytes,
+                                     addr: tuple[str, int]) -> None:
+            session = await self.proxy._create_session(addr)
+            session.handle_from_client(data)
+
+        def error_received(self, exc: Exception) -> None:
+            logging.error(f"Listener socket error: {exc}")
 
     class ServerProtocol(asyncio.DatagramProtocol):
-        def __init__(self, proxy: AbstractUDPProxy) -> None:
+        def __init__(self, proxy: UDPProxy,
-            self.proxy: AbstractUDPProxy = proxy
+                     client_addr: tuple[str, int]) -> None:
+            self.proxy = proxy
+            self.client_addr = client_addr
 
-        def datagram_received(self, data: bytes, addr: tuple) -> None:
+        def datagram_received(self, data: bytes,
-            if self.proxy.transport_client and self.proxy.client_address:
+                              addr: tuple[str, int]) -> None:
-                try:
+            session = self.proxy.sessions.get(self.client_addr)
-                    data = self.proxy.on_forward(data, to_server=False)
+            if session:
-                    self.proxy.transport_client.sendto(
+                session.handle_from_server(data)
-                        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:
+        def error_received(self, exc: Exception) -> None:
-            logging.error(f"Server socket error: {exc}")
+            logging.error("Server socket error for "
-            self.proxy.reset()
+                          f"{self.client_addr[0]}:{self.client_addr[1]}: {exc}")
+            self.proxy.remove_session(self.client_addr)
 
-        def connection_lost(self, exc: Exception|None) -> None:
+        def connection_lost(self, exc: Exception | None) -> None:
-            logging.error(f"Server connection lost: {exc}")
+            logging.error("Server connection lost for "
-            self.proxy.reset()
+                          f"{self.client_addr[0]}:{self.client_addr[1]}: {exc}")
+            self.proxy.remove_session(self.client_addr)

relog.py

@@ -19,7 +19,7 @@ def main():
             line = line.strip()
             if line.startswith("HD: "):
                 pdata = bytes.fromhex(line[4:])
-                log_packet("DIR", pdata, ('', ''))
+                log_packet("<=??=>", pdata)
 
 
 if __name__ == "__main__":

udpproxylogger.py

@@ -3,12 +3,11 @@ import asyncio
 import logging
 import traceback
 
-from dmrtools.udpproxy import AbstractUDPProxy
+from dmrtools.udpproxy import UDPProxy
 from dmrtools import DMRPPacketFactory
 from dmrtools import hexdump
-# from dmrtools.dmrproto import DMRPL2FullLC, DMRPL2VoiceBurst
 from dmrtools.dmrproto import DMRPPacketData
-from dmrtools.dmrproto import EmbLCAssembler, LCFactory
+from dmrtools.dmrproto import CallLCDecoder, LCFactory, LCTalkerAlias
 
 
 def setup_logger(log_file=None):
@@ -25,27 +24,20 @@ def setup_logger(log_file=None):
     )
 
 
-elca: dict[int, EmbLCAssembler] = dict()
+lcd: dict[int, CallLCDecoder] = dict()
 
 def get_packet_details(data):
     try:
         p = DMRPPacketFactory.fd(data)
         if type(p) is DMRPPacketData:
-            emblcinfo = ""
+            if p.stream_id not in lcd:
-            if (full_lc := p.get_full_lc()) is not None:
+                lcd[p.stream_id] = CallLCDecoder(p.stream_id)
-                lcdec = LCFactory.fd(full_lc)
+            if (lc := lcd[p.stream_id].process_voicedata(p)) is not None:
-                emblcinfo = f"FullLC:{hexdump(full_lc)}\n{lcdec}\n"
+                if type(lc) is LCTalkerAlias:
-                return f"{p:l2}\n" + emblcinfo
+                    if (ta := lcd[p.stream_id].ta) is not None:
-            if p.stream_id not in elca:
+                        return f"{p:l2}\nLCdata: {hexdump(lc._data)}\n{lc}\nTA: {ta}\n"
-                elca[p.stream_id] = EmbLCAssembler()
+                return f"{p:l2}\nLCdata: {hexdump(lc._data)}\n{lc}\n"
-            if elca[p.stream_id].process_voicedata(p):
+            return f"{p:l2}\n"
-                emblcdata = elca[p.stream_id].decode()
-                lcdec = LCFactory.fd(emblcdata)
-                emblcinfo = f"EmbLC:{hexdump(emblcdata)}\n{lcdec}\n"
-                del elca[p.stream_id]
-            if p.is_voice_term:
-                del elca[p.stream_id]
-            return f"{p:l2}\n" + emblcinfo
         return f"{p}\n"
     except:
         return f"Exception while decoding packet:\n{traceback.format_exc()}"
@@ -60,9 +52,11 @@ def log_packet(direction, data):
     logging.info(log_message)
 
 
-class UDPProxyLogger(AbstractUDPProxy):
+class UDPProxyLogger(UDPProxy):
-    def on_forward(self, data: bytes, to_server: bool) -> bytes:
+    def on_forward(self, data: bytes, to_server: bool,
-        direction = "OUT=>" if to_server else "<==IN"
+                   client_addr: tuple[str, int]) -> bytes:
+        ca_str = f"{client_addr[0]}:{client_addr[1]}"
+        direction = f"{ca_str} =>" if to_server else f"{ca_str} <="
         log_packet(direction, data)
         return data