diff --git a/dmrtools/udpproxy.py b/dmrtools/udpproxy.py new file mode 100644 index 0000000..a97f662 --- /dev/null +++ b/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() diff --git a/udpproxylogger.py b/udpproxylogger.py index 7a561c8..74708ec 100755 --- a/udpproxylogger.py +++ b/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): - def __init__(self, server_address, loop): - self.server_address = server_address - self.loop = loop - self.transport = None - 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") +class UDPProxyLogger(AbstractUDPProxy): + def on_forward(self, data: bytes, to_server: bool) -> bytes: + direction = "OUT=>" if to_server else "<==IN" + log_packet(direction, data) + return data -async def start_udp_proxy(listen_ip, listen_port, server_ip, server_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 - +async def start_udp_proxy(server_ip, server_port, listen_ip, listen_port): 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() - server_transport.close() + # await proxy.stop() + 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, - args.serverip, args.serverport)) + asyncio.run(start_udp_proxy(args.serverip, args.serverport, + args.listenip, args.listenport)) except KeyboardInterrupt: logging.info("Shutting down.")