connect.py
448 lines
| 14.8 KiB
| text/x-python
|
PythonLexer
/ jupyter_client / connect.py
Min RK
|
r21027 | """Utilities for connecting to jupyter kernels | ||
MinRK
|
r9351 | |||
Thomas Kluyver
|
r16640 | The :class:`ConnectionFileMixin` class in this module encapsulates the logic | ||
related to writing and reading connections files. | ||||
MinRK
|
r9351 | """ | ||
Min RK
|
r21027 | |||
# Copyright (c) Jupyter Development Team. | ||||
Paul Ivanov
|
r16497 | # Distributed under the terms of the Modified BSD License. | ||
MinRK
|
r9351 | |||
MinRK
|
r10324 | from __future__ import absolute_import | ||
MinRK
|
r9351 | import glob | ||
import json | ||||
import os | ||||
MinRK
|
r9352 | import socket | ||
MinRK
|
r9351 | from getpass import getpass | ||
MinRK
|
r9352 | import tempfile | ||
MinRK
|
r9351 | |||
MinRK
|
r10324 | import zmq | ||
MinRK
|
r9351 | |||
MinRK
|
r16732 | from IPython.config import LoggingConfigurable | ||
Min RK
|
r21098 | from .localinterfaces import localhost | ||
Min RK
|
r21027 | from IPython.utils.path import filefind | ||
Thomas Kluyver
|
r13353 | from IPython.utils.py3compat import (str_to_bytes, bytes_to_str, cast_bytes_py2, | ||
string_types) | ||||
MinRK
|
r10285 | from IPython.utils.traitlets import ( | ||
MinRK
|
r16731 | Bool, Integer, Unicode, CaselessStrEnum, Instance, | ||
MinRK
|
r10285 | ) | ||
MinRK
|
r9351 | |||
MinRK
|
r9352 | def write_connection_file(fname=None, shell_port=0, iopub_port=0, stdin_port=0, hb_port=0, | ||
MinRK
|
r12591 | control_port=0, ip='', key=b'', transport='tcp', | ||
MinRK
|
r11656 | signature_scheme='hmac-sha256', | ||
): | ||||
MinRK
|
r9352 | """Generates a JSON config file, including the selection of random ports. | ||
Min RK
|
r20949 | |||
MinRK
|
r9352 | Parameters | ||
---------- | ||||
fname : unicode | ||||
The path to the file to write | ||||
shell_port : int, optional | ||||
MinRK
|
r10328 | The port to use for ROUTER (shell) channel. | ||
MinRK
|
r9352 | |||
iopub_port : int, optional | ||||
The port to use for the SUB channel. | ||||
stdin_port : int, optional | ||||
MinRK
|
r10296 | The port to use for the ROUTER (raw input) channel. | ||
control_port : int, optional | ||||
MinRK
|
r10328 | The port to use for the ROUTER (control) channel. | ||
MinRK
|
r9352 | |||
hb_port : int, optional | ||||
MinRK
|
r10328 | The port to use for the heartbeat REP channel. | ||
MinRK
|
r9352 | |||
ip : str, optional | ||||
The ip address the kernel will bind to. | ||||
key : str, optional | ||||
MinRK
|
r11656 | The Session key used for message authentication. | ||
Min RK
|
r20949 | |||
MinRK
|
r11656 | signature_scheme : str, optional | ||
The scheme used for message authentication. | ||||
This has the form 'digest-hash', where 'digest' | ||||
is the scheme used for digests, and 'hash' is the name of the hash function | ||||
used by the digest scheme. | ||||
Currently, 'hmac' is the only supported digest scheme, | ||||
and 'sha256' is the default hash function. | ||||
MinRK
|
r9352 | |||
""" | ||||
MinRK
|
r12591 | if not ip: | ||
ip = localhost() | ||||
MinRK
|
r9352 | # default to temporary connector file | ||
if not fname: | ||||
Julian Taylor
|
r15372 | fd, fname = tempfile.mkstemp('.json') | ||
os.close(fd) | ||||
Min RK
|
r20949 | |||
MinRK
|
r9352 | # Find open ports as necessary. | ||
Min RK
|
r20949 | |||
MinRK
|
r9352 | ports = [] | ||
MinRK
|
r10296 | ports_needed = int(shell_port <= 0) + \ | ||
int(iopub_port <= 0) + \ | ||||
int(stdin_port <= 0) + \ | ||||
int(control_port <= 0) + \ | ||||
int(hb_port <= 0) | ||||
MinRK
|
r9352 | if transport == 'tcp': | ||
for i in range(ports_needed): | ||||
sock = socket.socket() | ||||
MinRK
|
r15549 | # struct.pack('ii', (0,0)) is 8 null bytes | ||
MinRK
|
r15527 | sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, b'\0' * 8) | ||
MinRK
|
r9352 | sock.bind(('', 0)) | ||
ports.append(sock) | ||||
for i, sock in enumerate(ports): | ||||
port = sock.getsockname()[1] | ||||
sock.close() | ||||
ports[i] = port | ||||
else: | ||||
N = 1 | ||||
for i in range(ports_needed): | ||||
while os.path.exists("%s-%s" % (ip, str(N))): | ||||
N += 1 | ||||
ports.append(N) | ||||
N += 1 | ||||
if shell_port <= 0: | ||||
shell_port = ports.pop(0) | ||||
if iopub_port <= 0: | ||||
iopub_port = ports.pop(0) | ||||
if stdin_port <= 0: | ||||
stdin_port = ports.pop(0) | ||||
MinRK
|
r10296 | if control_port <= 0: | ||
control_port = ports.pop(0) | ||||
MinRK
|
r9352 | if hb_port <= 0: | ||
hb_port = ports.pop(0) | ||||
Min RK
|
r20949 | |||
MinRK
|
r9352 | cfg = dict( shell_port=shell_port, | ||
iopub_port=iopub_port, | ||||
stdin_port=stdin_port, | ||||
MinRK
|
r10296 | control_port=control_port, | ||
MinRK
|
r9352 | hb_port=hb_port, | ||
) | ||||
cfg['ip'] = ip | ||||
cfg['key'] = bytes_to_str(key) | ||||
cfg['transport'] = transport | ||||
MinRK
|
r11656 | cfg['signature_scheme'] = signature_scheme | ||
Min RK
|
r20949 | |||
MinRK
|
r9352 | with open(fname, 'w') as f: | ||
f.write(json.dumps(cfg, indent=2)) | ||||
Min RK
|
r20949 | |||
MinRK
|
r9352 | return fname, cfg | ||
Min RK
|
r21027 | def find_connection_file(filename='kernel-*.json', path=None): | ||
MinRK
|
r9351 | """find a connection file, and return its absolute path. | ||
Min RK
|
r20949 | |||
MinRK
|
r9351 | The current working directory and the profile's security | ||
directory will be searched for the file if it is not given by | ||||
absolute path. | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | If profile is unspecified, then the current running application's | ||
profile will be used, or 'default', if not run from IPython. | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | If the argument does not match an existing file, it will be interpreted as a | ||
fileglob, and the matching file in the profile's security dir with | ||||
the latest access time will be used. | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | Parameters | ||
---------- | ||||
filename : str | ||||
The connection file or fileglob to search for. | ||||
Min RK
|
r21027 | path : str or list of strs[optional] | ||
Paths in which to search for connection files. | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | Returns | ||
------- | ||||
str : The absolute path of the connection file. | ||||
""" | ||||
Min RK
|
r21027 | if path is None: | ||
path = ['.'] | ||||
if isinstance(path, string_types): | ||||
path = [path] | ||||
MinRK
|
r9351 | try: | ||
# first, try explicit name | ||||
Min RK
|
r21027 | return filefind(filename, path) | ||
MinRK
|
r9351 | except IOError: | ||
pass | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | # not found by full name | ||
Min RK
|
r20949 | |||
MinRK
|
r9351 | if '*' in filename: | ||
# given as a glob already | ||||
pat = filename | ||||
else: | ||||
# accept any substring match | ||||
pat = '*%s*' % filename | ||||
Min RK
|
r21027 | |||
matches = [] | ||||
for p in path: | ||||
matches.extend(glob.glob(os.path.join(p, pat))) | ||||
MinRK
|
r9351 | if not matches: | ||
Min RK
|
r21027 | raise IOError("Could not find %r in %r" % (filename, path)) | ||
MinRK
|
r9351 | elif len(matches) == 1: | ||
return matches[0] | ||||
else: | ||||
# get most recent match, by access time: | ||||
return sorted(matches, key=lambda f: os.stat(f).st_atime)[-1] | ||||
MinRK
|
r9352 | |||
MinRK
|
r9351 | def tunnel_to_kernel(connection_info, sshserver, sshkey=None): | ||
"""tunnel connections to a kernel via ssh | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | 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. | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | 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. | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | Returns | ||
------- | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | (shell, iopub, stdin, hb) : ints | ||
The four ports on localhost that have been forwarded to the kernel. | ||||
""" | ||||
MinRK
|
r17054 | from zmq.ssh import tunnel | ||
Thomas Kluyver
|
r13353 | if isinstance(connection_info, string_types): | ||
MinRK
|
r9351 | # it's a path, unpack it | ||
with open(connection_info) as f: | ||||
connection_info = json.loads(f.read()) | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | cf = connection_info | ||
Min RK
|
r20949 | |||
MinRK
|
r9351 | lports = tunnel.select_random_ports(4) | ||
rports = cf['shell_port'], cf['iopub_port'], cf['stdin_port'], cf['hb_port'] | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | remote_ip = cf['ip'] | ||
Min RK
|
r20949 | |||
MinRK
|
r9351 | if tunnel.try_passwordless_ssh(sshserver, sshkey): | ||
password=False | ||||
else: | ||||
Thomas Kluyver
|
r12893 | password = getpass("SSH Password for %s: " % cast_bytes_py2(sshserver)) | ||
Min RK
|
r20949 | |||
MinRK
|
r9351 | for lp,rp in zip(lports, rports): | ||
tunnel.ssh_tunnel(lp, rp, sshserver, remote_ip, sshkey, password) | ||||
Min RK
|
r20949 | |||
MinRK
|
r9351 | return tuple(lports) | ||
MinRK
|
r9352 | |||
MinRK
|
r10285 | |||
#----------------------------------------------------------------------------- | ||||
MinRK
|
r10324 | # Mixin for classes that work with connection files | ||
MinRK
|
r10285 | #----------------------------------------------------------------------------- | ||
MinRK
|
r10324 | |||
channel_socket_types = { | ||||
'hb' : zmq.REQ, | ||||
'shell' : zmq.DEALER, | ||||
'iopub' : zmq.SUB, | ||||
'stdin' : zmq.DEALER, | ||||
'control': zmq.DEALER, | ||||
} | ||||
MinRK
|
r10296 | port_names = [ "%s_port" % channel for channel in ('shell', 'stdin', 'iopub', 'hb', 'control')] | ||
MinRK
|
r10285 | |||
MinRK
|
r16732 | class ConnectionFileMixin(LoggingConfigurable): | ||
MinRK
|
r10285 | """Mixin for configurable classes that work with connection files""" | ||
# The addresses for the communication channels | ||||
Min RK
|
r20949 | connection_file = Unicode('', config=True, | ||
Paul Ivanov
|
r16503 | help="""JSON file in which to store connection info [default: kernel-<pid>.json] | ||
Min RK
|
r20949 | |||
Paul Ivanov
|
r16503 | This file will contain the IP, ports, and authentication key needed to connect | ||
clients to this kernel. By default, this file will be created in the security dir | ||||
of the current profile, but can be specified by absolute path. | ||||
""") | ||||
MinRK
|
r10285 | _connection_file_written = Bool(False) | ||
transport = CaselessStrEnum(['tcp', 'ipc'], default_value='tcp', config=True) | ||||
MinRK
|
r12591 | ip = Unicode(config=True, | ||
MinRK
|
r10285 | help="""Set the kernel\'s IP address [default localhost]. | ||
If the IP address is something other than localhost, then | ||||
Consoles on other machines will be able to connect | ||||
to the Kernel, so be careful!""" | ||||
) | ||||
def _ip_default(self): | ||||
if self.transport == 'ipc': | ||||
if self.connection_file: | ||||
return os.path.splitext(self.connection_file)[0] + '-ipc' | ||||
else: | ||||
return 'kernel-ipc' | ||||
else: | ||||
MinRK
|
r12591 | return localhost() | ||
MinRK
|
r10285 | |||
def _ip_changed(self, name, old, new): | ||||
if new == '*': | ||||
self.ip = '0.0.0.0' | ||||
# protected traits | ||||
Paul Ivanov
|
r16504 | hb_port = Integer(0, config=True, | ||
help="set the heartbeat port [default: random]") | ||||
shell_port = Integer(0, config=True, | ||||
help="set the shell (ROUTER) port [default: random]") | ||||
iopub_port = Integer(0, config=True, | ||||
help="set the iopub (PUB) port [default: random]") | ||||
stdin_port = Integer(0, config=True, | ||||
help="set the stdin (ROUTER) port [default: random]") | ||||
control_port = Integer(0, config=True, | ||||
help="set the control (ROUTER) port [default: random]") | ||||
MinRK
|
r10285 | |||
MinRK
|
r10296 | @property | ||
def ports(self): | ||||
return [ getattr(self, name) for name in port_names ] | ||||
MinRK
|
r16731 | # The Session to use for communication with the kernel. | ||
Min RK
|
r20951 | session = Instance('jupyter_client.session.Session') | ||
MinRK
|
r16731 | def _session_default(self): | ||
Min RK
|
r20951 | from jupyter_client.session import Session | ||
MinRK
|
r16731 | return Session(parent=self) | ||
MinRK
|
r10285 | #-------------------------------------------------------------------------- | ||
# Connection and ipc file management | ||||
#-------------------------------------------------------------------------- | ||||
MinRK
|
r10296 | def get_connection_info(self): | ||
"""return the connection info as a dict""" | ||||
return dict( | ||||
transport=self.transport, | ||||
ip=self.ip, | ||||
shell_port=self.shell_port, | ||||
iopub_port=self.iopub_port, | ||||
stdin_port=self.stdin_port, | ||||
hb_port=self.hb_port, | ||||
control_port=self.control_port, | ||||
MinRK
|
r11840 | signature_scheme=self.session.signature_scheme, | ||
MinRK
|
r11816 | key=self.session.key, | ||
MinRK
|
r10296 | ) | ||
MinRK
|
r10285 | def cleanup_connection_file(self): | ||
"""Cleanup connection file *if we wrote it* | ||||
Will not raise if the connection file was already removed somehow. | ||||
""" | ||||
if self._connection_file_written: | ||||
# cleanup connection files on full shutdown of kernel we started | ||||
self._connection_file_written = False | ||||
try: | ||||
os.remove(self.connection_file) | ||||
except (IOError, OSError, AttributeError): | ||||
pass | ||||
def cleanup_ipc_files(self): | ||||
"""Cleanup ipc files if we wrote them.""" | ||||
if self.transport != 'ipc': | ||||
return | ||||
MinRK
|
r10296 | for port in self.ports: | ||
MinRK
|
r10285 | ipcfile = "%s-%i" % (self.ip, port) | ||
try: | ||||
os.remove(ipcfile) | ||||
except (IOError, OSError): | ||||
pass | ||||
def write_connection_file(self): | ||||
"""Write connection info to JSON dict in self.connection_file.""" | ||||
Pankaj Pandey
|
r15680 | if self._connection_file_written and os.path.exists(self.connection_file): | ||
MinRK
|
r10285 | return | ||
MinRK
|
r10296 | |||
self.connection_file, cfg = write_connection_file(self.connection_file, | ||||
MinRK
|
r10285 | transport=self.transport, ip=self.ip, key=self.session.key, | ||
stdin_port=self.stdin_port, iopub_port=self.iopub_port, | ||||
MinRK
|
r10296 | shell_port=self.shell_port, hb_port=self.hb_port, | ||
control_port=self.control_port, | ||||
MinRK
|
r11816 | signature_scheme=self.session.signature_scheme, | ||
MinRK
|
r10296 | ) | ||
MinRK
|
r10285 | # write_connection_file also sets default ports: | ||
MinRK
|
r10296 | for name in port_names: | ||
setattr(self, name, cfg[name]) | ||||
MinRK
|
r10285 | |||
self._connection_file_written = True | ||||
def load_connection_file(self): | ||||
"""Load connection info from JSON dict in self.connection_file.""" | ||||
Paul Ivanov
|
r16497 | self.log.debug(u"Loading connection file %s", self.connection_file) | ||
MinRK
|
r10285 | with open(self.connection_file) as f: | ||
Paul Ivanov
|
r16497 | cfg = json.load(f) | ||
Paul Ivanov
|
r16501 | self.transport = cfg.get('transport', self.transport) | ||
self.ip = cfg.get('ip', self._ip_default()) | ||||
Min RK
|
r20949 | |||
MinRK
|
r10296 | for name in port_names: | ||
Paul Ivanov
|
r16497 | if getattr(self, name) == 0 and name in cfg: | ||
# not overridden by config or cl_args | ||||
setattr(self, name, cfg[name]) | ||||
Min RK
|
r20949 | |||
MinRK
|
r11656 | if 'key' in cfg: | ||
MinRK
|
r16731 | self.session.key = str_to_bytes(cfg['key']) | ||
Paul Ivanov
|
r16497 | if 'signature_scheme' in cfg: | ||
MinRK
|
r16731 | self.session.signature_scheme = cfg['signature_scheme'] | ||
MinRK
|
r16729 | |||
MinRK
|
r10324 | #-------------------------------------------------------------------------- | ||
# Creating connected sockets | ||||
#-------------------------------------------------------------------------- | ||||
def _make_url(self, channel): | ||||
"""Make a ZeroMQ URL for a given channel.""" | ||||
transport = self.transport | ||||
ip = self.ip | ||||
port = getattr(self, '%s_port' % channel) | ||||
if transport == 'tcp': | ||||
return "tcp://%s:%i" % (ip, port) | ||||
else: | ||||
return "%s://%s-%s" % (transport, ip, port) | ||||
def _create_connected_socket(self, channel, identity=None): | ||||
"""Create a zmq Socket and connect it to the kernel.""" | ||||
url = self._make_url(channel) | ||||
socket_type = channel_socket_types[channel] | ||||
MinRK
|
r14645 | self.log.debug("Connecting to: %s" % url) | ||
MinRK
|
r10324 | sock = self.context.socket(socket_type) | ||
MinRK
|
r16526 | # set linger to 1s to prevent hangs at exit | ||
sock.linger = 1000 | ||||
MinRK
|
r10324 | if identity: | ||
sock.identity = identity | ||||
sock.connect(url) | ||||
return sock | ||||
def connect_iopub(self, identity=None): | ||||
"""return zmq Socket connected to the IOPub channel""" | ||||
sock = self._create_connected_socket('iopub', identity=identity) | ||||
sock.setsockopt(zmq.SUBSCRIBE, b'') | ||||
return sock | ||||
def connect_shell(self, identity=None): | ||||
"""return zmq Socket connected to the Shell channel""" | ||||
return self._create_connected_socket('shell', identity=identity) | ||||
def connect_stdin(self, identity=None): | ||||
"""return zmq Socket connected to the StdIn channel""" | ||||
return self._create_connected_socket('stdin', identity=identity) | ||||
def connect_hb(self, identity=None): | ||||
"""return zmq Socket connected to the Heartbeat channel""" | ||||
return self._create_connected_socket('hb', identity=identity) | ||||
def connect_control(self, identity=None): | ||||
Thomas Kluyver
|
r19172 | """return zmq Socket connected to the Control channel""" | ||
MinRK
|
r10324 | return self._create_connected_socket('control', identity=identity) | ||
MinRK
|
r10285 | |||
MinRK
|
r9376 | __all__ = [ | ||
'write_connection_file', | ||||
'find_connection_file', | ||||
'tunnel_to_kernel', | ||||
MinRK
|
r10285 | ] | ||