[proxy] Switch to AbstractUDPProxy
Этот коммит содержится в:
Alexander UR6LKW
@@ -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()
|
||||
+17
-80
@@ -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.")
|
||||
|
||||
|
||||
Ссылка в новой задаче
Block a user
© 2026 StackWorks