shaba/openuds
1
0
Fork 1
mirror of https://github.com/dkmstr/openuds.git synced 2025-01-10 01:17:59 +03:00
Code Issues Releases Wiki Activity

added keep_listening parameter for future and initial_payload (fix)

Browse Source
This commit is contained in:
Adolfo Gómez García 2023-05-10 16:45:39 +02:00
parent db70f02df0
commit 8c9b326c3c
No known key found for this signature in database
GPG Key ID: DD1ABF20724CDA23
1 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
client-py3/full/src/uds/tunnel.py
View File

@ -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,