forked from shaba/openuds
added local uds_tunnel work to openuds
This commit is contained in:
parent
f0bd3782d7
commit
e486d6708d
0
tunnel-server/.gitignore
vendored
Normal file
0
tunnel-server/.gitignore
vendored
Normal file
2
tunnel-server/requirements.txt
Normal file
2
tunnel-server/requirements.txt
Normal file
@ -0,0 +1,2 @@
|
||||
curio>=1.4
|
||||
psutil>=5.7.3
|
1
tunnel-server/src/.gitignore
vendored
Normal file
1
tunnel-server/src/.gitignore
vendored
Normal file
@ -0,0 +1 @@
|
||||
tests
|
180
tunnel-server/src/forwarder/uds_forwarder.py
Normal file
180
tunnel-server/src/forwarder/uds_forwarder.py
Normal file
@ -0,0 +1,180 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import socket
|
||||
import socketserver
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
import threading
|
||||
import select
|
||||
import typing
|
||||
import logging
|
||||
|
||||
HANDSHAKE_V1 = b'\x5AMGB\xA5\x01\x00'
|
||||
BUFFER_SIZE = 1024 * 16 # Max buffer length
|
||||
DEBUG = True
|
||||
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__)
|
||||
|
||||
class ForwardServer(socketserver.ThreadingTCPServer):
|
||||
daemon_threads = True
|
||||
allow_reuse_address = True
|
||||
|
||||
remote: typing.Tuple[str, int]
|
||||
ticket: str
|
||||
running: bool
|
||||
stop_flag: threading.Event
|
||||
timeout: int
|
||||
check_certificate: bool
|
||||
status: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
remote: typing.Tuple[str, int],
|
||||
ticket: str,
|
||||
timeout: int = 0,
|
||||
local_port: int = 0,
|
||||
check_certificate: bool = True,
|
||||
) -> None:
|
||||
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
|
||||
self.check_certificate = check_certificate
|
||||
self.stop_flag = threading.Event() # False initial
|
||||
self.running = True
|
||||
|
||||
self.status = TUNNEL_LISTENING
|
||||
|
||||
def stop(self) -> None:
|
||||
if not self.stop_flag.is_set():
|
||||
self.stop_flag.set()
|
||||
self.running = False
|
||||
self.shutdown()
|
||||
|
||||
|
||||
class Handler(socketserver.BaseRequestHandler):
|
||||
# Override Base type
|
||||
server: ForwardServer
|
||||
|
||||
# server: ForwardServer
|
||||
def handle(self) -> None:
|
||||
# If server processing is timed out...
|
||||
if self.server.timeout and int(time.time()) > self.server.timeout:
|
||||
self.request.close() # End connection without processing it
|
||||
return
|
||||
|
||||
# Open remote connection
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as rsocket:
|
||||
rsocket.connect(self.server.remote)
|
||||
|
||||
context = ssl.create_default_context()
|
||||
|
||||
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(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}')
|
||||
self.server.status = TUNNEL_ERROR
|
||||
|
||||
# Processes data forwarding
|
||||
def process(self, remote: ssl.SSLSocket):
|
||||
# Process data until stop requested or connection closed
|
||||
try:
|
||||
while not self.server.stop_flag.is_set():
|
||||
r, _w, _x = select.select([self.request, remote], [], [], 1.0)
|
||||
if self.request in r:
|
||||
data = self.request.recv(BUFFER_SIZE)
|
||||
if not data:
|
||||
break
|
||||
remote.sendall(data)
|
||||
if remote in r:
|
||||
data = remote.recv(BUFFER_SIZE)
|
||||
if not data:
|
||||
break
|
||||
self.request.sendall(data)
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
def _run(server: ForwardServer) -> None:
|
||||
logger.debug('Starting server')
|
||||
server.serve_forever()
|
||||
logger.debug('Stoped server')
|
||||
|
||||
def forward(
|
||||
remote: typing.Tuple[str, int],
|
||||
ticket: str,
|
||||
timeout: int = 0,
|
||||
local_port: int = 0,
|
||||
check_certificate=True,
|
||||
) -> ForwardServer:
|
||||
fs = ForwardServer(
|
||||
remote=remote,
|
||||
ticket=ticket,
|
||||
timeout=timeout,
|
||||
local_port=local_port,
|
||||
check_certificate=check_certificate,
|
||||
)
|
||||
# Starts a new thread
|
||||
threading.Thread(target=_run, args=(fs,)).start()
|
||||
|
||||
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()
|
||||
|
0
tunnel-server/src/uds_tunnel/__init__.py
Normal file
0
tunnel-server/src/uds_tunnel/__init__.py
Normal file
95
tunnel-server/src/uds_tunnel/config.py
Normal file
95
tunnel-server/src/uds_tunnel/config.py
Normal file
@ -0,0 +1,95 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import hashlib
|
||||
import multiprocessing
|
||||
import configparser
|
||||
import logging
|
||||
import typing
|
||||
|
||||
from .consts import CONFIGFILE
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ConfigurationType(typing.NamedTuple):
|
||||
log_level: str
|
||||
log_file: str
|
||||
|
||||
listen_address: str
|
||||
listen_port: int
|
||||
|
||||
workers: int
|
||||
|
||||
ssl_certificate: str
|
||||
ssl_certificate_key: str
|
||||
ssl_ciphers: str
|
||||
ssl_dhparam: str
|
||||
|
||||
uds_server: str
|
||||
|
||||
secret: str
|
||||
allow: typing.Set[str]
|
||||
|
||||
storage: str
|
||||
|
||||
|
||||
def read() -> ConfigurationType:
|
||||
with open(CONFIGFILE, 'r') as f:
|
||||
config_str = '[uds]\n' + f.read()
|
||||
|
||||
cfg = configparser.ConfigParser()
|
||||
cfg.read_string(config_str)
|
||||
|
||||
uds = cfg['uds']
|
||||
|
||||
h = hashlib.sha256()
|
||||
h.update(uds.get('secret', '').encode())
|
||||
secret = h.hexdigest()
|
||||
|
||||
try:
|
||||
return ConfigurationType(
|
||||
log_level=uds.get('loglevel', 'ERROR'),
|
||||
log_file=uds.get('logfile', ''),
|
||||
listen_address=uds.get('address', '0.0.0.0'),
|
||||
listen_port=int(uds.get('port', '443')),
|
||||
workers=int(uds.get('workers', '0')) or multiprocessing.cpu_count(),
|
||||
ssl_certificate=uds['ssl_certificate'],
|
||||
ssl_certificate_key=uds['ssl_certificate_key'],
|
||||
ssl_ciphers=uds.get('ssl_ciphers'),
|
||||
ssl_dhparam=uds.get('ssl_dhparam'),
|
||||
uds_server=uds['uds_server'],
|
||||
secret=secret,
|
||||
allow=set(uds.get('allow', '127.0.0.1').split(',')),
|
||||
storage=uds['storage']
|
||||
)
|
||||
except ValueError as e:
|
||||
raise Exception(f'Mandatory configuration file in incorrect format: {e.args[0]}. Please, revise {CONFIGFILE}')
|
||||
except KeyError as e:
|
||||
raise Exception(f'Mandatory configuration parameter not found: {e.args[0]}. Please, revise {CONFIGFILE}')
|
60
tunnel-server/src/uds_tunnel/consts.py
Normal file
60
tunnel-server/src/uds_tunnel/consts.py
Normal file
@ -0,0 +1,60 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import string
|
||||
|
||||
DEBUG = True
|
||||
|
||||
if DEBUG:
|
||||
CONFIGFILE = 'udstunnel.cfg'
|
||||
LOGFORMAT = '%(levelname)s %(asctime)s %(message)s'
|
||||
else:
|
||||
CONFIGFILE = '/etc/udstunnel.cfg'
|
||||
LOGFORMAT = '%(levelname)s %(asctime)s %(message)s'
|
||||
|
||||
# MAX Length of read buffer for proxy
|
||||
BUFFER_SIZE = 1024 * 16
|
||||
# Handshake for conversation start
|
||||
HANDSHAKE_V1 = b'\x5AMGB\xA5\x01\x00'
|
||||
# Ticket length
|
||||
TICKET_LENGTH = 64
|
||||
# Admin password length, (size of an hex sha256)
|
||||
PASSWORD_LENGTH = 64
|
||||
# Bandwidth calc time lapse
|
||||
BANDWIDTH_TIME = 10
|
||||
|
||||
# Commands LENGTH (all same lenght)
|
||||
COMMAND_LENGTH = 4
|
||||
|
||||
# Valid commands
|
||||
COMMAND_OPEN = b'OPEN'
|
||||
COMMAND_TEST = b'TEST'
|
||||
COMMAND_STAT = b'STAT' # full stats
|
||||
COMMAND_INFO = b'INFO' # Basic stats
|
45
tunnel-server/src/uds_tunnel/message.py
Normal file
45
tunnel-server/src/uds_tunnel/message.py
Normal file
@ -0,0 +1,45 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2021 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import enum
|
||||
import socket
|
||||
import typing
|
||||
|
||||
class Command(enum.IntEnum):
|
||||
TUNNEL = 0
|
||||
STATS = 1
|
||||
|
||||
class Message:
|
||||
command: Command
|
||||
connection: typing.Optional[typing.Tuple[socket.socket, typing.Any]]
|
||||
|
||||
def __init__(self, command: Command, connection: typing.Optional[typing.Tuple[socket.socket, typing.Any]]):
|
||||
self.command = command
|
||||
self.connection = connection
|
191
tunnel-server/src/uds_tunnel/proxy.py
Normal file
191
tunnel-server/src/uds_tunnel/proxy.py
Normal file
@ -0,0 +1,191 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import logging
|
||||
import typing
|
||||
|
||||
import curio
|
||||
import requests
|
||||
|
||||
from . import config
|
||||
from . import stats
|
||||
from . import consts
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class Proxy:
|
||||
cfg: config.ConfigurationType
|
||||
stat: stats.Stats
|
||||
|
||||
def __init__(self, cfg: config.ConfigurationType) -> None:
|
||||
self.cfg = cfg
|
||||
self.stat = stats.Stats()
|
||||
|
||||
@staticmethod
|
||||
def getFromUds(cfg: config.ConfigurationType, ticket: bytes) -> typing.MutableMapping[str, typing.Any]:
|
||||
# Sanity checks
|
||||
if len(ticket) != consts.TICKET_LENGTH:
|
||||
raise Exception(f'TICKET INVALID (len={len(ticket)})')
|
||||
|
||||
for n, i in enumerate(ticket.decode(errors='ignore')):
|
||||
if (i >= 'a' and i <= 'z') or (i >= '0' and i <= '9') or (i >= 'A' and i <= 'Z'):
|
||||
continue # Correctus
|
||||
raise Exception(f'TICKET INVALID (char {i} at pos {n})')
|
||||
|
||||
# Gets the UDS connection data
|
||||
# r = requests.get(f'{cfg.uds_server}/XXXX/ticket')
|
||||
# if not r.ok:
|
||||
# raise Exception(f'TICKET INVALID (check {r.json})')
|
||||
return {
|
||||
'host': ['172.27.1.15', '172.27.0.10'][int(ticket[0]) - 0x30],
|
||||
'port': '3389'
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def doProxy(source, destination, counter: stats.StatsSingleCounter) -> None:
|
||||
while True:
|
||||
data = await source.recv(consts.BUFFER_SIZE)
|
||||
if not data:
|
||||
break
|
||||
await destination.sendall(data)
|
||||
counter.add(len(data))
|
||||
|
||||
async def stats(self, full: bool, source, address: typing.Tuple[str, int]) -> None:
|
||||
# Check valid source ip
|
||||
if address[0] not in self.cfg.allow:
|
||||
# Invalid source
|
||||
await source.sendall(b'FORBIDDEN')
|
||||
return
|
||||
|
||||
# Check password
|
||||
passwd = await source.recv(consts.PASSWORD_LENGTH)
|
||||
if passwd.decode(errors='ignore') != self.cfg.secret:
|
||||
# Invalid password
|
||||
await source.sendall(b'FORBIDDEN')
|
||||
return
|
||||
|
||||
logger.info('STATS TO %s', address)
|
||||
|
||||
if full:
|
||||
data = self.stat.full_as_csv()
|
||||
else:
|
||||
data = self.stat.simple_as_csv()
|
||||
|
||||
async for v in data:
|
||||
await source.sendall(v.encode() + b'\n')
|
||||
|
||||
# Method responsible of proxying requests
|
||||
async def __call__(self, source, address: typing.Tuple[str, int]) -> None:
|
||||
await self.proxy(source, address)
|
||||
|
||||
async def proxy(self, source, address: typing.Tuple[str, int]) -> None:
|
||||
logger.info('OPEN FROM %s', address)
|
||||
|
||||
try:
|
||||
# First, ensure handshake (simple handshake) and command
|
||||
data: bytes = await source.recv(len(consts.HANDSHAKE_V1))
|
||||
|
||||
if data != consts.HANDSHAKE_V1:
|
||||
raise Exception()
|
||||
except Exception:
|
||||
if consts.DEBUG:
|
||||
logger.exception('HANDSHAKE')
|
||||
logger.error('HANDSHAKE from %s', address)
|
||||
await source.sendall(b'HANDSHAKE_ERROR')
|
||||
|
||||
# Closes connection now
|
||||
return
|
||||
|
||||
try:
|
||||
# Handshake correct, get the command (4 bytes)
|
||||
command: bytes = await source.recv(consts.COMMAND_LENGTH)
|
||||
if command == consts.COMMAND_TEST:
|
||||
await source.sendall(b'OK')
|
||||
return
|
||||
|
||||
if command in (consts.COMMAND_STAT, consts.COMMAND_INFO):
|
||||
# This is an stats requests
|
||||
await self.stats(full=command==consts.COMMAND_STAT, source=source, address=address)
|
||||
return
|
||||
|
||||
if command != consts.COMMAND_OPEN:
|
||||
# Invalid command
|
||||
raise Exception()
|
||||
|
||||
# Now, read a TICKET_LENGTH (64) bytes string, that must be [a-zA-Z0-9]{64}
|
||||
ticket: bytes = await source.recv(consts.TICKET_LENGTH)
|
||||
|
||||
# Ticket received, now process it with UDS
|
||||
try:
|
||||
result = await curio.run_in_thread(Proxy.getFromUds, self.cfg, ticket)
|
||||
except Exception as e:
|
||||
logger.error('%s', e.args[0] if e.args else e)
|
||||
raise
|
||||
|
||||
print(f'Result: {result}')
|
||||
|
||||
# Invalid result from UDS, not allowed to connect
|
||||
if not result:
|
||||
raise Exception()
|
||||
|
||||
except Exception:
|
||||
if consts.DEBUG:
|
||||
logger.exception('COMMAND')
|
||||
logger.error('COMMAND from %s', address)
|
||||
await source.sendall(b'COMMAND_ERROR')
|
||||
return
|
||||
|
||||
# Communicate source OPEN is ok
|
||||
await source.sendall(b'OK')
|
||||
|
||||
# Initialize own stats counter
|
||||
counter = await self.stat.new()
|
||||
|
||||
# Open remote server connection
|
||||
try:
|
||||
destination = await curio.open_connection(result['host'], int(result['port']))
|
||||
async with curio.TaskGroup(wait=any) as grp:
|
||||
await grp.spawn(Proxy.doProxy, source, destination, counter.as_sent_counter())
|
||||
await grp.spawn(Proxy.doProxy, destination, source, counter.as_recv_counter())
|
||||
logger.debug('Launched proxies')
|
||||
|
||||
logger.debug('Proxies finalized: %s', grp.exceptions)
|
||||
|
||||
except Exception as e:
|
||||
if consts.DEBUG:
|
||||
logger.exception('OPEN REMOTE')
|
||||
|
||||
logger.error('REMOTE from %s: %s', address, e)
|
||||
finally:
|
||||
await counter.close()
|
||||
|
||||
|
||||
logger.info('CLOSED FROM %s', address)
|
||||
logger.info('STATS: %s', counter.as_csv())
|
258
tunnel-server/src/uds_tunnel/stats.py
Normal file
258
tunnel-server/src/uds_tunnel/stats.py
Normal file
@ -0,0 +1,258 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import time
|
||||
import io
|
||||
import ssl
|
||||
import logging
|
||||
import typing
|
||||
|
||||
import curio
|
||||
import blist
|
||||
|
||||
from . import config
|
||||
from . import consts
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Locker for id assigner
|
||||
assignLock = curio.Lock()
|
||||
|
||||
# Tuple index for several stats
|
||||
SENT, RECV = 0, 1
|
||||
|
||||
# Subclasses for += operation to work
|
||||
class StatsSingleCounter:
|
||||
def __init__(self, parent: 'StatsConnection', for_receiving=True) -> None:
|
||||
if for_receiving:
|
||||
self.adder = parent.add_recv
|
||||
else:
|
||||
self.adder = parent.add_sent
|
||||
|
||||
def add(self, value: int):
|
||||
self.adder(value)
|
||||
return self
|
||||
|
||||
|
||||
class StatsConnection:
|
||||
id: int
|
||||
recv: int
|
||||
sent: int
|
||||
start_time: int
|
||||
parent: 'Stats'
|
||||
|
||||
# Bandwidth stats (SENT, RECV)
|
||||
last: typing.List[int]
|
||||
last_time: typing.List[float]
|
||||
|
||||
bandwidth: typing.List[int]
|
||||
max_bandwidth: typing.List[int]
|
||||
|
||||
def __init__(self, parent: 'Stats', id: int) -> None:
|
||||
self.id = id
|
||||
self.recv = self.sent = 0
|
||||
|
||||
now = time.time()
|
||||
self.start_time = int(now)
|
||||
self.parent = parent
|
||||
|
||||
self.last = [0, 0]
|
||||
self.last_time = [now, now]
|
||||
self.bandwidth = [0, 0]
|
||||
self.max_bandwidth = [0, 0]
|
||||
|
||||
def update_bandwidth(self, kind: int, counter: int):
|
||||
now = time.time()
|
||||
elapsed = now - self.last_time[kind]
|
||||
# Update only when enouth data
|
||||
if elapsed < consts.BANDWIDTH_TIME:
|
||||
return
|
||||
total = counter - self.last[kind]
|
||||
self.bandwidth[kind] = int(float(total) / elapsed)
|
||||
self.last[kind] = counter
|
||||
self.last_time[kind] = now
|
||||
|
||||
if self.bandwidth[kind] > self.max_bandwidth[kind]:
|
||||
self.max_bandwidth[kind] = self.bandwidth[kind]
|
||||
|
||||
def add_recv(self, size: int) -> None:
|
||||
self.recv += size
|
||||
self.update_bandwidth(RECV, counter=self.recv)
|
||||
self.parent.add_recv(size)
|
||||
|
||||
def add_sent(self, size: int) -> None:
|
||||
self.sent += size
|
||||
self.update_bandwidth(SENT, counter=self.sent)
|
||||
self.parent.add_sent(size)
|
||||
|
||||
def as_sent_counter(self) -> 'StatsSingleCounter':
|
||||
return StatsSingleCounter(self, False)
|
||||
|
||||
def as_recv_counter(self) -> 'StatsSingleCounter':
|
||||
return StatsSingleCounter(self, True)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self.id:
|
||||
logger.debug(f'STAT {self.id} closed')
|
||||
await self.parent.remove(self.id)
|
||||
self.id = 0
|
||||
|
||||
def as_csv(self, separator: typing.Optional[str] = None) -> str:
|
||||
separator = separator or ';'
|
||||
# With connections of less than a second, consider them as a second
|
||||
elapsed = (int(time.time()) - self.start_time)
|
||||
|
||||
return separator.join(
|
||||
str(i)
|
||||
for i in (
|
||||
self.id,
|
||||
self.start_time,
|
||||
elapsed,
|
||||
self.sent,
|
||||
self.bandwidth[SENT],
|
||||
self.max_bandwidth[SENT],
|
||||
self.recv,
|
||||
self.bandwidth[RECV],
|
||||
self.max_bandwidth[RECV],
|
||||
)
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f'{self.id} t:{int(time.time())-self.start_time}, r:{self.recv}, s:{self.sent}>'
|
||||
|
||||
# For sorted array
|
||||
def __lt__(self, other) -> bool:
|
||||
if isinstance(other, int):
|
||||
return self.id < other
|
||||
|
||||
if not isinstance(other, StatsConnection):
|
||||
raise NotImplemented
|
||||
|
||||
return self.id < other.id
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
if isinstance(other, int):
|
||||
return self.id == other
|
||||
|
||||
if not isinstance(other, StatsConnection):
|
||||
raise NotImplemented
|
||||
|
||||
return self.id == other.id
|
||||
|
||||
|
||||
class Stats:
|
||||
counter_id: int
|
||||
|
||||
total_sent: int
|
||||
total_received: int
|
||||
current_connections: blist.sortedlist
|
||||
|
||||
def __init__(self) -> None:
|
||||
# First connection will be 1
|
||||
self.counter_id = 0
|
||||
self.total_sent = self.total_received = 0
|
||||
self.current_connections = blist.sortedlist()
|
||||
|
||||
async def new(self) -> StatsConnection:
|
||||
"""Initializes a connection stats counter and returns it id
|
||||
|
||||
Returns:
|
||||
str: connection id
|
||||
"""
|
||||
async with assignLock:
|
||||
self.counter_id += 1
|
||||
connection = StatsConnection(self, self.counter_id)
|
||||
self.current_connections.add(connection)
|
||||
return connection
|
||||
|
||||
def add_sent(self, size: int) -> None:
|
||||
self.total_sent += size
|
||||
|
||||
def add_recv(self, size: int) -> None:
|
||||
self.total_received += size
|
||||
|
||||
async def remove(self, connection_id: int) -> None:
|
||||
async with assignLock:
|
||||
try:
|
||||
self.current_connections.remove(connection_id)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
'Tried to remove %s from connections but was not present',
|
||||
connection_id,
|
||||
)
|
||||
# Does not exists, ignore it
|
||||
pass
|
||||
|
||||
async def simple_as_csv(self, separator: typing.Optional[str] = None) -> typing.AsyncIterable[str]:
|
||||
separator = separator or ';'
|
||||
yield separator.join(
|
||||
str(i)
|
||||
for i in (
|
||||
self.counter_id,
|
||||
self.total_sent,
|
||||
self.total_received,
|
||||
len(self.current_connections),
|
||||
)
|
||||
)
|
||||
|
||||
async def full_as_csv(self, separator: typing.Optional[str] = None) -> typing.AsyncIterable[str]:
|
||||
for i in self.current_connections:
|
||||
yield i.as_csv(separator)
|
||||
|
||||
|
||||
# Stats processor, invoked from command line
|
||||
async def getServerStats(detailed: bool = False) -> None:
|
||||
cfg = config.read()
|
||||
|
||||
# Context for local connection (ignores cert hostname)
|
||||
context = ssl.create_default_context()
|
||||
context.check_hostname = False
|
||||
|
||||
try:
|
||||
host = cfg.listen_address if cfg.listen_address != '0.0.0.0' else 'localhost'
|
||||
sock = await curio.open_connection(
|
||||
host, cfg.listen_port, ssl=context, server_hostname='localhost'
|
||||
)
|
||||
tmpdata = io.BytesIO()
|
||||
cmd = consts.COMMAND_STAT if detailed else consts.COMMAND_INFO
|
||||
async with sock:
|
||||
await sock.sendall(consts.HANDSHAKE_V1 + cmd + cfg.secret.encode())
|
||||
while True:
|
||||
chunk = await sock.recv(consts.BUFFER_SIZE)
|
||||
if not chunk:
|
||||
break
|
||||
tmpdata.write(chunk)
|
||||
|
||||
# Now we can output chunk data
|
||||
print(tmpdata.getvalue().decode())
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return
|
34
tunnel-server/src/udstunnel.cfg
Normal file
34
tunnel-server/src/udstunnel.cfg
Normal file
@ -0,0 +1,34 @@
|
||||
# Sample testing UDS tunnel configuration
|
||||
|
||||
# Log level, valid are DEBUG, INFO, WARN, ERROR
|
||||
loglevel = DEBUG
|
||||
|
||||
# Listen address. Defaults to 0.0.0.0
|
||||
address = 0.0.0.0
|
||||
|
||||
# Number of workers. Defaults to 0 (means "as much as cores")
|
||||
workers = 2
|
||||
|
||||
# Listening port
|
||||
port = 7777
|
||||
|
||||
# SSL Related parameters
|
||||
ssl_certificate = tests/testing.pem
|
||||
ssl_certificate_key = tests/testing.key
|
||||
ssl_ciphers = ECDHE-RSA-AES256-GCM-SHA512:DHE-RSA-AES256-GCM-SHA512:ECDHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-SHA384
|
||||
ssl_dhparam = /etc/certs/dhparam.pem
|
||||
|
||||
# UDS server location
|
||||
uds_server = http://172.27.0.1:8000
|
||||
|
||||
# Secret to get access to admin commands
|
||||
# Admin commands and only allowed from localhost
|
||||
# So, in order to allow this commands, ensure listen address allows connections from localhost
|
||||
secret = MySecret
|
||||
|
||||
# List of af allowed admin commands ips (only IPs, no networks or whatever)
|
||||
# defaults to localhost (change if listen address is different from 0.0.0.0)
|
||||
allow = 127.0.0.1
|
||||
|
||||
# Local storage configuration, for stats, etc...
|
||||
storage = .
|
203
tunnel-server/src/udstunnel.py
Executable file
203
tunnel-server/src/udstunnel.py
Executable file
@ -0,0 +1,203 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
# Copyright (c) 2020 Virtual Cable S.L.U.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without modification,
|
||||
# are permitted provided that the following conditions are met:
|
||||
#
|
||||
# * Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
# * Neither the name of Virtual Cable S.L. nor the names of its contributors
|
||||
# may be used to endorse or promote products derived from this software
|
||||
# without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
'''
|
||||
@author: Adolfo Gómez, dkmaster at dkmon dot com
|
||||
'''
|
||||
import sys
|
||||
import argparse
|
||||
import multiprocessing
|
||||
import socket
|
||||
import logging
|
||||
import typing
|
||||
|
||||
import curio
|
||||
import psutil
|
||||
|
||||
from uds_tunnel import config
|
||||
from uds_tunnel import proxy
|
||||
from uds_tunnel import consts
|
||||
from uds_tunnel import message
|
||||
from uds_tunnel import stats
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from multiprocessing.connection import Connection
|
||||
|
||||
BACKLOG = 100
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def setup_log(cfg: config.ConfigurationType) -> None:
|
||||
# Setup basic logging
|
||||
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)
|
||||
|
||||
# Update logging if needed
|
||||
if cfg.log_file:
|
||||
fileh = logging.FileHandler(cfg.log_file, 'a')
|
||||
formatter = logging.Formatter(consts.LOGFORMAT)
|
||||
fileh.setFormatter(formatter)
|
||||
log = logging.getLogger()
|
||||
for hdlr in log.handlers[:]:
|
||||
log.removeHandler(hdlr)
|
||||
log.addHandler(fileh)
|
||||
|
||||
|
||||
async def tunnel_proc_async(pipe: 'Connection', cfg: config.ConfigurationType) -> None:
|
||||
def get_socket(pipe: 'Connection') -> typing.Tuple[socket.SocketType, typing.Any]:
|
||||
try:
|
||||
while True:
|
||||
msg: message.Message = pipe.recv()
|
||||
if msg.command == message.Command.TUNNEL and msg.connection:
|
||||
return msg.connection
|
||||
# Process other messages, and retry
|
||||
except Exception:
|
||||
logger.exception('Receiving data from parent process')
|
||||
return None, None
|
||||
|
||||
async def run_server(
|
||||
pipe: 'Connection', cfg: config.ConfigurationType, group: curio.TaskGroup
|
||||
) -> None:
|
||||
# Instantiate a proxy redirector for this process (we only need one per process!!)
|
||||
tunneler = proxy.Proxy(cfg)
|
||||
|
||||
# Generate SSL context
|
||||
context = curio.ssl.SSLContext(curio.ssl.PROTOCOL_TLS_SERVER)
|
||||
context.load_cert_chain(cfg.ssl_certificate, cfg.ssl_certificate_key)
|
||||
|
||||
if cfg.ssl_ciphers:
|
||||
context.set_ciphers(cfg.ssl_ciphers)
|
||||
|
||||
if cfg.ssl_dhparam:
|
||||
context.load_dh_params(cfg.ssl_dhparam)
|
||||
|
||||
while True:
|
||||
sock, address = await curio.run_in_thread(get_socket, pipe)
|
||||
if not sock:
|
||||
break
|
||||
logger.debug(
|
||||
f'{multiprocessing.current_process().pid!r}: Got new connection from {address!r}'
|
||||
)
|
||||
sock = await context.wrap_socket(curio.io.Socket(sock), server_side=True)
|
||||
await group.spawn(tunneler, sock, address)
|
||||
del sock
|
||||
|
||||
async with curio.TaskGroup() as tg:
|
||||
await tg.spawn(run_server, pipe, cfg, tg)
|
||||
# Reap all of the children tasks as they complete
|
||||
async for task in tg:
|
||||
logger.debug(f'Deleting {task!r}')
|
||||
task.joined = True
|
||||
del task
|
||||
|
||||
|
||||
def tunnel_main():
|
||||
cfg = config.read()
|
||||
setup_log(cfg)
|
||||
|
||||
# Creates as many processes and pipes as required
|
||||
child: typing.List[
|
||||
typing.Tuple['Connection', multiprocessing.Process, psutil.Process]
|
||||
] = []
|
||||
|
||||
for i in range(cfg.workers):
|
||||
own_conn, child_conn = multiprocessing.Pipe()
|
||||
task = multiprocessing.Process(
|
||||
target=curio.run, args=(tunnel_proc_async, child_conn, cfg)
|
||||
)
|
||||
task.start()
|
||||
child.append((own_conn, task, psutil.Process(task.pid)))
|
||||
|
||||
def best_child() -> 'Connection':
|
||||
best: typing.Tuple[float, 'Connection'] = (1000.0, child[0][0])
|
||||
for c in child:
|
||||
percent = c[2].cpu_percent()
|
||||
logger.debug('PID %s has %s', c[2].pid, percent)
|
||||
if percent < best[0]:
|
||||
best = (percent, c[0])
|
||||
return best[1]
|
||||
|
||||
sock = None
|
||||
try:
|
||||
# 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', exc_info=True)
|
||||
|
||||
sock.bind((cfg.listen_address, cfg.listen_port))
|
||||
sock.listen(BACKLOG)
|
||||
while True:
|
||||
client, addr = sock.accept()
|
||||
# Select BEST process for sending this new connection
|
||||
best_child().send(message.Message(message.Command.TUNNEL, (client, addr)))
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if sock:
|
||||
sock.close()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
group = parser.add_mutually_exclusive_group()
|
||||
group.add_argument(
|
||||
'-t', '--tunnel', help='Starts the tunnel server', action='store_true'
|
||||
)
|
||||
group.add_argument(
|
||||
'-s',
|
||||
'--stats',
|
||||
help='get current global stats from RUNNING tunnel',
|
||||
action='store_true',
|
||||
)
|
||||
group.add_argument(
|
||||
'-d',
|
||||
'--detailed-stats',
|
||||
help='get current detailed stats from RUNNING tunnel',
|
||||
action='store_true',
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.tunnel:
|
||||
tunnel_main()
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
Loading…
Reference in New Issue
Block a user