forked from shaba/openuds
Upgraded tunnel.py to last real one
This commit is contained in:
Adolfo Gómez García
parent
4424f2a497
commit
011145e911
@ -1,6 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# Copyright (c) 2021 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
@ -33,6 +33,7 @@ import socketserver
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
import random
|
||||
import threading
|
||||
import select
|
||||
import typing
|
||||
@ -48,6 +49,7 @@ TUNNEL_LISTENING, TUNNEL_OPENING, TUNNEL_PROCESSING, TUNNEL_ERROR = 0, 1, 2, 3
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ForwardServer(socketserver.ThreadingTCPServer):
|
||||
daemon_threads = True
|
||||
allow_reuse_address = True
|
||||
@ -55,6 +57,7 @@ class ForwardServer(socketserver.ThreadingTCPServer):
|
||||
remote: typing.Tuple[str, int]
|
||||
ticket: str
|
||||
stop_flag: threading.Event
|
||||
can_stop: bool
|
||||
timeout: int
|
||||
timer: typing.Optional[threading.Timer]
|
||||
check_certificate: bool
|
||||
@ -69,23 +72,30 @@ class ForwardServer(socketserver.ThreadingTCPServer):
|
||||
local_port: int = 0,
|
||||
check_certificate: bool = True,
|
||||
) -> None:
|
||||
|
||||
local_port = local_port or random.randrange(33000, 53000)
|
||||
|
||||
super().__init__(
|
||||
server_address=(LISTEN_ADDRESS, local_port), RequestHandlerClass=Handler
|
||||
)
|
||||
self.remote = remote
|
||||
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.stop_flag = threading.Event() # False initial
|
||||
self.current_connections = 0
|
||||
|
||||
self.status = TUNNEL_LISTENING
|
||||
self.can_stop = False
|
||||
|
||||
if timeout:
|
||||
self.timer = threading.Timer(timeout, ForwardServer.__checkStarted, args=(self,))
|
||||
self.timer.start()
|
||||
else:
|
||||
self.timer = None
|
||||
timeout = abs(timeout) or 60
|
||||
self.timer = threading.Timer(
|
||||
abs(timeout), ForwardServer.__checkStarted, args=(self,)
|
||||
)
|
||||
self.timer.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
if not self.stop_flag.is_set():
|
||||
@ -96,13 +106,52 @@ class ForwardServer(socketserver.ThreadingTCPServer):
|
||||
self.timer = None
|
||||
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
|
||||
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
|
||||
def __checkStarted(fs: 'ForwardServer') -> None:
|
||||
logger.debug('New connection limit reached')
|
||||
fs.timer = None
|
||||
fs.can_stop = True
|
||||
if fs.current_connections <= 0:
|
||||
fs.stop()
|
||||
|
||||
@ -113,46 +162,33 @@ class Handler(socketserver.BaseRequestHandler):
|
||||
|
||||
# server: ForwardServer
|
||||
def handle(self) -> None:
|
||||
self.server.current_connections += 1
|
||||
self.server.status = TUNNEL_OPENING
|
||||
|
||||
# If server processing is over time
|
||||
# If server processing is over time
|
||||
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
|
||||
return
|
||||
|
||||
self.server.current_connections += 1
|
||||
|
||||
# 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.ticket)
|
||||
logger.debug('Ticket %s', self.server.ticket)
|
||||
with self.server.connect() 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(errors="ignore")}'
|
||||
) # Notify error
|
||||
|
||||
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:
|
||||
# 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)
|
||||
# All is fine, now we can tunnel data
|
||||
self.process(remote=ssl_socket)
|
||||
except Exception as e:
|
||||
logger.error(f'Error connecting to {self.server.remote!s}: {e!s}')
|
||||
self.server.status = TUNNEL_ERROR
|
||||
@ -185,10 +221,17 @@ class Handler(socketserver.BaseRequestHandler):
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
|
||||
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()
|
||||
logger.debug('Stoped forwarded %s -> %s', server.server_address, server.remote)
|
||||
logger.debug('Stoped forwarder %s -> %s', server.server_address, server.remote)
|
||||
|
||||
|
||||
def forward(
|
||||
remote: typing.Tuple[str, int],
|
||||
@ -197,6 +240,7 @@ def forward(
|
||||
local_port: int = 0,
|
||||
check_certificate=True,
|
||||
) -> ForwardServer:
|
||||
|
||||
fs = ForwardServer(
|
||||
remote=remote,
|
||||
ticket=ticket,
|
||||
@ -209,8 +253,10 @@ def forward(
|
||||
|
||||
return fs
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
log = logging.getLogger()
|
||||
log.setLevel(logging.DEBUG)
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
@ -221,4 +267,12 @@ if __name__ == "__main__":
|
||||
handler.setFormatter(formatter)
|
||||
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,
|
||||
)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user