Update handler.py

fix: socket moved to ssh_connect

webssh/handler.py

@@ -401,7 +401,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         privatekey, filename = self.get_privatekey()
         passphrase = self.get_argument('passphrase', u'')
         totp = self.get_argument('totp', u'')
-        source_address = self.get_source_address()
+        source_address = self.get_argument('source_address', u'')
 
         if isinstance(self.policy, paramiko.RejectPolicy):
             self.lookup_hostname(hostname, port)
@@ -411,19 +411,8 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         else:
             pkey = None
 
-        if source_address:
-            logging.info("Binding socket for source ip {}".format(source_address))
-            sock = socket.socket()
-            sock.settimeout(options.timeout)  # Set a timeout on blocking socket operations
-            try:
-                sock.bind((source_address, 0))
-            except OSError:
-                raise InvalidValueError('Unable to bind source address {} socket'.format(source_address))
-        else:
-            sock = None
-
         self.ssh_client.totp = totp
-        args = (hostname, port, username, password, pkey, sock)
+        args = (hostname, port, username, password, pkey, source_address)
         logging.debug(args)
 
         return args
@@ -469,16 +458,23 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         dst_addr = args[:2]
         logging.info('Connecting to {}:{}'.format(*dst_addr))
 
-        sock = args[5]
-        if sock:
+        sock = None
+        source_address = args[5]
+        if source_address:
             logging.info('Connecting source address socket')
+            sock = socket.socket()
+            sock.settimeout(options.timeout)  # Set a timeout on blocking socket operations
+            try:
+                sock.bind((source_address, 0))
+            except OSError:
+                raise InvalidValueError('Unable to bind source address {} socket'.format(source_address))
             try:
                 sock.connect(dst_addr)
             except socket.error:
                 raise ValueError('Unable to connect source address socket to {}:{}'.format(*dst_addr))
 
         try:
-            ssh.connect(*args, timeout=options.timeout)
+            ssh.connect(*args, sock=sock, timeout=options.timeout)
         except socket.error:
             raise ValueError('Unable to connect to {}:{}'.format(*dst_addr))
         except paramiko.BadAuthenticationType:
@@ -626,4 +622,4 @@ class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
 
         worker = self.worker_ref() if self.worker_ref else None
         if worker:
-            worker.close(reason=self.close_reason)
+            worker.close(reason=self.close_reason)