Advancing on client for mac os x

client/src/UDSClient.py

@@ -178,6 +178,8 @@ if __name__ == "__main__":
 
         exitVal = app.exec_()
 
+        tools.waitForTasks()
+
         time.sleep(3)
         tools.unlinkFiles()
 

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:
+        try:
-            r, w, x = select.select([self.request, chan], [], [], 1)
+            while self.event.is_set() is False:
+                r, _w, _x = select.select([self.request, chan], [], [], 1)
 
-            if self.request in r:
+                if self.request in r:
-                data = self.request.recv(1024)
+                    data = self.request.recv(1024)
-                if len(data) == 0:
+                    if len(data) == 0:
-                    break
+                        break
-                chan.send(data)
+                    chan.send(data)
-            if chan in r:
+                if chan in r:
-                data = chan.recv(1024)
+                    data = chan.recv(1024)
-                if len(data) == 0:
+                    if len(data) == 0:
-                    break
+                        break
-                self.request.send(data)
+                    self.request.send(data)
+        except Exception:
+            pass
 
-        peername = self.request.getpeername()
+        try:
-        chan.close()
+            peername = self.request.getpeername()
-        self.request.close()
+            chan.close()
-        verbose('Tunnel closed from %r' % (peername,))
+            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)

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()