Repository: jupyter/jupyter_client
Pull request #335 (2018-01-10 10:39:00)

Migrate SSH tunneling from pyzmq
As discussed in https://github.com/ipython/ipython/issues/10934

As this part of code is mainly used in IPython we
agreed to move it from pyzmq. It will be easier to
maintain here and some planned changes in this code
will be easier to apply and release.

Also some code reformatting was applied in `connect.py` file.


Patch:
diff --git a/jupyter_client/connect.py b/jupyter_client/connect.py
index 91efbc4..2cc89de 100644
--- a/jupyter_client/connect.py
+++ b/jupyter_client/connect.py
@@ -35,9 +35,9 @@ from jupyter_core.paths import jupyter_data_dir, jupyter_runtime_dir
 
 
 def write_connection_file(fname=None, shell_port=0, iopub_port=0, stdin_port=0, hb_port=0,
-                         control_port=0, ip='', key=b'', transport='tcp',
-                         signature_scheme='hmac-sha256', kernel_name=''
-                         ):
+                          control_port=0, ip='', key=b'', transport='tcp',
+                          signature_scheme='hmac-sha256', kernel_name=''
+                          ):
     """Generates a JSON config file, including the selection of random ports.
 
     Parameters
@@ -193,7 +193,7 @@ def find_connection_file(filename='kernel-*.json', path=None, profile=None):
         path = ['.', jupyter_runtime_dir()]
     if isinstance(path, string_types):
         path = [path]
-    
+
     try:
         # first, try explicit name
         return filefind(filename, path)
@@ -208,11 +208,11 @@ def find_connection_file(filename='kernel-*.json', path=None, profile=None):
     else:
         # accept any substring match
         pat = '*%s*' % filename
-    
+
     matches = []
     for p in path:
         matches.extend(glob.glob(os.path.join(p, pat)))
-    
+
     matches = [ os.path.abspath(m) for m in matches ]
     if not matches:
         raise IOError("Could not find %r in %r" % (filename, path))
@@ -249,7 +249,7 @@ def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
     (shell, iopub, stdin, hb) : ints
         The four ports on localhost that have been forwarded to the kernel.
     """
-    from zmq.ssh import tunnel
+    from .ssh import tunnel
     if isinstance(connection_info, string_types):
         # it's a path, unpack it
         with open(connection_info) as f:
@@ -289,11 +289,11 @@ port_names = [ "%s_port" % channel for channel in ('shell', 'stdin', 'iopub', 'h
 
 class ConnectionFileMixin(LoggingConfigurable):
     """Mixin for configurable classes that work with connection files"""
-    
+
     data_dir = Unicode()
     def _data_dir_default(self):
         return jupyter_data_dir()
-    
+
     # The addresses for the communication channels
     connection_file = Unicode('', config=True,
     help="""JSON file in which to store connection info [default: kernel-<pid>.json]
@@ -480,7 +480,7 @@ class ConnectionFileMixin(LoggingConfigurable):
 
     def load_connection_file(self, connection_file=None):
         """Load connection info from JSON dict in self.connection_file.
-        
+
         Parameters
         ----------
         connection_file: unicode, optional
@@ -496,10 +496,10 @@ class ConnectionFileMixin(LoggingConfigurable):
 
     def load_connection_info(self, info):
         """Load connection info from a dict containing connection info.
-        
+
         Typically this data comes from a connection file
         and is called by load_connection_file.
-        
+
         Parameters
         ----------
         info: dict
diff --git a/jupyter_client/ssh/__init__.py b/jupyter_client/ssh/__init__.py
new file mode 100644
index 0000000..d7bc9d5
--- /dev/null
+++ b/jupyter_client/ssh/__init__.py
@@ -0,0 +1,1 @@
+from jupyter_client.ssh.tunnel import *
diff --git a/jupyter_client/ssh/forward.py b/jupyter_client/ssh/forward.py
new file mode 100644
index 0000000..a44c117
--- /dev/null
+++ b/jupyter_client/ssh/forward.py
@@ -0,0 +1,92 @@
+#
+# This file is adapted from a paramiko demo, and thus licensed under LGPL 2.1.
+# Original Copyright (C) 2003-2007  Robey Pointer <robeypointer@gmail.com>
+# Edits Copyright (C) 2010 The IPython Team
+#
+# Paramiko is free software; you can redistribute it and/or modify it under the
+# terms of the GNU Lesser General Public License as published by the Free
+# Software Foundation; either version 2.1 of the License, or (at your option)
+# any later version.
+#
+# Paramiko is distrubuted in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
+# details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Paramiko; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA  02111-1301  USA.
+
+"""
+Sample script showing how to do local port forwarding over paramiko.
+
+This script connects to the requested SSH server and sets up local port
+forwarding (the openssh -L option) from a local port through a tunneled
+connection to a destination reachable from the SSH server machine.
+"""
+
+from __future__ import print_function
+
+import logging
+import select
+try:  # Python 3
+    import socketserver
+except ImportError:  # Python 2
+    import SocketServer as socketserver
+
+logger = logging.getLogger('ssh')
+
+
+class ForwardServer (socketserver.ThreadingTCPServer):
+    daemon_threads = True
+    allow_reuse_address = True
+
+
+class Handler (socketserver.BaseRequestHandler):
+
+    def handle(self):
+        try:
+            chan = self.ssh_transport.open_channel('direct-tcpip',
+                                                   (self.chain_host, self.chain_port),
+                                                   self.request.getpeername())
+        except Exception as e:
+            logger.debug('Incoming request to %s:%d failed: %s' % (self.chain_host,
+                                                                   self.chain_port,
+                                                                   repr(e)))
+            return
+        if chan is None:
+            logger.debug('Incoming request to %s:%d was rejected by the SSH server.' %
+                    (self.chain_host, self.chain_port))
+            return
+
+        logger.debug('Connected!  Tunnel open %r -> %r -> %r' % (self.request.getpeername(),
+                                                            chan.getpeername(), (self.chain_host, self.chain_port)))
+        while True:
+            r, w, x = select.select([self.request, chan], [], [])
+            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)
+        chan.close()
+        self.request.close()
+        logger.debug('Tunnel closed ')
+
+
+def forward_tunnel(local_port, remote_host, remote_port, transport):
+    # this is a little convoluted, but lets me configure things for the Handler
+    # object.  (SocketServer doesn't give Handlers any way to access the outer
+    # server normally.)
+    class SubHander (Handler):
+        chain_host = remote_host
+        chain_port = remote_port
+        ssh_transport = transport
+    ForwardServer(('127.0.0.1', local_port), SubHander).serve_forever()
+
+
+__all__ = ['forward_tunnel']
diff --git a/jupyter_client/ssh/tunnel.py b/jupyter_client/ssh/tunnel.py
new file mode 100644
index 0000000..e1cd080
--- /dev/null
+++ b/jupyter_client/ssh/tunnel.py
@@ -0,0 +1,375 @@
+"""Basic ssh tunnel utilities, and convenience functions for tunneling
+zeromq connections.
+"""
+
+# Copyright (C) 2010-2011  IPython Development Team
+# Copyright (C) 2011- PyZMQ Developers
+#
+# Redistributed from IPython under the terms of the BSD License.
+
+
+from __future__ import print_function
+
+import atexit
+import os
+import re
+import signal
+import socket
+import sys
+import warnings
+from getpass import getpass, getuser
+from multiprocessing import Process
+
+try:
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore', DeprecationWarning)
+        import paramiko
+        SSHException = paramiko.ssh_exception.SSHException
+except ImportError:
+    paramiko = None
+    class SSHException(Exception):
+        pass
+else:
+    from .forward import forward_tunnel
+
+try:
+    import pexpect
+except ImportError:
+    pexpect = None
+
+from zmq.utils.strtypes import b
+
+
+def select_random_ports(n):
+    """Select and return n random ports that are available."""
+    ports = []
+    sockets = []
+    for i in range(n):
+        sock = socket.socket()
+        sock.bind(('', 0))
+        ports.append(sock.getsockname()[1])
+        sockets.append(sock)
+    for sock in sockets:
+        sock.close()
+    return ports
+
+
+#-----------------------------------------------------------------------------
+# Check for passwordless login
+#-----------------------------------------------------------------------------
+_password_pat = re.compile(b(r'pass(word|phrase):'), re.IGNORECASE)
+
+
+def try_passwordless_ssh(server, keyfile, paramiko=None):
+    """Attempt to make an ssh connection without a password.
+    This is mainly used for requiring password input only once
+    when many tunnels may be connected to the same server.
+
+    If paramiko is None, the default for the platform is chosen.
+    """
+    if paramiko is None:
+        paramiko = sys.platform == 'win32'
+    if not paramiko:
+        f = _try_passwordless_openssh
+    else:
+        f = _try_passwordless_paramiko
+    return f(server, keyfile)
+
+
+def _try_passwordless_openssh(server, keyfile):
+    """Try passwordless login with shell ssh command."""
+    if pexpect is None:
+        raise ImportError("pexpect unavailable, use paramiko")
+    cmd = 'ssh -f ' + server
+    if keyfile:
+        cmd += ' -i ' + keyfile
+    cmd += ' exit'
+
+    # pop SSH_ASKPASS from env
+    env = os.environ.copy()
+    env.pop('SSH_ASKPASS', None)
+
+    ssh_newkey = 'Are you sure you want to continue connecting'
+    p = pexpect.spawn(cmd, env=env)
+    while True:
+        try:
+            i = p.expect([ssh_newkey, _password_pat], timeout=.1)
+            if i == 0:
+                raise SSHException('The authenticity of the host can\'t be established.')
+        except pexpect.TIMEOUT:
+            continue
+        except pexpect.EOF:
+            return True
+        else:
+            return False
+
+
+def _try_passwordless_paramiko(server, keyfile):
+    """Try passwordless login with paramiko."""
+    if paramiko is None:
+        msg = "Paramiko unavailable, "
+        if sys.platform == 'win32':
+            msg += "Paramiko is required for ssh tunneled connections on Windows."
+        else:
+            msg += "use OpenSSH."
+        raise ImportError(msg)
+    username, server, port = _split_server(server)
+    client = paramiko.SSHClient()
+    client.load_system_host_keys()
+    client.set_missing_host_key_policy(paramiko.WarningPolicy())
+    try:
+        client.connect(server, port, username=username, key_filename=keyfile,
+               look_for_keys=True)
+    except paramiko.AuthenticationException:
+        return False
+    else:
+        client.close()
+        return True
+
+
+def tunnel_connection(socket, addr, server, keyfile=None, password=None, paramiko=None, timeout=60):
+    """Connect a socket to an address via an ssh tunnel.
+
+    This is a wrapper for socket.connect(addr), when addr is not accessible
+    from the local machine.  It simply creates an ssh tunnel using the remaining args,
+    and calls socket.connect('tcp://localhost:lport') where lport is the randomly
+    selected local port of the tunnel.
+
+    """
+    new_url, tunnel = open_tunnel(addr, server, keyfile=keyfile, password=password, paramiko=paramiko, timeout=timeout)
+    socket.connect(new_url)
+    return tunnel
+
+
+def open_tunnel(addr, server, keyfile=None, password=None, paramiko=None, timeout=60):
+    """Open a tunneled connection from a 0MQ url.
+
+    For use inside tunnel_connection.
+
+    Returns
+    -------
+
+    (url, tunnel) : (str, object)
+        The 0MQ url that has been forwarded, and the tunnel object
+    """
+
+    lport = select_random_ports(1)[0]
+    transport, addr = addr.split('://')
+    ip, rport = addr.split(':')
+    rport = int(rport)
+    if paramiko is None:
+        paramiko = sys.platform == 'win32'
+    if paramiko:
+        tunnelf = paramiko_tunnel
+    else:
+        tunnelf = openssh_tunnel
+
+    tunnel = tunnelf(lport, rport, server, remoteip=ip, keyfile=keyfile, password=password, timeout=timeout)
+    return 'tcp://127.0.0.1:%i' % lport, tunnel
+
+
+def openssh_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, password=None, timeout=60):
+    """Create an ssh tunnel using command-line ssh that connects port lport
+    on this machine to localhost:rport on server.  The tunnel
+    will automatically close when not in use, remaining open
+    for a minimum of timeout seconds for an initial connection.
+
+    This creates a tunnel redirecting `localhost:lport` to `remoteip:rport`,
+    as seen from `server`.
+
+    keyfile and password may be specified, but ssh config is checked for defaults.
+
+    Parameters
+    ----------
+
+    lport : int
+        local port for connecting to the tunnel from this machine.
+    rport : int
+        port on the remote machine to connect to.
+    server : str
+        The ssh server to connect to. The full ssh server string will be parsed.
+        user@server:port
+    remoteip : str [Default: 127.0.0.1]
+        The remote ip, specifying the destination of the tunnel.
+        Default is localhost, which means that the tunnel would redirect
+        localhost:lport on this machine to localhost:rport on the *server*.
+
+    keyfile : str; path to public key file
+        This specifies a key to be used in ssh login, default None.
+        Regular default ssh keys will be used without specifying this argument.
+    password : str;
+        Your ssh password to the ssh server. Note that if this is left None,
+        you will be prompted for it if passwordless key based login is unavailable.
+    timeout : int [default: 60]
+        The time (in seconds) after which no activity will result in the tunnel
+        closing.  This prevents orphaned tunnels from running forever.
+    """
+    if pexpect is None:
+        raise ImportError("pexpect unavailable, use paramiko_tunnel")
+    ssh = "ssh "
+    if keyfile:
+        ssh += "-i " + keyfile
+
+    if ':' in server:
+        server, port = server.split(':')
+        ssh += " -p %s" % port
+
+    cmd = "%s -O check %s" % (ssh, server)
+    (output, exitstatus) = pexpect.run(cmd, withexitstatus=True)
+    if not exitstatus:
+        pid = int(output[output.find(b"(pid=")+5:output.find(b")")])
+        cmd = "%s -O forward -L 127.0.0.1:%i:%s:%i %s" % (
+            ssh, lport, remoteip, rport, server)
+        (output, exitstatus) = pexpect.run(cmd, withexitstatus=True)
+        if not exitstatus:
+            atexit.register(_stop_tunnel, cmd.replace("-O forward", "-O cancel", 1))
+            return pid
+    cmd = "%s -f -S none -L 127.0.0.1:%i:%s:%i %s sleep %i" % (
+        ssh, lport, remoteip, rport, server, timeout)
+
+    # pop SSH_ASKPASS from env
+    env = os.environ.copy()
+    env.pop('SSH_ASKPASS', None)
+
+    ssh_newkey = 'Are you sure you want to continue connecting'
+    tunnel = pexpect.spawn(cmd, env=env)
+    failed = False
+    while True:
+        try:
+            i = tunnel.expect([ssh_newkey, _password_pat], timeout=.1)
+            if i == 0:
+                raise SSHException('The authenticity of the host can\'t be established.')
+        except pexpect.TIMEOUT:
+            continue
+        except pexpect.EOF:
+            if tunnel.exitstatus:
+                print(tunnel.exitstatus)
+                print(tunnel.before)
+                print(tunnel.after)
+                raise RuntimeError("tunnel '%s' failed to start" % (cmd))
+            else:
+                return tunnel.pid
+        else:
+            if failed:
+                print("Password rejected, try again")
+                password = None
+            if password is None:
+                password = getpass("%s's password: " % (server))
+            tunnel.sendline(password)
+            failed = True
+
+
+def _stop_tunnel(cmd):
+    pexpect.run(cmd)
+
+
+def _split_server(server):
+    if '@' in server:
+        username, server = server.split('@', 1)
+    else:
+        username = getuser()
+    if ':' in server:
+        server, port = server.split(':')
+        port = int(port)
+    else:
+        port = 22
+    return username, server, port
+
+
+def paramiko_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, password=None, timeout=60):
+    """launch a tunner with paramiko in a subprocess. This should only be used
+    when shell ssh is unavailable (e.g. Windows).
+
+    This creates a tunnel redirecting `localhost:lport` to `remoteip:rport`,
+    as seen from `server`.
+
+    If you are familiar with ssh tunnels, this creates the tunnel:
+
+    ssh server -L localhost:lport:remoteip:rport
+
+    keyfile and password may be specified, but ssh config is checked for defaults.
+
+
+    Parameters
+    ----------
+
+    lport : int
+        local port for connecting to the tunnel from this machine.
+    rport : int
+        port on the remote machine to connect to.
+    server : str
+        The ssh server to connect to. The full ssh server string will be parsed.
+        user@server:port
+    remoteip : str [Default: 127.0.0.1]
+        The remote ip, specifying the destination of the tunnel.
+        Default is localhost, which means that the tunnel would redirect
+        localhost:lport on this machine to localhost:rport on the *server*.
+
+    keyfile : str; path to public key file
+        This specifies a key to be used in ssh login, default None.
+        Regular default ssh keys will be used without specifying this argument.
+    password : str;
+        Your ssh password to the ssh server. Note that if this is left None,
+        you will be prompted for it if passwordless key based login is unavailable.
+    timeout : int [default: 60]
+        The time (in seconds) after which no activity will result in the tunnel
+        closing.  This prevents orphaned tunnels from running forever.
+
+    """
+    if paramiko is None:
+        raise ImportError("Paramiko not available")
+
+    if password is None:
+        if not _try_passwordless_paramiko(server, keyfile):
+            password = getpass("%s's password: " % (server))
+
+    p = Process(target=_paramiko_tunnel,
+                args=(lport, rport, server, remoteip),
+                kwargs=dict(keyfile=keyfile, password=password))
+    p.daemon = True
+    p.start()
+    return p
+
+
+def _paramiko_tunnel(lport, rport, server, remoteip, keyfile=None, password=None):
+    """Function for actually starting a paramiko tunnel, to be passed
+    to multiprocessing.Process(target=this), and not called directly.
+    """
+    username, server, port = _split_server(server)
+    client = paramiko.SSHClient()
+    client.load_system_host_keys()
+    client.set_missing_host_key_policy(paramiko.WarningPolicy())
+
+    try:
+        client.connect(server, port, username=username, key_filename=keyfile,
+                       look_for_keys=True, password=password)
+#    except paramiko.AuthenticationException:
+#        if password is None:
+#            password = getpass("%s@%s's password: "%(username, server))
+#            client.connect(server, port, username=username, password=password)
+#        else:
+#            raise
+    except Exception as e:
+        print('*** Failed to connect to %s:%d: %r' % (server, port, e))
+        sys.exit(1)
+
+    # Don't let SIGINT kill the tunnel subprocess
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    try:
+        forward_tunnel(lport, remoteip, rport, client.get_transport())
+    except KeyboardInterrupt:
+        print('SIGINT: Port forwarding stopped cleanly')
+        sys.exit(0)
+    except Exception as e:
+        print("Port forwarding stopped uncleanly: %s" % e)
+        sys.exit(255)
+
+
+if sys.platform == 'win32':
+    ssh_tunnel = paramiko_tunnel
+else:
+    ssh_tunnel = openssh_tunnel
+
+
+__all__ = ['tunnel_connection', 'ssh_tunnel', 'openssh_tunnel', 'paramiko_tunnel', 'try_passwordless_ssh']


Test patch:
diff --git a/jupyter_client/tests/test_ssh.py b/jupyter_client/tests/test_ssh.py
new file mode 100644
index 0000000..e1673f9
--- /dev/null
+++ b/jupyter_client/tests/test_ssh.py
@@ -0,0 +1,8 @@
+from jupyter_client.ssh.tunnel import select_random_ports
+
+def test_random_ports():
+    for i in range(4096):
+        ports = select_random_ports(10)
+        assert len(ports) == 10
+        for p in ports:
+            assert ports.count(p) == 1