[proxy] multi-client proxy
Этот коммит содержится в:
+94
-77
@@ -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_port: int = listen_port
|
||||
self.server_host: str = server_host
|
||||
self.server_port: int = server_port
|
||||
self.listen_host = listen_host
|
||||
self.listen_port = listen_port
|
||||
self.server_host = server_host
|
||||
self.server_port = server_port
|
||||
|
||||
self.transport_client: asyncio.DatagramTransport|None = None
|
||||
self.transport_server: asyncio.DatagramTransport|None = None
|
||||
self.listener_transport: 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) -> bytes:
|
||||
def on_forward(self, data: bytes, to_server: bool,
|
||||
client_addr: tuple[str, int]) -> 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.transport_client, _ = await loop.create_datagram_endpoint(
|
||||
lambda: self.ClientProtocol(self),
|
||||
self.listener_transport, _ = await loop.create_datagram_endpoint(
|
||||
lambda: self.ListenerProtocol(self),
|
||||
local_addr=(self.listen_host, self.listen_port)
|
||||
)
|
||||
|
||||
logging.info(
|
||||
f"Listening for on {self.listen_host}:{self.listen_port}")
|
||||
logging.info(f"Listening for clients 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):
|
||||
async def _create_session(self, client_addr: tuple[str, int]) -> Session:
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
self.transport_server, _ = await loop.create_datagram_endpoint(
|
||||
lambda: self.ServerProtocol(self),
|
||||
transport, _ = await loop.create_datagram_endpoint(
|
||||
lambda: self.ServerProtocol(self, client_addr),
|
||||
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}")
|
||||
session = UDPProxy.Session(proxy=self,
|
||||
client_addr=client_addr,
|
||||
server_transport=transport)
|
||||
|
||||
class ClientProtocol(asyncio.DatagramProtocol):
|
||||
def __init__(self, proxy: AbstractUDPProxy) -> None:
|
||||
self.proxy: AbstractUDPProxy = proxy
|
||||
self.sessions[client_addr] = session
|
||||
|
||||
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())
|
||||
logging.info(f"Proxying {client_addr[0]}:{client_addr[1]}"
|
||||
f" <=> {self.server_host}:{self.server_port}")
|
||||
|
||||
if addr != self.proxy.client_address:
|
||||
logging.error("Ignoring packet from unknown client"
|
||||
f" {addr[0]}:{addr[1]}")
|
||||
return session
|
||||
|
||||
def remove_session(self, client_addr: tuple[str, int]) -> None:
|
||||
session = self.sessions.pop(client_addr, None)
|
||||
if session:
|
||||
session.server_transport.close()
|
||||
logging.info(f"Closed session for {client_addr}")
|
||||
|
||||
class Session:
|
||||
def __init__(self, proxy: UDPProxy, client_addr: tuple[str, int],
|
||||
server_transport: asyncio.DatagramTransport) -> None:
|
||||
self.proxy = proxy
|
||||
self.client_addr = client_addr
|
||||
self.server_transport = server_transport
|
||||
|
||||
def handle_from_client(self, data: bytes) -> None:
|
||||
try:
|
||||
data = self.proxy.on_forward(data, to_server=True,
|
||||
client_addr=self.client_addr)
|
||||
|
||||
self.server_transport.sendto(data)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
"Error forwarding to server for "
|
||||
f"{self.client_addr[0]}:{self.client_addr[1]}: {e}")
|
||||
self.proxy.remove_session(self.client_addr)
|
||||
|
||||
def handle_from_server(self, data: bytes) -> None:
|
||||
try:
|
||||
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)
|
||||
|
||||
class ListenerProtocol(asyncio.DatagramProtocol):
|
||||
def __init__(self, proxy: UDPProxy) -> None:
|
||||
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)
|
||||
|
||||
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()
|
||||
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) -> 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()
|
||||
def error_received(self, exc: Exception) -> None:
|
||||
logging.error(f"Listener socket error: {exc}")
|
||||
|
||||
class ServerProtocol(asyncio.DatagramProtocol):
|
||||
def __init__(self, proxy: AbstractUDPProxy) -> None:
|
||||
self.proxy: AbstractUDPProxy = proxy
|
||||
def __init__(self, proxy: UDPProxy,
|
||||
client_addr: tuple[str, int]) -> None:
|
||||
self.proxy = proxy
|
||||
self.client_addr = client_addr
|
||||
|
||||
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 datagram_received(self, data: bytes,
|
||||
addr: tuple[str, int]) -> None:
|
||||
session = self.proxy.sessions.get(self.client_addr)
|
||||
if session:
|
||||
session.handle_from_server(data)
|
||||
|
||||
def error_received(self, exc: Exception|None) -> None:
|
||||
logging.error(f"Server socket error: {exc}")
|
||||
self.proxy.reset()
|
||||
def error_received(self, exc: Exception) -> None:
|
||||
logging.error("Server socket error for "
|
||||
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:
|
||||
logging.error(f"Server connection lost: {exc}")
|
||||
self.proxy.reset()
|
||||
def connection_lost(self, exc: Exception | None) -> None:
|
||||
logging.error("Server connection lost for "
|
||||
f"{self.client_addr[0]}:{self.client_addr[1]}: {exc}")
|
||||
self.proxy.remove_session(self.client_addr)
|
||||
|
||||
+1
-1
@@ -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__":
|
||||
|
||||
+16
-22
@@ -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 (full_lc := p.get_full_lc()) is not None:
|
||||
lcdec = LCFactory.fd(full_lc)
|
||||
emblcinfo = f"FullLC:{hexdump(full_lc)}\n{lcdec}\n"
|
||||
return f"{p:l2}\n" + emblcinfo
|
||||
if p.stream_id not in elca:
|
||||
elca[p.stream_id] = EmbLCAssembler()
|
||||
if elca[p.stream_id].process_voicedata(p):
|
||||
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
|
||||
if p.stream_id not in lcd:
|
||||
lcd[p.stream_id] = CallLCDecoder(p.stream_id)
|
||||
if (lc := lcd[p.stream_id].process_voicedata(p)) is not None:
|
||||
if type(lc) is LCTalkerAlias:
|
||||
if (ta := lcd[p.stream_id].ta) is not None:
|
||||
return f"{p:l2}\nLCdata: {hexdump(lc._data)}\n{lc}\nTA: {ta}\n"
|
||||
return f"{p:l2}\nLCdata: {hexdump(lc._data)}\n{lc}\n"
|
||||
return f"{p:l2}\n"
|
||||
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):
|
||||
def on_forward(self, data: bytes, to_server: bool) -> bytes:
|
||||
direction = "OUT=>" if to_server else "<==IN"
|
||||
class UDPProxyLogger(UDPProxy):
|
||||
def on_forward(self, data: bytes, to_server: bool,
|
||||
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
|
||||
|
||||
|
||||
Ссылка в новой задаче
Block a user