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

Advancing on forwarder

Browse Source
This commit is contained in:
Adolfo Gómez García 2021-01-14 08:10:28 +01:00
parent d6a8639b18
commit 40abfb6014
4 changed files with 79 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

81
tunnel-server/src/forwarder/uds_forwarder.py
View File

@ -54,10 +54,11 @@ class ForwardServer(socketserver.ThreadingTCPServer):
remote: typing.Tuple[str, int]
ticket: str
running: bool
stop_flag: threading.Event
timeout: int
timer: typing.Optional[threading.Timer]
check_certificate: bool
current_connections: int
status: int
def __init__(
@ -76,16 +77,35 @@ class ForwardServer(socketserver.ThreadingTCPServer):
self.timeout = int(time.time()) + timeout if timeout else 0
self.check_certificate = check_certificate
self.stop_flag = threading.Event() # False initial
self.running = True
self.current_connections = 0
self.status = TUNNEL_LISTENING
if timeout:
self.timer = threading.Timer(timeout, ForwardServer.__checkStarted, args=(self,))
self.timer.start()
else:
self.timer = None
def stop(self) -> None:
if not self.stop_flag.is_set():
logger.debug('Stopping servers')
self.stop_flag.set()
self.running = False
if self.timer:
self.timer.cancel()
self.timer = None
self.shutdown()
@property
def stoppable(self) -> bool:
return self.timeout != 0 and int(time.time()) > self.timeout
@staticmethod
def __checkStarted(fs: 'ForwardServer') -> None:
fs.timer = None
if fs.current_connections <= 0:
fs.stop()
class Handler(socketserver.BaseRequestHandler):
# Override Base type
@ -93,18 +113,31 @@ class Handler(socketserver.BaseRequestHandler):
# server: ForwardServer
def handle(self) -> None:
# If server processing is timed out...
if self.server.timeout and int(time.time()) > self.server.timeout:
self.server.current_connections += 1
self.server.status = TUNNEL_OPENING
# If server processing is over time
if self.server.stoppable:
logger.info('Rejected timedout connection try')
self.request.close() # End connection without processing it
return
# Open remote connection
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as rsocket:
logger.info('CONNECT to %s', self.server.remote)
logger.debug('Ticket %s', self.server.remote, self.server.ticket)
rsocket.connect(self.server.remote)
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:
@ -116,19 +149,24 @@ class Handler(socketserver.BaseRequestHandler):
data = ssl_socket.recv(2)
if data != b'OK':
data += ssl_socket.recv(128)
raise Exception(data.decode()) # Notify error
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:
# TODO log error connecting...
if DEBUG:
logger.exception('Processing')
logger.error(f'Error connecting: {e!s}')
logger.error(f'Error connecting to {self.server.remote!s}: {e!s}')
self.server.status = TUNNEL_ERROR
self.server.stop()
finally:
self.server.current_connections -= 1
if self.server.current_connections <= 0 and self.server.stoppable:
self.server.stop()
# Processes data forwarding
def process(self, remote: ssl.SSLSocket):
self.server.status = TUNNEL_PROCESSING
logger.debug('Processing tunnel with ticket %s', self.server.ticket)
# Process data until stop requested or connection closed
try:
while not self.server.stop_flag.is_set():
@ -143,14 +181,14 @@ class Handler(socketserver.BaseRequestHandler):
if not data:
break
self.request.sendall(data)
logger.debug('Finished process')
logger.debug('Finished tunnel with ticekt %s', self.server.ticket)
except Exception as e:
pass
def _run(server: ForwardServer) -> None:
logger.debug('Starting server')
logger.debug('Starting forwarder: %s -> %s, timeout: %d', server.server_address, server.remote, server.timeout)
server.serve_forever()
logger.debug('Stoped server')
logger.debug('Stoped forwarded %s -> %s', server.server_address, server.remote)
def forward(
remote: typing.Tuple[str, int],
@ -172,10 +210,15 @@ def forward(
return fs
if __name__ == "__main__":
fs1 = forward(('fake.udsenterprise.com', 7777), '0'*64, local_port=49998)
print(f'Listening on {fs1.server_address}')
fs2 = forward(('fake.udsenterprise.com', 7777), '1'*64, local_port=49999)
print(f'Listening on {fs2.server_address}')
# time.sleep(30)
# fs.stop()
import sys
log = logging.getLogger()
log.setLevel(logging.DEBUG)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.DEBUG)
formatter = logging.Formatter(
'%(levelname)s - %(message)s'
) # Basic log format, nice for syslog
handler.setFormatter(formatter)
log.addHandler(handler)
fs = forward(('172.27.0.1', 7777), '1'*64, local_port=49999, timeout=10, check_certificate=False)

7
tunnel-server/src/uds_tunnel/proxy.py
View File

@ -158,11 +158,8 @@ class Proxy:
result = await curio.run_in_thread(Proxy.getFromUds, self.cfg, ticket)
except Exception as e:
logger.error('ERROR %s', e.args[0] if e.args else e)
raise
# Invalid result from UDS, not allowed to connect
if not result:
raise Exception('INVALID TICKET')
await source.sendall(b'ERROR INVALID TICKET')
return
logger.info('OPEN TUNNEL FROM %s to %s:%s', pretty_adress, result['host'], result['port'])

12
tunnel-server/src/uds_tunnel/stats.py
View File

@ -64,7 +64,9 @@ class StatsSingleCounter:
class Stats:
ns: 'Namespace'
last_sent: int
sent: int
last_recv: int
recv: int
last: float
@ -72,16 +74,18 @@ class Stats:
self.ns = ns
self.ns.current += 1
self.ns.total += 1
self.sent = 0
self.recv = 0
self.sent = self.last_sent = 0
self.recv = self.last_recv = 0
self.last = time.monotonic()
def update(self, force: bool = False):
now = time.monotonic()
if force or now - self.last > INTERVAL:
self.last = now
self.ns.recv = self.recv
self.ns.sent = self.sent
self.ns.recv += self.recv - self.last_recv
self.ns.sent += self.sent - self.last_sent
self.last_sent = self.sent
self.last_recv = self.recv
def add_recv(self, size: int) -> None:
self.recv += size

12
tunnel-server/src/udstunnel.py
View File

@ -198,10 +198,12 @@ def tunnel_main():
# Wait for socket incoming connections and spread them
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
try:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, True)
except (AttributeError, OSError) as e:
logger.warning('socket.REUSEPORT not available')
# We will not reuse port, we only want a UDS tunnel server running on a port
# but this may change on future...
# try:
# sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, True)
# except (AttributeError, OSError) as e:
# logger.warning('socket.REUSEPORT not available')
sock.settimeout(3.0) # So we can check for stop from time to time
sock.bind((cfg.listen_address, cfg.listen_port))
@ -219,8 +221,8 @@ def tunnel_main():
except Exception as e:
logger.error('LOOP: %s', e)
except Exception as e:
sys.stderr.write(f'Error: {e}\n')
logger.error('MAIN: %s', e)
pass
if sock:
sock.close()