added preliminary ssh tunneling support for clients

IPython/zmq/forward.py

@@ -1,183 +1,183
 #!/usr/bin/env python
 
 #
-# This file is adapted from a paramiko demo, and thus LGPL 2.1.
+# This file is adapted from a paramiko demo, and thus licensed under LGPL 2.1.
 # Original Copyright (C) 2003-2007  Robey Pointer <[e-mail]>
 # 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.,
 # 59 Temple Place, Suite 330, Boston, MA  02111-1307  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 getpass
 import os
 import socket
 import select
 import SocketServer
 import sys
 from optparse import OptionParser
 
 import paramiko
 
 SSH_PORT = 22
 DEFAULT_PORT = 4000
 
 g_verbose = False
 
 
 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, e:
             verbose('Incoming request to %s:%d failed: %s' % (self.chain_host,
                                                               self.chain_port,
                                                               repr(e)))
             return
         if chan is None:
             verbose('Incoming request to %s:%d was rejected by the SSH server.' %
                     (self.chain_host, self.chain_port))
             return
 
         verbose('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()
-        verbose('Tunnel closed from %r' % (self.request.getpeername(),))
+        verbose('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(('', local_port), SubHander).serve_forever()
+    ForwardServer(('127.0.0.1', local_port), SubHander).serve_forever()
 
 
 def verbose(s):
     if g_verbose:
         print (s)
 
 
 HELP = """\
 Set up a forward tunnel across an SSH server, using paramiko. A local port
 (given with -p) is forwarded across an SSH session to an address:port from
 the SSH server. This is similar to the openssh -L option.
 """
 
 
 def get_host_port(spec, default_port):
     "parse 'hostname:22' into a host and port, with the port optional"
     args = (spec.split(':', 1) + [default_port])[:2]
     args[1] = int(args[1])
     return args[0], args[1]
 
 
 def parse_options():
     global g_verbose
     
     parser = OptionParser(usage='usage: %prog [options] <ssh-server>[:<server-port>]',
                           version='%prog 1.0', description=HELP)
     parser.add_option('-q', '--quiet', action='store_false', dest='verbose', default=True,
                       help='squelch all informational output')
     parser.add_option('-p', '--local-port', action='store', type='int', dest='port',
                       default=DEFAULT_PORT,
                       help='local port to forward (default: %d)' % DEFAULT_PORT)
     parser.add_option('-u', '--user', action='store', type='string', dest='user',
                       default=getpass.getuser(),
                       help='username for SSH authentication (default: %s)' % getpass.getuser())
     parser.add_option('-K', '--key', action='store', type='string', dest='keyfile',
                       default=None,
                       help='private key file to use for SSH authentication')
     parser.add_option('', '--no-key', action='store_false', dest='look_for_keys', default=True,
                       help='don\'t look for or use a private key file')
     parser.add_option('-P', '--password', action='store_true', dest='readpass', default=False,
                       help='read password (for key or password auth) from stdin')
     parser.add_option('-r', '--remote', action='store', type='string', dest='remote', default=None, metavar='host:port',
                       help='remote host and port to forward to')
     options, args = parser.parse_args()
 
     if len(args) != 1:
         parser.error('Incorrect number of arguments.')
     if options.remote is None:
         parser.error('Remote address required (-r).')
     
     g_verbose = options.verbose
     server_host, server_port = get_host_port(args[0], SSH_PORT)
     remote_host, remote_port = get_host_port(options.remote, SSH_PORT)
     return options, (server_host, server_port), (remote_host, remote_port)
 
 
 def main():
     options, server, remote = parse_options()
     
     password = None
     if options.readpass:
         password = getpass.getpass('Enter SSH password: ')
     
     client = paramiko.SSHClient()
     client.load_system_host_keys()
     client.set_missing_host_key_policy(paramiko.WarningPolicy())
 
     verbose('Connecting to ssh host %s:%d ...' % (server[0], server[1]))
     try:
         client.connect(server[0], server[1], username=options.user, key_filename=options.keyfile,
                        look_for_keys=options.look_for_keys, password=password)
     except Exception as e:
         print ('*** Failed to connect to %s:%d: %r' % (server[0], server[1], e))
         sys.exit(1)
 
     verbose('Now forwarding port %d to %s:%d ...' % (options.port, remote[0], remote[1]))
 
     try:
         forward_tunnel(options.port, remote[0], remote[1], client.get_transport())
     except KeyboardInterrupt:
         print ('C-c: Port forwarding stopped.')
         sys.exit(0)
 
 
 if __name__ == '__main__':
     main()

IPython/zmq/parallel/client.py

@@ -1,855 +1,905
 """A semi-synchronous Client for the ZMQ controller"""
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 from __future__ import print_function
 
 import time
 from pprint import pprint
 
 import zmq
 from zmq.eventloop import ioloop, zmqstream
 
 from IPython.external.decorator import decorator
+from IPython.zmq import tunnel
 
 import streamsession as ss
 # from remotenamespace import RemoteNamespace
 from view import DirectView, LoadBalancedView
 from dependency import Dependency, depend, require
 
 def _push(ns):
     globals().update(ns)
 
 def _pull(keys):
     g = globals()
     if isinstance(keys, (list,tuple, set)):
         for key in keys:
             if not g.has_key(key):
                 raise NameError("name '%s' is not defined"%key)
         return map(g.get, keys)
     else:
         if not g.has_key(keys):
             raise NameError("name '%s' is not defined"%keys)
         return g.get(keys)
 
 def _clear():
     globals().clear()
 
 def execute(code):
     exec code in globals()
 
 #--------------------------------------------------------------------------
 # Decorators for Client methods
 #--------------------------------------------------------------------------
 
 @decorator
 def spinfirst(f, self, *args, **kwargs):
     """Call spin() to sync state prior to calling the method."""
     self.spin()
     return f(self, *args, **kwargs)
 
 @decorator
 def defaultblock(f, self, *args, **kwargs):
     """Default to self.block; preserve self.block."""
     block = kwargs.get('block',None)
     block = self.block if block is None else block
     saveblock = self.block
     self.block = block
     ret = f(self, *args, **kwargs)
     self.block = saveblock
     return ret
 
 def remote(client, bound=False, block=None, targets=None):
     """Turn a function into a remote function.
     
     This method can be used for map:
     
     >>> @remote(client,block=True)
         def func(a)
     """
     def remote_function(f):
         return RemoteFunction(client, f, bound, block, targets)
     return remote_function
 
 #--------------------------------------------------------------------------
 # Classes
 #--------------------------------------------------------------------------
 
 class RemoteFunction(object):
     """Turn an existing function into a remote function"""
     
     def __init__(self, client, f, bound=False, block=None, targets=None):
         self.client = client
         self.func = f
         self.block=block
         self.bound=bound
         self.targets=targets
     
     def __call__(self, *args, **kwargs):
         return self.client.apply(self.func, args=args, kwargs=kwargs,
                 block=self.block, targets=self.targets, bound=self.bound)
     
 
 class AbortedTask(object):
     """A basic wrapper object describing an aborted task."""
     def __init__(self, msg_id):
         self.msg_id = msg_id
 
 class ControllerError(Exception):
     def __init__(self, etype, evalue, tb):
         self.etype = etype
         self.evalue = evalue
         self.traceback=tb
     
 class Client(object):
     """A semi-synchronous client to the IPython ZMQ controller
     
     Parameters
     ----------
     
     addr : bytes; zmq url, e.g. 'tcp://127.0.0.1:10101'
         The address of the controller's registration socket.
-    
+        [Default: 'tcp://127.0.0.1:10101']
+    context : zmq.Context
+        Pass an existing zmq.Context instance, otherwise the client will create its own
+    username : bytes
+        set username to be passed to the Session object
+    debug : bool
+        flag for lots of message printing for debug purposes
+ 
+    #-------------- ssh related args ----------------
+    # These are args for configuring the ssh tunnel to be used
+    # credentials are used to forward connections over ssh to the Controller
+    # Note that the ip given in `addr` needs to be relative to sshserver
+    # The most basic case is to leave addr as pointing to localhost (127.0.0.1),
+    # and set sshserver as the same machine the Controller is on. However, 
+    # the only requirement is that sshserver is able to see the Controller
+    # (i.e. is within the same trusted network).
+    
+    sshserver : str
+        A string of the form passed to ssh, i.e. 'server.tld' or '[e-mail]:port'
+        If keyfile or password is specified, and this is not, it will default to
+        the ip given in addr.
+    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 sshserver. Note that if this is left None,
+        you will be prompted for it if passwordless key based login is unavailable.
     
     Attributes
     ----------
     ids : set of int engine IDs
         requesting the ids attribute always synchronizes
         the registration state. To request ids without synchronization,
         use semi-private _ids.
     
     history : list of msg_ids
         a list of msg_ids, keeping track of all the execution
         messages you have submitted in order.
     
     outstanding : set of msg_ids
         a set of msg_ids that have been submitted, but whose
         results have not yet been received.
     
     results : dict
         a dict of all our results, keyed by msg_id
     
     block : bool
         determines default behavior when block not specified
         in execution methods
     
     Methods
     -------
     spin : flushes incoming results and registration state changes
             control methods spin, and requesting `ids` also ensures up to date
             
     barrier : wait on one or more msg_ids
     
     execution methods: apply/apply_bound/apply_to/applu_bount
         legacy: execute, run
     
     query methods: queue_status, get_result, purge
     
     control methods: abort, kill
     
     """
     
     
     _connected=False
+    _ssh=False
     _engines=None
     _addr='tcp://127.0.0.1:10101'
     _registration_socket=None
     _query_socket=None
     _control_socket=None
     _notification_socket=None
     _mux_socket=None
     _task_socket=None
     block = False
     outstanding=None
     results = None
     history = None
     debug = False
     
-    def __init__(self, addr='tcp://127.0.0.1:10101', context=None, username=None, debug=False):
+    def __init__(self, addr='tcp://127.0.0.1:10101', context=None, username=None, debug=False, 
+            sshserver=None, keyfile=None, password=None, paramiko=None):
         if context is None:
             context = zmq.Context()
         self.context = context
         self._addr = addr
+        self._ssh = bool(sshserver or keyfile or password)
+        if self._ssh and sshserver is None:
+            # default to the same
+            sshserver = addr.split('://')[1].split(':')[0]
+        if self._ssh and password is None:
+            if tunnel.try_passwordless_ssh(sshserver, keyfile, paramiko):
+                password=False
+            else:
+                password = getpass("SSH Password for %s: "%sshserver)
+        ssh_kwargs = dict(keyfile=keyfile, password=password, paramiko=paramiko)
+
         if username is None:
             self.session = ss.StreamSession()
         else:
             self.session = ss.StreamSession(username)
-        self._registration_socket = self.context.socket(zmq.PAIR)
+        self._registration_socket = self.context.socket(zmq.XREQ)
         self._registration_socket.setsockopt(zmq.IDENTITY, self.session.session)
+        if self._ssh:
+            tunnel.tunnel_connection(self._registration_socket, addr, sshserver, **ssh_kwargs)
+        else:
         self._registration_socket.connect(addr)
         self._engines = {}
         self._ids = set()
         self.outstanding=set()
         self.results = {}
         self.history = []
         self.debug = debug
         self.session.debug = debug
         
         self._notification_handlers = {'registration_notification' : self._register_engine,
                                     'unregistration_notification' : self._unregister_engine,
                                     }
         self._queue_handlers = {'execute_reply' : self._handle_execute_reply,
                                 'apply_reply' : self._handle_apply_reply}
-        self._connect()
+        self._connect(sshserver, ssh_kwargs)
         
     
     @property
     def ids(self):
         """Always up to date ids property."""
         self._flush_notifications()
         return self._ids
     
     def _update_engines(self, engines):
         """Update our engines dict and _ids from a dict of the form: {id:uuid}."""
         for k,v in engines.iteritems():
             eid = int(k)
             self._engines[eid] = bytes(v) # force not unicode
             self._ids.add(eid)
     
     def _build_targets(self, targets):
         """Turn valid target IDs or 'all' into two lists:
         (int_ids, uuids).
         """
         if targets is None:
             targets = self._ids
         elif isinstance(targets, str):
             if targets.lower() == 'all':
                 targets = self._ids
             else:
                 raise TypeError("%r not valid str target, must be 'all'"%(targets))
         elif isinstance(targets, int):
             targets = [targets]
         return [self._engines[t] for t in targets], list(targets)
     
-    def _connect(self):
+    def _connect(self, sshserver, ssh_kwargs):
         """setup all our socket connections to the controller. This is called from
         __init__."""
         if self._connected:
             return
         self._connected=True
+
+        def connect_socket(s, addr):
+            if self._ssh:
+                return tunnel.tunnel_connection(s, addr, sshserver, **ssh_kwargs)
+            else:
+                return s.connect(addr)
+            
         self.session.send(self._registration_socket, 'connection_request')
         idents,msg = self.session.recv(self._registration_socket,mode=0)
         if self.debug:
             pprint(msg)
         msg = ss.Message(msg)
         content = msg.content
         if content.status == 'ok':
             if content.queue:
                 self._mux_socket = self.context.socket(zmq.PAIR)
                 self._mux_socket.setsockopt(zmq.IDENTITY, self.session.session)
-                self._mux_socket.connect(content.queue)
+                connect_socket(self._mux_socket, content.queue)
             if content.task:
                 self._task_socket = self.context.socket(zmq.PAIR)
                 self._task_socket.setsockopt(zmq.IDENTITY, self.session.session)
-                self._task_socket.connect(content.task)
+                connect_socket(self._task_socket, content.task)
             if content.notification:
                 self._notification_socket = self.context.socket(zmq.SUB)
-                self._notification_socket.connect(content.notification)
+                connect_socket(self._notification_socket, content.notification)
                 self._notification_socket.setsockopt(zmq.SUBSCRIBE, "")
             if content.query:
                 self._query_socket = self.context.socket(zmq.PAIR)
                 self._query_socket.setsockopt(zmq.IDENTITY, self.session.session)
-                self._query_socket.connect(content.query)
+                connect_socket(self._query_socket, content.query)
             if content.control:
                 self._control_socket = self.context.socket(zmq.PAIR)
                 self._control_socket.setsockopt(zmq.IDENTITY, self.session.session)
-                self._control_socket.connect(content.control)
+                connect_socket(self._control_socket, content.control)
             self._update_engines(dict(content.engines))
                 
         else:
             self._connected = False
             raise Exception("Failed to connect!")
     
     #--------------------------------------------------------------------------
     # handlers and callbacks for incoming messages
     #--------------------------------------------------------------------------
     
     def _register_engine(self, msg):
         """Register a new engine, and update our connection info."""
         content = msg['content']
         eid = content['id']
         d = {eid : content['queue']}
         self._update_engines(d)
         self._ids.add(int(eid))
 
     def _unregister_engine(self, msg):
         """Unregister an engine that has died."""
         content = msg['content']
         eid = int(content['id'])
         if eid in self._ids:
             self._ids.remove(eid)
             self._engines.pop(eid)
         
     def _handle_execute_reply(self, msg):
         """Save the reply to an execute_request into our results."""
         parent = msg['parent_header']
         msg_id = parent['msg_id']
         if msg_id not in self.outstanding:
             print("got unknown result: %s"%msg_id)
         else:
             self.outstanding.remove(msg_id)
         self.results[msg_id] = ss.unwrap_exception(msg['content'])
     
     def _handle_apply_reply(self, msg):
         """Save the reply to an apply_request into our results."""
         parent = msg['parent_header']
         msg_id = parent['msg_id']
         if msg_id not in self.outstanding:
             print ("got unknown result: %s"%msg_id)
         else:
             self.outstanding.remove(msg_id)
         content = msg['content']
         if content['status'] == 'ok':
             self.results[msg_id] = ss.unserialize_object(msg['buffers'])
         elif content['status'] == 'aborted':
             self.results[msg_id] = AbortedTask(msg_id)
         elif content['status'] == 'resubmitted':
             # TODO: handle resubmission
             pass
         else:
             self.results[msg_id] = ss.unwrap_exception(content)
     
     def _flush_notifications(self):
         """Flush notifications of engine registrations waiting
         in ZMQ queue."""
         msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
         while msg is not None:
             if self.debug:
                 pprint(msg)
             msg = msg[-1]
             msg_type = msg['msg_type']
             handler = self._notification_handlers.get(msg_type, None)
             if handler is None:
                 raise Exception("Unhandled message type: %s"%msg.msg_type)
             else:
                 handler(msg)
             msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
     
     def _flush_results(self, sock):
         """Flush task or queue results waiting in ZMQ queue."""
         msg = self.session.recv(sock, mode=zmq.NOBLOCK)
         while msg is not None:
             if self.debug:
                 pprint(msg)
             msg = msg[-1]
             msg_type = msg['msg_type']
             handler = self._queue_handlers.get(msg_type, None)
             if handler is None:
                 raise Exception("Unhandled message type: %s"%msg.msg_type)
             else:
                 handler(msg)
             msg = self.session.recv(sock, mode=zmq.NOBLOCK)
     
     def _flush_control(self, sock):
         """Flush replies from the control channel waiting
         in the ZMQ queue.
         
         Currently: ignore them."""
         msg = self.session.recv(sock, mode=zmq.NOBLOCK)
         while msg is not None:
             if self.debug:
                 pprint(msg)
             msg = self.session.recv(sock, mode=zmq.NOBLOCK)
     
     #--------------------------------------------------------------------------
     # getitem
     #--------------------------------------------------------------------------
     
     def __getitem__(self, key):
         """Dict access returns DirectView multiplexer objects or,
         if key is None, a LoadBalancedView."""
         if key is None:
             return LoadBalancedView(self)
         if isinstance(key, int):
             if key not in self.ids:
                 raise IndexError("No such engine: %i"%key)
             return DirectView(self, key)
         
         if isinstance(key, slice):
             indices = range(len(self.ids))[key]
             ids = sorted(self._ids)
             key = [ ids[i] for i in indices ]
             # newkeys = sorted(self._ids)[thekeys[k]]
         
         if isinstance(key, (tuple, list, xrange)):
             _,targets = self._build_targets(list(key))
             return DirectView(self, targets)
         else:
             raise TypeError("key by int/iterable of ints only, not %s"%(type(key)))
     
     #--------------------------------------------------------------------------
     # Begin public methods
     #--------------------------------------------------------------------------
     
     def spin(self):
         """Flush any registration notifications and execution results
         waiting in the ZMQ queue.
         """
         if self._notification_socket:
             self._flush_notifications()
         if self._mux_socket:
             self._flush_results(self._mux_socket)
         if self._task_socket:
             self._flush_results(self._task_socket)
         if self._control_socket:
             self._flush_control(self._control_socket)
     
     def barrier(self, msg_ids=None, timeout=-1):
         """waits on one or more `msg_ids`, for up to `timeout` seconds.
         
         Parameters
         ----------
         msg_ids : int, str, or list of ints and/or strs
                 ints are indices to self.history
                 strs are msg_ids
                 default: wait on all outstanding messages
         timeout : float
                 a time in seconds, after which to give up.
                 default is -1, which means no timeout
         
         Returns
         -------
         True : when all msg_ids are done
         False : timeout reached, some msg_ids still outstanding
         """
         tic = time.time()
         if msg_ids is None:
             theids = self.outstanding
         else:
             if isinstance(msg_ids, (int, str)):
                 msg_ids = [msg_ids]
             theids = set()
             for msg_id in msg_ids:
                 if isinstance(msg_id, int):
                     msg_id = self.history[msg_id]
                 theids.add(msg_id)
         self.spin()
         while theids.intersection(self.outstanding):
             if timeout >= 0 and ( time.time()-tic ) > timeout:
                 break
             time.sleep(1e-3)
             self.spin()
         return len(theids.intersection(self.outstanding)) == 0
     
     #--------------------------------------------------------------------------
     # Control methods
     #--------------------------------------------------------------------------
     
     @spinfirst
     @defaultblock
     def clear(self, targets=None, block=None):
         """Clear the namespace in target(s)."""
         targets = self._build_targets(targets)[0]
         for t in targets:
             self.session.send(self._control_socket, 'clear_request', content={}, ident=t)
         error = False
         if self.block:
             for i in range(len(targets)):
                 idents,msg = self.session.recv(self._control_socket,0)
                 if self.debug:
                     pprint(msg)
                 if msg['content']['status'] != 'ok':
                     error = ss.unwrap_exception(msg['content'])
         if error:
             return error
         
     
     @spinfirst
     @defaultblock
     def abort(self, msg_ids = None, targets=None, block=None):
         """Abort the execution queues of target(s)."""
         targets = self._build_targets(targets)[0]
         if isinstance(msg_ids, basestring):
             msg_ids = [msg_ids]
         content = dict(msg_ids=msg_ids)
         for t in targets:
             self.session.send(self._control_socket, 'abort_request', 
                     content=content, ident=t)
         error = False
         if self.block:
             for i in range(len(targets)):
                 idents,msg = self.session.recv(self._control_socket,0)
                 if self.debug:
                     pprint(msg)
                 if msg['content']['status'] != 'ok':
                     error = ss.unwrap_exception(msg['content'])
         if error:
             return error
     
     @spinfirst
     @defaultblock
     def kill(self, targets=None, block=None):
         """Terminates one or more engine processes."""
         targets = self._build_targets(targets)[0]
         for t in targets:
             self.session.send(self._control_socket, 'kill_request', content={},ident=t)
         error = False
         if self.block:
             for i in range(len(targets)):
                 idents,msg = self.session.recv(self._control_socket,0)
                 if self.debug:
                     pprint(msg)
                 if msg['content']['status'] != 'ok':
                     error = ss.unwrap_exception(msg['content'])
         if error:
             return error
     
     #--------------------------------------------------------------------------
     # Execution methods
     #--------------------------------------------------------------------------
     
     @defaultblock
     def execute(self, code, targets='all', block=None):
         """Executes `code` on `targets` in blocking or nonblocking manner.
         
         Parameters
         ----------
         code : str
                 the code string to be executed
         targets : int/str/list of ints/strs
                 the engines on which to execute
                 default : all
         block : bool
                 whether or not to wait until done to return
                 default: self.block
         """
         # block = self.block if block is None else block
         # saveblock = self.block
         # self.block = block
         result = self.apply(execute, (code,), targets=targets, block=block, bound=True)
         # self.block = saveblock
         return result
     
     def run(self, code, block=None):
         """Runs `code` on an engine. 
         
         Calls to this are load-balanced.
         
         Parameters
         ----------
         code : str
                 the code string to be executed
         block : bool
                 whether or not to wait until done
         
         """
         result = self.apply(execute, (code,), targets=None, block=block, bound=False)
         return result
     
     def apply(self, f, args=None, kwargs=None, bound=True, block=None, targets=None,
                         after=None, follow=None):
         """Call `f(*args, **kwargs)` on a remote engine(s), returning the result.
         
         This is the central execution command for the client.
         
         Parameters
         ----------
         
         f : function
             The fuction to be called remotely
         args : tuple/list
             The positional arguments passed to `f`
         kwargs : dict
             The keyword arguments passed to `f`
         bound : bool (default: True)
             Whether to execute in the Engine(s) namespace, or in a clean
             namespace not affecting the engine.
         block : bool (default: self.block)
             Whether to wait for the result, or return immediately.
             False:
                 returns msg_id(s)
                 if multiple targets:
                     list of ids
             True:
                 returns actual result(s) of f(*args, **kwargs)
                 if multiple targets:
                     dict of results, by engine ID
         targets : int,list of ints, 'all', None
             Specify the destination of the job.
             if None:
                 Submit via Task queue for load-balancing.
             if 'all':
                 Run on all active engines
             if list:
                 Run on each specified engine
             if int:
                 Run on single engine
             
         after : Dependency or collection of msg_ids
             Only for load-balanced execution (targets=None)
             Specify a list of msg_ids as a time-based dependency.
             This job will only be run *after* the dependencies
             have been met.
             
         follow : Dependency or collection of msg_ids
             Only for load-balanced execution (targets=None)
             Specify a list of msg_ids as a location-based dependency.
             This job will only be run on an engine where this dependency
             is met.
         
         Returns
         -------
         if block is False:
             if single target:
                 return msg_id
             else:
                 return list of msg_ids
                 ? (should this be dict like block=True) ?
         else:
             if single target:
                 return result of f(*args, **kwargs)
             else:
                 return dict of results, keyed by engine
         """
         
         # defaults:
         block = block if block is not None else self.block
         args = args if args is not None else []
         kwargs = kwargs if kwargs is not None else {}
         
         # enforce types of f,args,kwrags
         if not callable(f):
             raise TypeError("f must be callable, not %s"%type(f))
         if not isinstance(args, (tuple, list)):
             raise TypeError("args must be tuple or list, not %s"%type(args))
         if not isinstance(kwargs, dict):
             raise TypeError("kwargs must be dict, not %s"%type(kwargs))
             
         options  = dict(bound=bound, block=block, after=after, follow=follow)
             
         if targets is None:
             return self._apply_balanced(f, args, kwargs, **options)
         else:
             return self._apply_direct(f, args, kwargs, targets=targets, **options)
     
     def _apply_balanced(self, f, args, kwargs, bound=True, block=None,
                             after=None, follow=None):
         """The underlying method for applying functions in a load balanced
         manner, via the task queue."""
         if isinstance(after, Dependency):
             after = after.as_dict()
         elif after is None:
             after = []
         if isinstance(follow, Dependency):
             follow = follow.as_dict()
         elif follow is None:
             follow = []
         subheader = dict(after=after, follow=follow)
         
         bufs = ss.pack_apply_message(f,args,kwargs)
         content = dict(bound=bound)
         msg = self.session.send(self._task_socket, "apply_request", 
                 content=content, buffers=bufs, subheader=subheader)
         msg_id = msg['msg_id']
         self.outstanding.add(msg_id)
         self.history.append(msg_id)
         if block:
             self.barrier(msg_id)
             return self.results[msg_id]
         else:
             return msg_id
     
     def _apply_direct(self, f, args, kwargs, bound=True, block=None, targets=None,
                                 after=None, follow=None):
         """Then underlying method for applying functions to specific engines
         via the MUX queue."""
         
         queues,targets = self._build_targets(targets)
         bufs = ss.pack_apply_message(f,args,kwargs)
         if isinstance(after, Dependency):
             after = after.as_dict()
         elif after is None:
             after = []
         if isinstance(follow, Dependency):
             follow = follow.as_dict()
         elif follow is None:
             follow = []
         subheader = dict(after=after, follow=follow)
         content = dict(bound=bound)
         msg_ids = []
         for queue in queues:
             msg = self.session.send(self._mux_socket, "apply_request", 
                     content=content, buffers=bufs,ident=queue, subheader=subheader)
             msg_id = msg['msg_id']
             self.outstanding.add(msg_id)
             self.history.append(msg_id)
             msg_ids.append(msg_id)
         if block:
             self.barrier(msg_ids)
         else:
             if len(msg_ids) == 1:
                 return msg_ids[0]
             else:
                 return msg_ids
         if len(msg_ids) == 1:
             return self.results[msg_ids[0]]
         else:
             result = {}
             for target,mid in zip(targets, msg_ids):
                     result[target] = self.results[mid]
             return result
     
     #--------------------------------------------------------------------------
     # Data movement
     #--------------------------------------------------------------------------
     
     @defaultblock
     def push(self, ns, targets=None, block=None):
         """Push the contents of `ns` into the namespace on `target`"""
         if not isinstance(ns, dict):
             raise TypeError("Must be a dict, not %s"%type(ns))
         result = self.apply(_push, (ns,), targets=targets, block=block, bound=True)
         return result
     
     @defaultblock
     def pull(self, keys, targets=None, block=True):
         """Pull objects from `target`'s namespace by `keys`"""
         if isinstance(keys, str):
             pass
         elif isistance(keys, (list,tuple,set)):
             for key in keys:
                 if not isinstance(key, str):
                     raise TypeError
         result = self.apply(_pull, (keys,), targets=targets, block=block, bound=True)
         return result
     
     #--------------------------------------------------------------------------
     # Query methods
     #--------------------------------------------------------------------------
     
     @spinfirst
     def get_results(self, msg_ids, status_only=False):
         """Returns the result of the execute or task request with `msg_ids`.
         
         Parameters
         ----------
         msg_ids : list of ints or msg_ids
             if int:
                 Passed as index to self.history for convenience.
         status_only : bool (default: False)
             if False:
                 return the actual results
         """
         if not isinstance(msg_ids, (list,tuple)):
             msg_ids = [msg_ids]
         theids = []
         for msg_id in msg_ids:
             if isinstance(msg_id, int):
                 msg_id = self.history[msg_id]
             if not isinstance(msg_id, str):
                 raise TypeError("msg_ids must be str, not %r"%msg_id)
             theids.append(msg_id)
         
         completed = []
         local_results = {}
         for msg_id in list(theids):
             if msg_id in self.results:
                 completed.append(msg_id)
                 local_results[msg_id] = self.results[msg_id]
                 theids.remove(msg_id)
         
         if theids: # some not locally cached
             content = dict(msg_ids=theids, status_only=status_only)
             msg = self.session.send(self._query_socket, "result_request", content=content)
             zmq.select([self._query_socket], [], [])
             idents,msg = self.session.recv(self._query_socket, zmq.NOBLOCK)
             if self.debug:
                 pprint(msg)
             content = msg['content']
             if content['status'] != 'ok':
                 raise ss.unwrap_exception(content)
         else:
             content = dict(completed=[],pending=[])
         if not status_only:
             # load cached results into result:
             content['completed'].extend(completed)
             content.update(local_results)
             # update cache with results:
             for msg_id in msg_ids:
                 if msg_id in content['completed']:
                     self.results[msg_id] = content[msg_id]
         return content
 
     @spinfirst
     def queue_status(self, targets=None, verbose=False):
         """Fetch the status of engine queues.
         
         Parameters
         ----------
         targets : int/str/list of ints/strs
                 the engines on which to execute
                 default : all
         verbose : bool
                 whether to return lengths only, or lists of ids for each element
         """
         targets = self._build_targets(targets)[1]
         content = dict(targets=targets, verbose=verbose)
         self.session.send(self._query_socket, "queue_request", content=content)
         idents,msg = self.session.recv(self._query_socket, 0)
         if self.debug:
             pprint(msg)
         content = msg['content']
         status = content.pop('status')
         if status != 'ok':
             raise ss.unwrap_exception(content)
         return content
         
     @spinfirst
     def purge_results(self, msg_ids=[], targets=[]):
         """Tell the controller to forget results.
         
         Individual results can be purged by msg_id, or the entire
         history of specific targets can 
         
         Parameters
         ----------
         targets : int/str/list of ints/strs
                 the targets
                 default : None
         """
         if not targets and not msg_ids:
             raise ValueError
         if targets:
             targets = self._build_targets(targets)[1]
         content = dict(targets=targets, msg_ids=msg_ids)
         self.session.send(self._query_socket, "purge_request", content=content)
         idents, msg = self.session.recv(self._query_socket, 0)
         if self.debug:
             pprint(msg)
         content = msg['content']
         if content['status'] != 'ok':
             raise ss.unwrap_exception(content)
 
 class AsynClient(Client):
     """An Asynchronous client, using the Tornado Event Loop.
     !!!unfinished!!!"""
     io_loop = None
     _queue_stream = None
     _notifier_stream = None
     _task_stream = None
     _control_stream = None
     
     def __init__(self, addr, context=None, username=None, debug=False, io_loop=None):
         Client.__init__(self, addr, context, username, debug)
         if io_loop is None:
             io_loop = ioloop.IOLoop.instance()
         self.io_loop = io_loop
         
         self._queue_stream = zmqstream.ZMQStream(self._mux_socket, io_loop)
         self._control_stream = zmqstream.ZMQStream(self._control_socket, io_loop)
         self._task_stream = zmqstream.ZMQStream(self._task_socket, io_loop)
         self._notification_stream = zmqstream.ZMQStream(self._notification_socket, io_loop)
     
     def spin(self):
         for stream in (self.queue_stream, self.notifier_stream, 
                         self.task_stream, self.control_stream):
             stream.flush()
     

IPython/zmq/tunnel.py

@@ -1,123 +1,300
+"""Basic ssh tunneling utilities."""
 
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2008-2010  The IPython Development Team
+#
+#  Distributed under the terms of the BSD License.  The full license is in
+#  the file COPYING, distributed as part of this software.
+#-----------------------------------------------------------------------------
 
-#-----------------------------------------
+
+
+#-----------------------------------------------------------------------------
 # Imports
-#-----------------------------------------
+#-----------------------------------------------------------------------------
 
 from __future__ import print_function
 
-import os,sys
+import os,sys, atexit
 from multiprocessing import Process
 from getpass import getpass, getuser
 
 try:
     import paramiko
 except ImportError:
     paramiko = None
 else:
     from forward import forward_tunnel
     
+try:
 from IPython.external import pexpect
+except ImportError:
+    pexpect = None
+
+from IPython.zmq.parallel.entry_point import select_random_ports
+
+#-----------------------------------------------------------------------------
+# Code
+#-----------------------------------------------------------------------------
 
+#-----------------------------------------------------------------------------
+# Check for passwordless login
+#-----------------------------------------------------------------------------
+
+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'
+    p = pexpect.spawn(cmd)
+    while True:
+        try:
+            p.expect('[Ppassword]:', timeout=.1)
+        except pexpect.TIMEOUT:
+            continue
+        except pexpect.EOF:
+            return True
+        else:
+            return False
 
-def launch_ssh_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, timeout=15):
+def _try_passwordless_paramiko(server, keyfile):
+    """Try passwordless login with paramiko."""
+    if paramiko is None:
+        raise ImportError("paramiko unavailable, use openssh")
+    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):
+    """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.
+    
+    """
+    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)
+    socket.connect('tcp://127.0.0.1:%i'%lport)
+    return tunnel
+
+def openssh_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, password=None, timeout=15):
     """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.
+    
     """
+    if pexpect is None:
+        raise ImportError("pexpect unavailable, use paramiko_tunnel")
     ssh="ssh "
     if keyfile:
         ssh += "-i " + keyfile 
-    cmd = ssh + " -f -L %i:127.0.0.1:%i %s sleep %i"%(lport, rport, server, timeout)
+    cmd = ssh + " -f -L 127.0.0.1:%i:127.0.0.1:%i %s sleep %i"%(lport, rport, server, timeout)
     tunnel = pexpect.spawn(cmd)
     failed = False
     while True:
         try:
             tunnel.expect('[Pp]assword:', timeout=.1)
         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")
-            tunnel.sendline(getpass())
+                password=None
+            if password is None:
+                password = getpass("%s's password: "%(server))
+            tunnel.sendline(password)
             failed = True
     
 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 launch_paramiko_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None):
-    """launch a tunner with paramiko in a subprocess"""
+def paramiko_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, password=None, timeout=15):
+    """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`.
+    
+    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.
+    
+    """
     if paramiko is None:
         raise ImportError("Paramiko not available")
-    server = _split_server(server)
-    if keyfile is None:
-        passwd = getpass("%s@%s's password: "%(server[0], server[1]))
-    else:
-        passwd = None
+    
+    if password is None:
+        if not _check_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=passwd))
+            kwargs=dict(keyfile=keyfile, password=password))
     p.daemon=False
     p.start()
+    atexit.register(_shutdown_process, p)
     return p
     
+def _shutdown_process(p):
+    if p.isalive():
+        p.terminate()
 
 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).
     """
-    username, server, port = server
+    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)
 
-    print ('Now forwarding port %d to %s:%d ...' % (lport, server, rport))
+    # print ('Now forwarding port %d to %s:%d ...' % (lport, server, rport))
 
     try:
         forward_tunnel(lport, remoteip, rport, client.get_transport())
     except KeyboardInterrupt:
-        print ('C-c: Port forwarding stopped.')
+        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__ = ['launch_ssh_tunnel', 'launch_paramiko_tunnel']
+__all__ = ['tunnel_connection', 'ssh_tunnel', 'openssh_tunnel', 'paramiko_tunnel', 'try_passwordless_ssh']