diff --git a/client/src/UDSClient.py b/client/src/UDSClient.py index 0a845f7c..b037d578 100644 --- a/client/src/UDSClient.py +++ b/client/src/UDSClient.py @@ -178,6 +178,8 @@ if __name__ == "__main__": exitVal = app.exec_() + tools.waitForTasks() + time.sleep(3) tools.unlinkFiles() diff --git a/client/src/uds/forward.py b/client/src/uds/forward.py index 9df02fd4..010cc398 100644 --- a/client/src/uds/forward.py +++ b/client/src/uds/forward.py @@ -7,10 +7,9 @@ from __future__ import unicode_literals import select import SocketServer -import sys - import paramiko import threading +import random g_verbose = True @@ -23,6 +22,8 @@ class ForwardServer (SocketServer.ThreadingTCPServer): class Handler (SocketServer.BaseRequestHandler): def handle(self): + self.thread.alreadyConnected = True + try: chan = self.ssh_transport.open_channel('direct-tcpip', (self.chain_host, self.chain_port), @@ -39,24 +40,32 @@ class Handler (SocketServer.BaseRequestHandler): verbose('Connected! Tunnel open %r -> %r -> %r' % (self.request.getpeername(), chan.getpeername(), (self.chain_host, self.chain_port))) - while self.event.is_set() is False: - r, w, x = select.select([self.request, chan], [], [], 1) + try: + while self.event.is_set() is False: + r, _w, _x = select.select([self.request, chan], [], [], 1) - if self.request in r: - data = self.request.recv(1024) - if len(data) == 0: - break - chan.send(data) - if chan in r: - data = chan.recv(1024) - if len(data) == 0: - break - self.request.send(data) + if self.request in r: + data = self.request.recv(1024) + if len(data) == 0: + break + chan.send(data) + if chan in r: + data = chan.recv(1024) + if len(data) == 0: + break + self.request.send(data) + except Exception: + pass - peername = self.request.getpeername() - chan.close() - self.request.close() - verbose('Tunnel closed from %r' % (peername,)) + try: + peername = self.request.getpeername() + chan.close() + self.request.close() + verbose('Tunnel closed from %r' % (peername,)) + except Exception: + pass + + self.thread.stop() def verbose(s): @@ -65,7 +74,7 @@ def verbose(s): class ForwardThread(threading.Thread): - def __init__(self, server, port, username, password, localPort, redirectHost, redirectPort): + def __init__(self, server, port, username, password, localPort, redirectHost, redirectPort, waitTime): threading.Thread.__init__(self) self.client = None self.fs = None @@ -79,8 +88,19 @@ class ForwardThread(threading.Thread): self.redirectHost = redirectHost self.redirectPort = redirectPort + self.waitTime = waitTime + self.stopEvent = threading.Event() + self.timer = None + self.alreadyConnected = False + + def _timerFnc(self): + self.timer = None + verbose('Timer fnc: {}'.format(self.alreadyConnected)) + if self.alreadyConnected is False: + self.stop() + def run(self): self.client = paramiko.SSHClient() self.client.load_system_host_keys() @@ -90,17 +110,25 @@ class ForwardThread(threading.Thread): self.client.connect(self.server, self.port, username=self.username, password=self.password) - class SubHander (Handler): + class SubHandler (Handler): chain_host = self.redirectHost chain_port = self.redirectPort ssh_transport = self.client.get_transport() event = self.stopEvent + thread = self - self.fs = ForwardServer(('', self.redirectPort), SubHander) + verbose('Wait Time: {}'.format(self.waitTime)) + self.timer = threading.Timer(self.waitTime, self._timerFnc) + self.timer.start() + + self.fs = ForwardServer(('', self.localPort), SubHandler) self.fs.serve_forever() def stop(self): try: + if self.timer: + self.timer.cancel() + self.stopEvent.set() self.fs.shutdown() @@ -110,14 +138,21 @@ class ForwardThread(threading.Thread): pass -def forward(server, port, username, password, localPort, redirectHost, redirectPort): +def forward(server, port, username, password, redirectHost, redirectPort, localPort=None, waitTime=10): + ''' + Instantiates an ssh connection to server:port + Returns the Thread created and the local redirected port as a list: (thread, port) + ''' port, redirectPort = int(port), int(redirectPort) + if localPort is None: + localPort = random.randrange(40000, 50000) - verbose('Connected') + verbose('Connecting to {}:{} using {}/{} redirecting to {}:{}, listening on 127.0.0.1:{}'.format( + server, port, username, password, redirectHost, redirectPort, localPort)) - ft = ForwardThread(server, port, username, password, localPort, redirectHost, redirectPort) + ft = ForwardThread(server, port, username, password, localPort, redirectHost, redirectPort, waitTime) ft.start() - return ft + return (ft, localPort) diff --git a/client/src/uds/tools.py b/client/src/uds/tools.py index 352d0706..a514a313 100644 --- a/client/src/uds/tools.py +++ b/client/src/uds/tools.py @@ -38,6 +38,8 @@ import random import os _unlinkFiles = [] +_tasksToWait = [] +_execBeforeExit = [] def saveTempFile(content, filename=None): @@ -52,12 +54,36 @@ def saveTempFile(content, filename=None): def addFileToUnlink(filename): + ''' + Adds a file to the wait-and-unlink list + ''' _unlinkFiles.append(filename) def unlinkFiles(): + ''' + Removes all wait-and-unlink files + ''' for f in _unlinkFiles: try: os.unlink(f) except Exception: pass + + +def addTaskToWait(taks): + _tasksToWait.append(taks) + + +def waitForTasks(): + for t in _tasksToWait: + t.join() + + +def addExecBeforeExit(fnc): + _execBeforeExit.append(fnc) + + +def execBeforeExit(): + for fnc in _execBeforeExit: + fnc.call()