From 8c9b326c3c869aa09ad1841c52845d63ccb3665b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adolfo=20G=C3=B3mez=20Garc=C3=ADa?= Date: Wed, 10 May 2023 16:45:39 +0200 Subject: [PATCH] added keep_listening parameter for future and initial_payload (fix) --- client-py3/full/src/uds/tunnel.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/client-py3/full/src/uds/tunnel.py b/client-py3/full/src/uds/tunnel.py index c5ad1cfdf..8fc177e7d 100644 --- a/client-py3/full/src/uds/tunnel.py +++ b/client-py3/full/src/uds/tunnel.py @@ -49,8 +49,10 @@ LISTEN_ADDRESS = '0.0.0.0' if DEBUG else '127.0.0.1' # ForwarServer states TUNNEL_LISTENING, TUNNEL_OPENING, TUNNEL_PROCESSING, TUNNEL_ERROR = 0, 1, 2, 3 + logger = logging.getLogger(__name__) +PayLoadType = typing.Optional[typing.Tuple[typing.Optional[bytes], typing.Optional[bytes]]] class ForwardServer(socketserver.ThreadingTCPServer): daemon_threads = True @@ -66,7 +68,7 @@ class ForwardServer(socketserver.ThreadingTCPServer): keep_listening: bool current_connections: int status: int - initial_payload: typing.Optional[bytes] + initial_payload: PayLoadType def __init__( self, @@ -76,7 +78,7 @@ class ForwardServer(socketserver.ThreadingTCPServer): local_port: int = 0, check_certificate: bool = True, keep_listening: bool = False, - initial_payload: typing.Optional[bytes] = None, + initial_payload: PayLoadType = None, ) -> None: local_port = local_port or random.randrange(33000, 53000) @@ -200,7 +202,13 @@ class Handler(socketserver.BaseRequestHandler): # If we have a payload, send it if self.server.initial_payload: - ssl_socket.sendall(self.server.initial_payload) + to_send, to_receive = self.server.initial_payload + if to_send: # To send + ssl_socket.sendall(to_send) + if to_receive: + temp = ssl_socket.recv(len(to_receive)) + if temp != to_receive: + raise Exception(f'Invalid response: {temp!s} != {to_receive!s}') self.process(remote=ssl_socket) except Exception as e: @@ -254,7 +262,7 @@ def forward( local_port: int = 0, check_certificate=True, keep_listening=True, - initial_payload: typing.Optional[bytes] = None, + initial_payload: PayLoadType = None, ) -> ForwardServer: fs = ForwardServer( remote=remote,