burmatovan/openuds
1
0
Fork 0
forked from shaba/openuds
Code Issues Pull Requests Projects Releases Wiki Activity

Upgraded tunnel.py to last real one

Browse Source
This commit is contained in:
Adolfo Gómez García 2021-06-23 16:17:23 +02:00
parent 4424f2a497
commit 011145e911
1 changed files with 96 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,6 +1,6 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
# #
# Copyright (c) 2020 Virtual Cable S.L.U. # Copyright (c) 2021 Virtual Cable S.L.U.
# All rights reserved. # All rights reserved.
# #
# Redistribution and use in source and binary forms, with or without modification, # Redistribution and use in source and binary forms, with or without modification,
@ -33,6 +33,7 @@ import socketserver
import ssl import ssl
import threading import threading
import time import time
import random
import threading import threading
import select import select
import typing import typing
@ -48,6 +49,7 @@ TUNNEL_LISTENING, TUNNEL_OPENING, TUNNEL_PROCESSING, TUNNEL_ERROR = 0, 1, 2, 3
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class ForwardServer(socketserver.ThreadingTCPServer): class ForwardServer(socketserver.ThreadingTCPServer):
daemon_threads = True daemon_threads = True
allow_reuse_address = True allow_reuse_address = True
@ -55,6 +57,7 @@ class ForwardServer(socketserver.ThreadingTCPServer):
remote: typing.Tuple[str, int] remote: typing.Tuple[str, int]
ticket: str ticket: str
stop_flag: threading.Event stop_flag: threading.Event
can_stop: bool
timeout: int timeout: int
timer: typing.Optional[threading.Timer] timer: typing.Optional[threading.Timer]
check_certificate: bool check_certificate: bool
@ -69,23 +72,30 @@ class ForwardServer(socketserver.ThreadingTCPServer):
local_port: int = 0, local_port: int = 0,
check_certificate: bool = True, check_certificate: bool = True,
) -> None: ) -> None:
local_port = local_port or random.randrange(33000, 53000)
super().__init__( super().__init__(
server_address=(LISTEN_ADDRESS, local_port), RequestHandlerClass=Handler server_address=(LISTEN_ADDRESS, local_port), RequestHandlerClass=Handler
) )
self.remote = remote self.remote = remote
self.ticket = ticket self.ticket = ticket
self.timeout = int(time.time()) + timeout if timeout else 0 # Negative values for timeout, means "accept always connections"
# "but if no connection is stablished on timeout (positive)"
# "stop the listener"
self.timeout = int(time.time()) + timeout if timeout > 0 else 0
self.check_certificate = check_certificate self.check_certificate = check_certificate
self.stop_flag = threading.Event() # False initial self.stop_flag = threading.Event() # False initial
self.current_connections = 0 self.current_connections = 0
self.status = TUNNEL_LISTENING self.status = TUNNEL_LISTENING
self.can_stop = False
if timeout: timeout = abs(timeout) or 60
self.timer = threading.Timer(timeout, ForwardServer.__checkStarted, args=(self,)) self.timer = threading.Timer(
self.timer.start() abs(timeout), ForwardServer.__checkStarted, args=(self,)
else: )
self.timer = None self.timer.start()
def stop(self) -> None: def stop(self) -> None:
if not self.stop_flag.is_set(): if not self.stop_flag.is_set():
@ -96,13 +106,52 @@ class ForwardServer(socketserver.ThreadingTCPServer):
self.timer = None self.timer = None
self.shutdown() self.shutdown()
def connect(self) -> ssl.SSLSocket:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as rsocket:
logger.info('CONNECT to %s', self.remote)
rsocket.connect(self.remote)
context = ssl.create_default_context()
# Do not "recompress" data, use only "base protocol" compression
context.options |= ssl.OP_NO_COMPRESSION
# If ignore remote certificate
if self.check_certificate is False:
context.check_hostname = False
context.verify_mode = ssl.CERT_NONE
logger.warning('Certificate checking is disabled!')
return context.wrap_socket(rsocket, server_hostname=self.remote[0])
def check(self) -> bool:
if self.status == TUNNEL_ERROR:
return False
try:
with self.connect() as ssl_socket:
ssl_socket.sendall(HANDSHAKE_V1 + b'TEST')
resp = ssl_socket.recv(2)
if resp != b'OK':
raise Exception({'Invalid tunnelresponse: {resp}'})
return True
except Exception as e:
logger.error(
'Error connecting to tunnel server %s: %s', self.server_address, e
)
return False
@property @property
def stoppable(self) -> bool: def stoppable(self) -> bool:
return self.timeout != 0 and int(time.time()) > self.timeout logger.debug('Is stoppable: %s', self.can_stop)
return self.can_stop or (self.timeout != 0 and int(time.time()) > self.timeout)
@staticmethod @staticmethod
def __checkStarted(fs: 'ForwardServer') -> None: def __checkStarted(fs: 'ForwardServer') -> None:
logger.debug('New connection limit reached')
fs.timer = None fs.timer = None
fs.can_stop = True
if fs.current_connections <= 0: if fs.current_connections <= 0:
fs.stop() fs.stop()
@ -113,46 +162,33 @@ class Handler(socketserver.BaseRequestHandler):
# server: ForwardServer # server: ForwardServer
def handle(self) -> None: def handle(self) -> None:
self.server.current_connections += 1
self.server.status = TUNNEL_OPENING self.server.status = TUNNEL_OPENING
# If server processing is over time # If server processing is over time
if self.server.stoppable: if self.server.stoppable:
logger.info('Rejected timedout connection try') self.server.status = TUNNEL_ERROR
logger.info('Rejected timedout connection')
self.request.close() # End connection without processing it self.request.close() # End connection without processing it
return return
self.server.current_connections += 1
# Open remote connection # Open remote connection
try: try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as rsocket: logger.debug('Ticket %s', self.server.ticket)
logger.info('CONNECT to %s', self.server.remote) with self.server.connect() as ssl_socket:
logger.debug('Ticket %s', self.server.ticket) # Send handhshake + command + ticket
ssl_socket.sendall(HANDSHAKE_V1 + b'OPEN' + self.server.ticket.encode())
# Check response is OK
data = ssl_socket.recv(2)
if data != b'OK':
data += ssl_socket.recv(128)
raise Exception(
f'Error received: {data.decode(errors="ignore")}'
) # Notify error
rsocket.connect(self.server.remote) # All is fine, now we can tunnel data
self.process(remote=ssl_socket)
context = ssl.create_default_context()
# If ignore remote certificate
if self.server.check_certificate is False:
context.check_hostname = False
context.verify_mode = ssl.CERT_NONE
logger.warning('Certificate checking is disabled!')
with context.wrap_socket(
rsocket, server_hostname=self.server.remote[0]
) as ssl_socket:
# Send handhshake + command + ticket
ssl_socket.sendall(
HANDSHAKE_V1 + b'OPEN' + self.server.ticket.encode()
)
# Check response is OK
data = ssl_socket.recv(2)
if data != b'OK':
data += ssl_socket.recv(128)
raise Exception(f'Error received: {data.decode()}') # Notify error
# All is fine, now we can tunnel data
self.process(remote=ssl_socket)
except Exception as e: except Exception as e:
logger.error(f'Error connecting to {self.server.remote!s}: {e!s}') logger.error(f'Error connecting to {self.server.remote!s}: {e!s}')
self.server.status = TUNNEL_ERROR self.server.status = TUNNEL_ERROR
@ -185,10 +221,17 @@ class Handler(socketserver.BaseRequestHandler):
except Exception as e: except Exception as e:
pass pass
def _run(server: ForwardServer) -> None: def _run(server: ForwardServer) -> None:
logger.debug('Starting forwarder: %s -> %s, timeout: %d', server.server_address, server.remote, server.timeout) logger.debug(
'Starting forwarder: %s -> %s, timeout: %d',
server.server_address,
server.remote,
server.timeout,
)
server.serve_forever() server.serve_forever()
logger.debug('Stoped forwarded %s -> %s', server.server_address, server.remote) logger.debug('Stoped forwarder %s -> %s', server.server_address, server.remote)
def forward( def forward(
remote: typing.Tuple[str, int], remote: typing.Tuple[str, int],
@ -197,6 +240,7 @@ def forward(
local_port: int = 0, local_port: int = 0,
check_certificate=True, check_certificate=True,
) -> ForwardServer: ) -> ForwardServer:
fs = ForwardServer( fs = ForwardServer(
remote=remote, remote=remote,
ticket=ticket, ticket=ticket,
@ -209,8 +253,10 @@ def forward(
return fs return fs
if __name__ == "__main__": if __name__ == "__main__":
import sys import sys
log = logging.getLogger() log = logging.getLogger()
log.setLevel(logging.DEBUG) log.setLevel(logging.DEBUG)
handler = logging.StreamHandler(sys.stdout) handler = logging.StreamHandler(sys.stdout)
@ -221,4 +267,12 @@ if __name__ == "__main__":
handler.setFormatter(formatter) handler.setFormatter(formatter)
log.addHandler(handler) log.addHandler(handler)
fs = forward(('172.27.0.1', 7777), '1'*64, local_port=49999, timeout=10, check_certificate=False) ticket = 'mffqg7q4s61fvx0ck2pe0zke6k0c5ipb34clhbkbs4dasb4g'
fs = forward(
('172.27.0.1', 7777),
ticket,
local_port=49999,
timeout=-20,
check_certificate=False,
)