1"""Basic ssh tunnel utilities, and convenience functions for tunneling
2zeromq connections.
3"""
4
5# Copyright (C) 2010-2011  IPython Development Team
6# Copyright (C) 2011- PyZMQ Developers
7#
8# Redistributed from IPython under the terms of the BSD License.
9
10import atexit
11import os
12import re
13import signal
14import socket
15import sys
16import warnings
17from getpass import getpass, getuser
18from multiprocessing import Process
19
20try:
21    with warnings.catch_warnings():
22        warnings.simplefilter('ignore', DeprecationWarning)
23        import paramiko
24        SSHException = paramiko.ssh_exception.SSHException
25except ImportError:
26    paramiko = None
27    class SSHException(Exception):
28        pass
29else:
30    from .forward import forward_tunnel
31
32try:
33    import pexpect
34except ImportError:
35    pexpect = None
36
37from zmq.utils.strtypes import b
38
39
40def select_random_ports(n):
41    """Select and return n random ports that are available."""
42    ports = []
43    sockets = []
44    for i in range(n):
45        sock = socket.socket()
46        sock.bind(('', 0))
47        ports.append(sock.getsockname()[1])
48        sockets.append(sock)
49    for sock in sockets:
50        sock.close()
51    return ports
52
53
54#-----------------------------------------------------------------------------
55# Check for passwordless login
56#-----------------------------------------------------------------------------
57_password_pat = re.compile(b(r'pass(word|phrase):'), re.IGNORECASE)
58
59
60def try_passwordless_ssh(server, keyfile, paramiko=None):
61    """Attempt to make an ssh connection without a password.
62    This is mainly used for requiring password input only once
63    when many tunnels may be connected to the same server.
64
65    If paramiko is None, the default for the platform is chosen.
66    """
67    if paramiko is None:
68        paramiko = sys.platform == 'win32'
69    if not paramiko:
70        f = _try_passwordless_openssh
71    else:
72        f = _try_passwordless_paramiko
73    return f(server, keyfile)
74
75
76def _try_passwordless_openssh(server, keyfile):
77    """Try passwordless login with shell ssh command."""
78    if pexpect is None:
79        raise ImportError("pexpect unavailable, use paramiko")
80    cmd = 'ssh -f ' + server
81    if keyfile:
82        cmd += ' -i ' + keyfile
83    cmd += ' exit'
84
85    # pop SSH_ASKPASS from env
86    env = os.environ.copy()
87    env.pop('SSH_ASKPASS', None)
88
89    ssh_newkey = 'Are you sure you want to continue connecting'
90    p = pexpect.spawn(cmd, env=env)
91    while True:
92        try:
93            i = p.expect([ssh_newkey, _password_pat], timeout=.1)
94            if i == 0:
95                raise SSHException('The authenticity of the host can\'t be established.')
96        except pexpect.TIMEOUT:
97            continue
98        except pexpect.EOF:
99            return True
100        else:
101            return False
102
103
104def _try_passwordless_paramiko(server, keyfile):
105    """Try passwordless login with paramiko."""
106    if paramiko is None:
107        msg = "Paramiko unavailable, "
108        if sys.platform == 'win32':
109            msg += "Paramiko is required for ssh tunneled connections on Windows."
110        else:
111            msg += "use OpenSSH."
112        raise ImportError(msg)
113    username, server, port = _split_server(server)
114    client = paramiko.SSHClient()
115    client.load_system_host_keys()
116    client.set_missing_host_key_policy(paramiko.WarningPolicy())
117    try:
118        client.connect(server, port, username=username, key_filename=keyfile,
119               look_for_keys=True)
120    except paramiko.AuthenticationException:
121        return False
122    else:
123        client.close()
124        return True
125
126
127def tunnel_connection(socket, addr, server, keyfile=None, password=None, paramiko=None, timeout=60):
128    """Connect a socket to an address via an ssh tunnel.
129
130    This is a wrapper for socket.connect(addr), when addr is not accessible
131    from the local machine.  It simply creates an ssh tunnel using the remaining args,
132    and calls socket.connect('tcp://localhost:lport') where lport is the randomly
133    selected local port of the tunnel.
134
135    """
136    new_url, tunnel = open_tunnel(addr, server, keyfile=keyfile, password=password, paramiko=paramiko, timeout=timeout)
137    socket.connect(new_url)
138    return tunnel
139
140
141def open_tunnel(addr, server, keyfile=None, password=None, paramiko=None, timeout=60):
142    """Open a tunneled connection from a 0MQ url.
143
144    For use inside tunnel_connection.
145
146    Returns
147    -------
148
149    (url, tunnel) : (str, object)
150        The 0MQ url that has been forwarded, and the tunnel object
151    """
152
153    lport = select_random_ports(1)[0]
154    transport, addr = addr.split('://')
155    ip, rport = addr.split(':')
156    rport = int(rport)
157    if paramiko is None:
158        paramiko = sys.platform == 'win32'
159    if paramiko:
160        tunnelf = paramiko_tunnel
161    else:
162        tunnelf = openssh_tunnel
163
164    tunnel = tunnelf(lport, rport, server, remoteip=ip, keyfile=keyfile, password=password, timeout=timeout)
165    return 'tcp://127.0.0.1:%i' % lport, tunnel
166
167
168def openssh_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, password=None, timeout=60):
169    """Create an ssh tunnel using command-line ssh that connects port lport
170    on this machine to localhost:rport on server.  The tunnel
171    will automatically close when not in use, remaining open
172    for a minimum of timeout seconds for an initial connection.
173
174    This creates a tunnel redirecting `localhost:lport` to `remoteip:rport`,
175    as seen from `server`.
176
177    keyfile and password may be specified, but ssh config is checked for defaults.
178
179    Parameters
180    ----------
181
182    lport : int
183        local port for connecting to the tunnel from this machine.
184    rport : int
185        port on the remote machine to connect to.
186    server : str
187        The ssh server to connect to. The full ssh server string will be parsed.
188        user@server:port
189    remoteip : str [Default: 127.0.0.1]
190        The remote ip, specifying the destination of the tunnel.
191        Default is localhost, which means that the tunnel would redirect
192        localhost:lport on this machine to localhost:rport on the *server*.
193
194    keyfile : str; path to public key file
195        This specifies a key to be used in ssh login, default None.
196        Regular default ssh keys will be used without specifying this argument.
197    password : str;
198        Your ssh password to the ssh server. Note that if this is left None,
199        you will be prompted for it if passwordless key based login is unavailable.
200    timeout : int [default: 60]
201        The time (in seconds) after which no activity will result in the tunnel
202        closing.  This prevents orphaned tunnels from running forever.
203    """
204    if pexpect is None:
205        raise ImportError("pexpect unavailable, use paramiko_tunnel")
206    ssh = "ssh "
207    if keyfile:
208        ssh += "-i " + keyfile
209
210    if ':' in server:
211        server, port = server.split(':')
212        ssh += " -p %s" % port
213
214    cmd = "%s -O check %s" % (ssh, server)
215    (output, exitstatus) = pexpect.run(cmd, withexitstatus=True)
216    if not exitstatus:
217        pid = int(output[output.find(b"(pid=")+5:output.find(b")")])
218        cmd = "%s -O forward -L 127.0.0.1:%i:%s:%i %s" % (
219            ssh, lport, remoteip, rport, server)
220        (output, exitstatus) = pexpect.run(cmd, withexitstatus=True)
221        if not exitstatus:
222            atexit.register(_stop_tunnel, cmd.replace("-O forward", "-O cancel", 1))
223            return pid
224    cmd = "%s -f -S none -L 127.0.0.1:%i:%s:%i %s sleep %i" % (
225        ssh, lport, remoteip, rport, server, timeout)
226
227    # pop SSH_ASKPASS from env
228    env = os.environ.copy()
229    env.pop('SSH_ASKPASS', None)
230
231    ssh_newkey = 'Are you sure you want to continue connecting'
232    tunnel = pexpect.spawn(cmd, env=env)
233    failed = False
234    while True:
235        try:
236            i = tunnel.expect([ssh_newkey, _password_pat], timeout=.1)
237            if i == 0:
238                raise SSHException('The authenticity of the host can\'t be established.')
239        except pexpect.TIMEOUT:
240            continue
241        except pexpect.EOF as e:
242            if tunnel.exitstatus:
243                print(tunnel.exitstatus)
244                print(tunnel.before)
245                print(tunnel.after)
246                raise RuntimeError("tunnel '%s' failed to start" % (cmd)) from e
247            else:
248                return tunnel.pid
249        else:
250            if failed:
251                print("Password rejected, try again")
252                password = None
253            if password is None:
254                password = getpass("%s's password: " % (server))
255            tunnel.sendline(password)
256            failed = True
257
258
259def _stop_tunnel(cmd):
260    pexpect.run(cmd)
261
262
263def _split_server(server):
264    if '@' in server:
265        username, server = server.split('@', 1)
266    else:
267        username = getuser()
268    if ':' in server:
269        server, port = server.split(':')
270        port = int(port)
271    else:
272        port = 22
273    return username, server, port
274
275
276def paramiko_tunnel(lport, rport, server, remoteip='127.0.0.1', keyfile=None, password=None, timeout=60):
277    """launch a tunner with paramiko in a subprocess. This should only be used
278    when shell ssh is unavailable (e.g. Windows).
279
280    This creates a tunnel redirecting `localhost:lport` to `remoteip:rport`,
281    as seen from `server`.
282
283    If you are familiar with ssh tunnels, this creates the tunnel:
284
285    ssh server -L localhost:lport:remoteip:rport
286
287    keyfile and password may be specified, but ssh config is checked for defaults.
288
289
290    Parameters
291    ----------
292
293    lport : int
294        local port for connecting to the tunnel from this machine.
295    rport : int
296        port on the remote machine to connect to.
297    server : str
298        The ssh server to connect to. The full ssh server string will be parsed.
299        user@server:port
300    remoteip : str [Default: 127.0.0.1]
301        The remote ip, specifying the destination of the tunnel.
302        Default is localhost, which means that the tunnel would redirect
303        localhost:lport on this machine to localhost:rport on the *server*.
304
305    keyfile : str; path to public key file
306        This specifies a key to be used in ssh login, default None.
307        Regular default ssh keys will be used without specifying this argument.
308    password : str;
309        Your ssh password to the ssh server. Note that if this is left None,
310        you will be prompted for it if passwordless key based login is unavailable.
311    timeout : int [default: 60]
312        The time (in seconds) after which no activity will result in the tunnel
313        closing.  This prevents orphaned tunnels from running forever.
314
315    """
316    if paramiko is None:
317        raise ImportError("Paramiko not available")
318
319    if password is None:
320        if not _try_passwordless_paramiko(server, keyfile):
321            password = getpass("%s's password: " % (server))
322
323    p = Process(target=_paramiko_tunnel,
324                args=(lport, rport, server, remoteip),
325                kwargs=dict(keyfile=keyfile, password=password))
326    p.daemon = True
327    p.start()
328    return p
329
330
331def _paramiko_tunnel(lport, rport, server, remoteip, keyfile=None, password=None):
332    """Function for actually starting a paramiko tunnel, to be passed
333    to multiprocessing.Process(target=this), and not called directly.
334    """
335    username, server, port = _split_server(server)
336    client = paramiko.SSHClient()
337    client.load_system_host_keys()
338    client.set_missing_host_key_policy(paramiko.WarningPolicy())
339
340    try:
341        client.connect(server, port, username=username, key_filename=keyfile,
342                       look_for_keys=True, password=password)
343#    except paramiko.AuthenticationException:
344#        if password is None:
345#            password = getpass("%s@%s's password: "%(username, server))
346#            client.connect(server, port, username=username, password=password)
347#        else:
348#            raise
349    except Exception as e:
350        print('*** Failed to connect to %s:%d: %r' % (server, port, e))
351        sys.exit(1)
352
353    # Don't let SIGINT kill the tunnel subprocess
354    signal.signal(signal.SIGINT, signal.SIG_IGN)
355
356    try:
357        forward_tunnel(lport, remoteip, rport, client.get_transport())
358    except KeyboardInterrupt:
359        print('SIGINT: Port forwarding stopped cleanly')
360        sys.exit(0)
361    except Exception as e:
362        print("Port forwarding stopped uncleanly: %s" % e)
363        sys.exit(255)
364
365
366if sys.platform == 'win32':
367    ssh_tunnel = paramiko_tunnel
368else:
369    ssh_tunnel = openssh_tunnel
370
371
372__all__ = ['tunnel_connection', 'ssh_tunnel', 'openssh_tunnel', 'paramiko_tunnel', 'try_passwordless_ssh']
373