From 46eb4e128df94f741cf0ba5911736f34c76c9f5e 2011-10-13 17:18:01
From: MinRK <benjaminrk@gmail.com>
Date: 2011-10-13 17:18:01
Subject: [PATCH] split QtConsole.init_ssh into generic tunnel_to_kernel

added to new IPython.lib.kernel with other kernel connection utilities.

---

diff --git a/IPython/frontend/qt/console/qtconsoleapp.py b/IPython/frontend/qt/console/qtconsoleapp.py
index e4c4d91..0a4e93f 100644
--- a/IPython/frontend/qt/console/qtconsoleapp.py
+++ b/IPython/frontend/qt/console/qtconsoleapp.py
@@ -21,26 +21,22 @@ import glob
 import os
 import signal
 import sys
-from getpass import getpass
 
 # System library imports
 from IPython.external.qt import QtGui
 from pygments.styles import get_all_styles
 from zmq.utils import jsonapi as json
 
-# external imports
-from IPython.external.ssh import tunnel
-
 # Local imports
 from IPython.config.application import boolean_flag
 from IPython.core.application import BaseIPythonApplication
 from IPython.core.profiledir import ProfileDir
+from IPython.lib.kernel import tunnel_to_kernel
 from IPython.frontend.qt.console.frontend_widget import FrontendWidget
 from IPython.frontend.qt.console.ipython_widget import IPythonWidget
 from IPython.frontend.qt.console.rich_ipython_widget import RichIPythonWidget
 from IPython.frontend.qt.console import styles
 from IPython.frontend.qt.kernelmanager import QtKernelManager
-from IPython.parallel.util import select_random_ports
 from IPython.utils.path import filefind
 from IPython.utils.py3compat import str_to_bytes
 from IPython.utils.traitlets import (
@@ -424,24 +420,30 @@ class IPythonQtConsoleApp(BaseIPythonApplication):
             return
         
         if self.sshkey and not self.sshserver:
+            # specifying just the key implies that we are connecting directly
             self.sshserver = self.ip
-            self.ip=LOCALHOST
+            self.ip = LOCALHOST
         
-        lports = select_random_ports(4)
-        rports = self.shell_port, self.iopub_port, self.stdin_port, self.hb_port
-        self.shell_port, self.iopub_port, self.stdin_port, self.hb_port = lports
+        # build connection dict for tunnels:
+        info = dict(ip=self.ip,
+                    shell_port=self.shell_port,
+                    iopub_port=self.iopub_port,
+                    stdin_port=self.stdin_port,
+                    hb_port=self.hb_port
+        )
         
-        remote_ip = self.ip
-        self.ip = LOCALHOST
-        self.log.info("Forwarding connections to %s via %s"%(remote_ip, self.sshserver))
+        self.log.info("Forwarding connections to %s via %s"%(self.ip, self.sshserver))
         
-        if tunnel.try_passwordless_ssh(self.sshserver, self.sshkey):
-            password=False
-        else:
-            password = getpass("SSH Password for %s: "%self.sshserver)
+        # tunnels return a new set of ports, which will be on localhost:
+        self.ip = LOCALHOST
+        try:
+            newports = tunnel_to_kernel(info, self.sshserver, self.sshkey)
+        except:
+            # even catch KeyboardInterrupt
+            self.log.error("Could not setup tunnels", exc_info=True)
+            self.exit(1)
         
-        for lp,rp in zip(lports, rports):
-            tunnel.ssh_tunnel(lp, rp, self.sshserver, remote_ip, self.sshkey, password)
+        self.shell_port, self.iopub_port, self.stdin_port, self.hb_port = newports
         
         cf = self.connection_file
         base,ext = os.path.splitext(cf)
diff --git a/IPython/lib/kernel.py b/IPython/lib/kernel.py
index 1293344..4abfeba 100644
--- a/IPython/lib/kernel.py
+++ b/IPython/lib/kernel.py
@@ -19,8 +19,13 @@ Authors:
 
 import json
 import sys
+from getpass import getpass
 from subprocess import Popen, PIPE
 
+# external imports
+from IPython.external.ssh import tunnel
+
+# IPython imports
 from IPython.utils.path import filefind
 from IPython.utils.py3compat import str_to_bytes
 
@@ -96,4 +101,52 @@ def connect_qtconsole(argv=None):
     
     return Popen([sys.executable, '-c', cmd, '--existing', cf] + argv, stdout=PIPE, stderr=PIPE)
 
+def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
+    """tunnel connections to a kernel via ssh
+    
+    This will open four SSH tunnels from localhost on this machine to the
+    ports associated with the kernel.  They can be either direct
+    localhost-localhost tunnels, or if an intermediate server is necessary,
+    the kernel must be listening on a public IP.
+    
+    Parameters
+    ----------
+    connection_info : dict or str (path)
+        Either a connection dict, or the path to a JSON connection file
+    sshserver : str
+        The ssh sever to use to tunnel to the kernel. Can be a full
+        `user@server:port` string. ssh config aliases are respected.
+    sshkey : str [optional]
+        Path to file containing ssh key to use for authentication.
+        Only necessary if your ssh config does not already associate
+        a keyfile with the host.
+    
+    Returns
+    -------
+    
+    (shell, iopub, stdin, hb) : ints
+        The four ports on localhost that have been forwarded to the kernel.
+    """
+    if isinstance(connection_info, basestring):
+        # it's a path, unpack it
+        with open(connection_info) as f:
+            connection_info = json.loads(f.read())
+    
+    cf = connection_info
+    
+    lports = tunnel.select_random_ports(4)
+    rports = cf['shell_port'], cf['iopub_port'], cf['stdin_port'], cf['hb_port']
+    
+    remote_ip = cf['ip']
+    
+    if tunnel.try_passwordless_ssh(sshserver, sshkey):
+        password=False
+    else:
+        password = getpass("SSH Password for %s: "%sshserver)
+    
+    for lp,rp in zip(lports, rports):
+        tunnel.ssh_tunnel(lp, rp, sshserver, remote_ip, sshkey, password)
+    
+    return tuple(lports)
+