Some minor cosmetic changes for UDSClient

This commit is contained in:
Adolfo Gómez García 2021-06-22 17:12:31 +02:00
parent 174d836f45
commit d1fb59ab77
3 changed files with 13 additions and 54 deletions

View File

@ -9,6 +9,7 @@ import random
import time import time
import select import select
import socketserver import socketserver
import typing
import paramiko import paramiko
@ -36,6 +37,11 @@ class ForwardServer(socketserver.ThreadingTCPServer):
class Handler(socketserver.BaseRequestHandler): class Handler(socketserver.BaseRequestHandler):
event: threading.Event
thread: 'ForwardThread'
ssh_transport: paramiko.Transport
chain_host: str
chain_port: int
def handle(self): def handle(self):
self.thread.currentConnections += 1 self.thread.currentConnections += 1
@ -86,6 +92,7 @@ class Handler(socketserver.BaseRequestHandler):
class ForwardThread(threading.Thread): class ForwardThread(threading.Thread):
status = 0 # Connecting status = 0 # Connecting
client: typing.Optional[paramiko.SSHClient]
def __init__(self, server, port, username, password, localPort, redirectHost, redirectPort, waitTime, fingerPrints): def __init__(self, server, port, username, password, localPort, redirectHost, redirectPort, waitTime, fingerPrints):
threading.Thread.__init__(self) threading.Thread.__init__(self)
@ -118,7 +125,7 @@ class ForwardThread(threading.Thread):
ft = ForwardThread(self.server, self.port, self.username, self.password, localPort, redirectHost, redirectPort, self.waitTime, self.fingerPrints) ft = ForwardThread(self.server, self.port, self.username, self.password, localPort, redirectHost, redirectPort, self.waitTime, self.fingerPrints)
ft.client = self.client ft.client = self.client
self.client.useCount += 1 # One more using this client self.client.useCount += 1 # type: ignore
ft.start() ft.start()
while ft.status == 0: while ft.status == 0:
@ -138,7 +145,7 @@ class ForwardThread(threading.Thread):
if self.client is None: if self.client is None:
try: try:
self.client = paramiko.SSHClient() self.client = paramiko.SSHClient()
self.client.useCount = 1 # Custom added variable, to keep track on when to close tunnel self.client.useCount = 1 # type: ignore
self.client.load_system_host_keys() self.client.load_system_host_keys()
self.client.set_missing_host_key_policy(CheckfingerPrints(self.fingerPrints)) self.client.set_missing_host_key_policy(CheckfingerPrints(self.fingerPrints))
@ -176,8 +183,8 @@ class ForwardThread(threading.Thread):
self.fs.shutdown() self.fs.shutdown()
if self.client is not None: if self.client is not None:
self.client.useCount -= 1 self.client.useCount -= 1 # type: ignore
if self.client.useCount == 0: if self.client.useCount == 0: # type: ignore
self.client.close() self.client.close()
self.client = None # Clean up self.client = None # Clean up
except Exception: except Exception:

View File

@ -1,48 +0,0 @@
# -*- coding: utf-8 -*-
#
# Copyright (c) 2017-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
'''
from __future__ import unicode_literals
import sys
LINUX = 'Linux'
WINDOWS = 'Windows'
MAC_OS_X = 'Mac os x'
def getOs():
if sys.platform.startswith('linux'):
return LINUX
if sys.platform.startswith('win'):
return WINDOWS
if sys.platform.startswith('darwin'):
return MAC_OS_X
return 'other'

View File

@ -46,7 +46,7 @@ import certifi
from cryptography import x509 from cryptography import x509
from cryptography.hazmat.backends import default_backend from cryptography.hazmat.backends import default_backend
from . import osDetector from . import os_detector
from . import tools from . import tools
from . import VERSION from . import VERSION
from .log import logger from .log import logger
@ -206,7 +206,7 @@ class RestApi:
def urlopen(url: str): def urlopen(url: str):
# Generate the request with the headers # Generate the request with the headers
req = urllib.request.Request(url, headers={ req = urllib.request.Request(url, headers={
'User-Agent': osDetector.getOs() + " - UDS Connector " + VERSION 'User-Agent': os_detector.getOs() + " - UDS Connector " + VERSION
}) })
return urllib.request.urlopen(req, context=ctx) return urllib.request.urlopen(req, context=ctx)