##// END OF EJS Templates
ConnectionFileMixin is Configurable...
MinRK -
Show More
@@ -1,557 +1,556 b''
1 1 """Utilities for connecting to kernels
2 2
3 3 Authors:
4 4
5 5 * Min Ragan-Kelley
6 6
7 7 """
8 8
9 9 #-----------------------------------------------------------------------------
10 10 # Copyright (C) 2013 The IPython Development Team
11 11 #
12 12 # Distributed under the terms of the BSD License. The full license is in
13 13 # the file COPYING, distributed as part of this software.
14 14 #-----------------------------------------------------------------------------
15 15
16 16 #-----------------------------------------------------------------------------
17 17 # Imports
18 18 #-----------------------------------------------------------------------------
19 19
20 20 from __future__ import absolute_import
21 21
22 22 import glob
23 23 import json
24 24 import os
25 25 import socket
26 26 import sys
27 27 from getpass import getpass
28 28 from subprocess import Popen, PIPE
29 29 import tempfile
30 30
31 31 import zmq
32 32
33 33 # external imports
34 34 from IPython.external.ssh import tunnel
35 35
36 36 # IPython imports
37 # from IPython.config import Configurable
37 from IPython.config import Configurable
38 38 from IPython.core.profiledir import ProfileDir
39 39 from IPython.utils.localinterfaces import LOCALHOST
40 40 from IPython.utils.path import filefind, get_ipython_dir
41 41 from IPython.utils.py3compat import str_to_bytes, bytes_to_str
42 42 from IPython.utils.traitlets import (
43 43 Bool, Integer, Unicode, CaselessStrEnum,
44 HasTraits,
45 44 )
46 45
47 46
48 47 #-----------------------------------------------------------------------------
49 48 # Working with Connection Files
50 49 #-----------------------------------------------------------------------------
51 50
52 51 def write_connection_file(fname=None, shell_port=0, iopub_port=0, stdin_port=0, hb_port=0,
53 52 control_port=0, ip=LOCALHOST, key=b'', transport='tcp',
54 53 signature_scheme='hmac-sha256',
55 54 ):
56 55 """Generates a JSON config file, including the selection of random ports.
57 56
58 57 Parameters
59 58 ----------
60 59
61 60 fname : unicode
62 61 The path to the file to write
63 62
64 63 shell_port : int, optional
65 64 The port to use for ROUTER (shell) channel.
66 65
67 66 iopub_port : int, optional
68 67 The port to use for the SUB channel.
69 68
70 69 stdin_port : int, optional
71 70 The port to use for the ROUTER (raw input) channel.
72 71
73 72 control_port : int, optional
74 73 The port to use for the ROUTER (control) channel.
75 74
76 75 hb_port : int, optional
77 76 The port to use for the heartbeat REP channel.
78 77
79 78 ip : str, optional
80 79 The ip address the kernel will bind to.
81 80
82 81 key : str, optional
83 82 The Session key used for message authentication.
84 83
85 84 signature_scheme : str, optional
86 85 The scheme used for message authentication.
87 86 This has the form 'digest-hash', where 'digest'
88 87 is the scheme used for digests, and 'hash' is the name of the hash function
89 88 used by the digest scheme.
90 89 Currently, 'hmac' is the only supported digest scheme,
91 90 and 'sha256' is the default hash function.
92 91
93 92 """
94 93 # default to temporary connector file
95 94 if not fname:
96 95 fname = tempfile.mktemp('.json')
97 96
98 97 # Find open ports as necessary.
99 98
100 99 ports = []
101 100 ports_needed = int(shell_port <= 0) + \
102 101 int(iopub_port <= 0) + \
103 102 int(stdin_port <= 0) + \
104 103 int(control_port <= 0) + \
105 104 int(hb_port <= 0)
106 105 if transport == 'tcp':
107 106 for i in range(ports_needed):
108 107 sock = socket.socket()
109 108 sock.bind(('', 0))
110 109 ports.append(sock)
111 110 for i, sock in enumerate(ports):
112 111 port = sock.getsockname()[1]
113 112 sock.close()
114 113 ports[i] = port
115 114 else:
116 115 N = 1
117 116 for i in range(ports_needed):
118 117 while os.path.exists("%s-%s" % (ip, str(N))):
119 118 N += 1
120 119 ports.append(N)
121 120 N += 1
122 121 if shell_port <= 0:
123 122 shell_port = ports.pop(0)
124 123 if iopub_port <= 0:
125 124 iopub_port = ports.pop(0)
126 125 if stdin_port <= 0:
127 126 stdin_port = ports.pop(0)
128 127 if control_port <= 0:
129 128 control_port = ports.pop(0)
130 129 if hb_port <= 0:
131 130 hb_port = ports.pop(0)
132 131
133 132 cfg = dict( shell_port=shell_port,
134 133 iopub_port=iopub_port,
135 134 stdin_port=stdin_port,
136 135 control_port=control_port,
137 136 hb_port=hb_port,
138 137 )
139 138 cfg['ip'] = ip
140 139 cfg['key'] = bytes_to_str(key)
141 140 cfg['transport'] = transport
142 141 cfg['signature_scheme'] = signature_scheme
143 142
144 143 with open(fname, 'w') as f:
145 144 f.write(json.dumps(cfg, indent=2))
146 145
147 146 return fname, cfg
148 147
149 148
150 149 def get_connection_file(app=None):
151 150 """Return the path to the connection file of an app
152 151
153 152 Parameters
154 153 ----------
155 154 app : IPKernelApp instance [optional]
156 155 If unspecified, the currently running app will be used
157 156 """
158 157 if app is None:
159 158 from IPython.kernel.zmq.kernelapp import IPKernelApp
160 159 if not IPKernelApp.initialized():
161 160 raise RuntimeError("app not specified, and not in a running Kernel")
162 161
163 162 app = IPKernelApp.instance()
164 163 return filefind(app.connection_file, ['.', app.profile_dir.security_dir])
165 164
166 165
167 166 def find_connection_file(filename, profile=None):
168 167 """find a connection file, and return its absolute path.
169 168
170 169 The current working directory and the profile's security
171 170 directory will be searched for the file if it is not given by
172 171 absolute path.
173 172
174 173 If profile is unspecified, then the current running application's
175 174 profile will be used, or 'default', if not run from IPython.
176 175
177 176 If the argument does not match an existing file, it will be interpreted as a
178 177 fileglob, and the matching file in the profile's security dir with
179 178 the latest access time will be used.
180 179
181 180 Parameters
182 181 ----------
183 182 filename : str
184 183 The connection file or fileglob to search for.
185 184 profile : str [optional]
186 185 The name of the profile to use when searching for the connection file,
187 186 if different from the current IPython session or 'default'.
188 187
189 188 Returns
190 189 -------
191 190 str : The absolute path of the connection file.
192 191 """
193 192 from IPython.core.application import BaseIPythonApplication as IPApp
194 193 try:
195 194 # quick check for absolute path, before going through logic
196 195 return filefind(filename)
197 196 except IOError:
198 197 pass
199 198
200 199 if profile is None:
201 200 # profile unspecified, check if running from an IPython app
202 201 if IPApp.initialized():
203 202 app = IPApp.instance()
204 203 profile_dir = app.profile_dir
205 204 else:
206 205 # not running in IPython, use default profile
207 206 profile_dir = ProfileDir.find_profile_dir_by_name(get_ipython_dir(), 'default')
208 207 else:
209 208 # find profiledir by profile name:
210 209 profile_dir = ProfileDir.find_profile_dir_by_name(get_ipython_dir(), profile)
211 210 security_dir = profile_dir.security_dir
212 211
213 212 try:
214 213 # first, try explicit name
215 214 return filefind(filename, ['.', security_dir])
216 215 except IOError:
217 216 pass
218 217
219 218 # not found by full name
220 219
221 220 if '*' in filename:
222 221 # given as a glob already
223 222 pat = filename
224 223 else:
225 224 # accept any substring match
226 225 pat = '*%s*' % filename
227 226 matches = glob.glob( os.path.join(security_dir, pat) )
228 227 if not matches:
229 228 raise IOError("Could not find %r in %r" % (filename, security_dir))
230 229 elif len(matches) == 1:
231 230 return matches[0]
232 231 else:
233 232 # get most recent match, by access time:
234 233 return sorted(matches, key=lambda f: os.stat(f).st_atime)[-1]
235 234
236 235
237 236 def get_connection_info(connection_file=None, unpack=False, profile=None):
238 237 """Return the connection information for the current Kernel.
239 238
240 239 Parameters
241 240 ----------
242 241 connection_file : str [optional]
243 242 The connection file to be used. Can be given by absolute path, or
244 243 IPython will search in the security directory of a given profile.
245 244 If run from IPython,
246 245
247 246 If unspecified, the connection file for the currently running
248 247 IPython Kernel will be used, which is only allowed from inside a kernel.
249 248 unpack : bool [default: False]
250 249 if True, return the unpacked dict, otherwise just the string contents
251 250 of the file.
252 251 profile : str [optional]
253 252 The name of the profile to use when searching for the connection file,
254 253 if different from the current IPython session or 'default'.
255 254
256 255
257 256 Returns
258 257 -------
259 258 The connection dictionary of the current kernel, as string or dict,
260 259 depending on `unpack`.
261 260 """
262 261 if connection_file is None:
263 262 # get connection file from current kernel
264 263 cf = get_connection_file()
265 264 else:
266 265 # connection file specified, allow shortnames:
267 266 cf = find_connection_file(connection_file, profile=profile)
268 267
269 268 with open(cf) as f:
270 269 info = f.read()
271 270
272 271 if unpack:
273 272 info = json.loads(info)
274 273 # ensure key is bytes:
275 274 info['key'] = str_to_bytes(info.get('key', ''))
276 275 return info
277 276
278 277
279 278 def connect_qtconsole(connection_file=None, argv=None, profile=None):
280 279 """Connect a qtconsole to the current kernel.
281 280
282 281 This is useful for connecting a second qtconsole to a kernel, or to a
283 282 local notebook.
284 283
285 284 Parameters
286 285 ----------
287 286 connection_file : str [optional]
288 287 The connection file to be used. Can be given by absolute path, or
289 288 IPython will search in the security directory of a given profile.
290 289 If run from IPython,
291 290
292 291 If unspecified, the connection file for the currently running
293 292 IPython Kernel will be used, which is only allowed from inside a kernel.
294 293 argv : list [optional]
295 294 Any extra args to be passed to the console.
296 295 profile : str [optional]
297 296 The name of the profile to use when searching for the connection file,
298 297 if different from the current IPython session or 'default'.
299 298
300 299
301 300 Returns
302 301 -------
303 302 subprocess.Popen instance running the qtconsole frontend
304 303 """
305 304 argv = [] if argv is None else argv
306 305
307 306 if connection_file is None:
308 307 # get connection file from current kernel
309 308 cf = get_connection_file()
310 309 else:
311 310 cf = find_connection_file(connection_file, profile=profile)
312 311
313 312 cmd = ';'.join([
314 313 "from IPython.qt.console import qtconsoleapp",
315 314 "qtconsoleapp.main()"
316 315 ])
317 316
318 317 return Popen([sys.executable, '-c', cmd, '--existing', cf] + argv,
319 318 stdout=PIPE, stderr=PIPE, close_fds=(sys.platform != 'win32'),
320 319 )
321 320
322 321
323 322 def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
324 323 """tunnel connections to a kernel via ssh
325 324
326 325 This will open four SSH tunnels from localhost on this machine to the
327 326 ports associated with the kernel. They can be either direct
328 327 localhost-localhost tunnels, or if an intermediate server is necessary,
329 328 the kernel must be listening on a public IP.
330 329
331 330 Parameters
332 331 ----------
333 332 connection_info : dict or str (path)
334 333 Either a connection dict, or the path to a JSON connection file
335 334 sshserver : str
336 335 The ssh sever to use to tunnel to the kernel. Can be a full
337 336 `user@server:port` string. ssh config aliases are respected.
338 337 sshkey : str [optional]
339 338 Path to file containing ssh key to use for authentication.
340 339 Only necessary if your ssh config does not already associate
341 340 a keyfile with the host.
342 341
343 342 Returns
344 343 -------
345 344
346 345 (shell, iopub, stdin, hb) : ints
347 346 The four ports on localhost that have been forwarded to the kernel.
348 347 """
349 348 if isinstance(connection_info, basestring):
350 349 # it's a path, unpack it
351 350 with open(connection_info) as f:
352 351 connection_info = json.loads(f.read())
353 352
354 353 cf = connection_info
355 354
356 355 lports = tunnel.select_random_ports(4)
357 356 rports = cf['shell_port'], cf['iopub_port'], cf['stdin_port'], cf['hb_port']
358 357
359 358 remote_ip = cf['ip']
360 359
361 360 if tunnel.try_passwordless_ssh(sshserver, sshkey):
362 361 password=False
363 362 else:
364 363 password = getpass("SSH Password for %s: "%sshserver)
365 364
366 365 for lp,rp in zip(lports, rports):
367 366 tunnel.ssh_tunnel(lp, rp, sshserver, remote_ip, sshkey, password)
368 367
369 368 return tuple(lports)
370 369
371 370
372 371 #-----------------------------------------------------------------------------
373 372 # Mixin for classes that work with connection files
374 373 #-----------------------------------------------------------------------------
375 374
376 375 channel_socket_types = {
377 376 'hb' : zmq.REQ,
378 377 'shell' : zmq.DEALER,
379 378 'iopub' : zmq.SUB,
380 379 'stdin' : zmq.DEALER,
381 380 'control': zmq.DEALER,
382 381 }
383 382
384 383 port_names = [ "%s_port" % channel for channel in ('shell', 'stdin', 'iopub', 'hb', 'control')]
385 384
386 class ConnectionFileMixin(HasTraits):
385 class ConnectionFileMixin(Configurable):
387 386 """Mixin for configurable classes that work with connection files"""
388 387
389 388 # The addresses for the communication channels
390 389 connection_file = Unicode('')
391 390 _connection_file_written = Bool(False)
392 391
393 392 transport = CaselessStrEnum(['tcp', 'ipc'], default_value='tcp', config=True)
394 393
395 394 ip = Unicode(LOCALHOST, config=True,
396 395 help="""Set the kernel\'s IP address [default localhost].
397 396 If the IP address is something other than localhost, then
398 397 Consoles on other machines will be able to connect
399 398 to the Kernel, so be careful!"""
400 399 )
401 400
402 401 def _ip_default(self):
403 402 if self.transport == 'ipc':
404 403 if self.connection_file:
405 404 return os.path.splitext(self.connection_file)[0] + '-ipc'
406 405 else:
407 406 return 'kernel-ipc'
408 407 else:
409 408 return LOCALHOST
410 409
411 410 def _ip_changed(self, name, old, new):
412 411 if new == '*':
413 412 self.ip = '0.0.0.0'
414 413
415 414 # protected traits
416 415
417 416 shell_port = Integer(0)
418 417 iopub_port = Integer(0)
419 418 stdin_port = Integer(0)
420 419 control_port = Integer(0)
421 420 hb_port = Integer(0)
422 421
423 422 @property
424 423 def ports(self):
425 424 return [ getattr(self, name) for name in port_names ]
426 425
427 426 #--------------------------------------------------------------------------
428 427 # Connection and ipc file management
429 428 #--------------------------------------------------------------------------
430 429
431 430 def get_connection_info(self):
432 431 """return the connection info as a dict"""
433 432 return dict(
434 433 transport=self.transport,
435 434 ip=self.ip,
436 435 shell_port=self.shell_port,
437 436 iopub_port=self.iopub_port,
438 437 stdin_port=self.stdin_port,
439 438 hb_port=self.hb_port,
440 439 control_port=self.control_port,
441 440 signature_scheme=self.session.signature_scheme,
442 441 key=self.session.key,
443 442 )
444 443
445 444 def cleanup_connection_file(self):
446 445 """Cleanup connection file *if we wrote it*
447 446
448 447 Will not raise if the connection file was already removed somehow.
449 448 """
450 449 if self._connection_file_written:
451 450 # cleanup connection files on full shutdown of kernel we started
452 451 self._connection_file_written = False
453 452 try:
454 453 os.remove(self.connection_file)
455 454 except (IOError, OSError, AttributeError):
456 455 pass
457 456
458 457 def cleanup_ipc_files(self):
459 458 """Cleanup ipc files if we wrote them."""
460 459 if self.transport != 'ipc':
461 460 return
462 461 for port in self.ports:
463 462 ipcfile = "%s-%i" % (self.ip, port)
464 463 try:
465 464 os.remove(ipcfile)
466 465 except (IOError, OSError):
467 466 pass
468 467
469 468 def write_connection_file(self):
470 469 """Write connection info to JSON dict in self.connection_file."""
471 470 if self._connection_file_written:
472 471 return
473 472
474 473 self.connection_file, cfg = write_connection_file(self.connection_file,
475 474 transport=self.transport, ip=self.ip, key=self.session.key,
476 475 stdin_port=self.stdin_port, iopub_port=self.iopub_port,
477 476 shell_port=self.shell_port, hb_port=self.hb_port,
478 477 control_port=self.control_port,
479 478 signature_scheme=self.session.signature_scheme,
480 479 )
481 480 # write_connection_file also sets default ports:
482 481 for name in port_names:
483 482 setattr(self, name, cfg[name])
484 483
485 484 self._connection_file_written = True
486 485
487 486 def load_connection_file(self):
488 487 """Load connection info from JSON dict in self.connection_file."""
489 488 with open(self.connection_file) as f:
490 489 cfg = json.loads(f.read())
491 490
492 491 self.transport = cfg.get('transport', 'tcp')
493 492 self.ip = cfg['ip']
494 493 for name in port_names:
495 494 setattr(self, name, cfg[name])
496 495 if 'key' in cfg:
497 496 self.session.key = str_to_bytes(cfg['key'])
498 497 if cfg.get('signature_scheme'):
499 498 self.session.signature_scheme = cfg['signature_scheme']
500 499
501 500 #--------------------------------------------------------------------------
502 501 # Creating connected sockets
503 502 #--------------------------------------------------------------------------
504 503
505 504 def _make_url(self, channel):
506 505 """Make a ZeroMQ URL for a given channel."""
507 506 transport = self.transport
508 507 ip = self.ip
509 508 port = getattr(self, '%s_port' % channel)
510 509
511 510 if transport == 'tcp':
512 511 return "tcp://%s:%i" % (ip, port)
513 512 else:
514 513 return "%s://%s-%s" % (transport, ip, port)
515 514
516 515 def _create_connected_socket(self, channel, identity=None):
517 516 """Create a zmq Socket and connect it to the kernel."""
518 517 url = self._make_url(channel)
519 518 socket_type = channel_socket_types[channel]
520 519 self.log.info("Connecting to: %s" % url)
521 520 sock = self.context.socket(socket_type)
522 521 if identity:
523 522 sock.identity = identity
524 523 sock.connect(url)
525 524 return sock
526 525
527 526 def connect_iopub(self, identity=None):
528 527 """return zmq Socket connected to the IOPub channel"""
529 528 sock = self._create_connected_socket('iopub', identity=identity)
530 529 sock.setsockopt(zmq.SUBSCRIBE, b'')
531 530 return sock
532 531
533 532 def connect_shell(self, identity=None):
534 533 """return zmq Socket connected to the Shell channel"""
535 534 return self._create_connected_socket('shell', identity=identity)
536 535
537 536 def connect_stdin(self, identity=None):
538 537 """return zmq Socket connected to the StdIn channel"""
539 538 return self._create_connected_socket('stdin', identity=identity)
540 539
541 540 def connect_hb(self, identity=None):
542 541 """return zmq Socket connected to the Heartbeat channel"""
543 542 return self._create_connected_socket('hb', identity=identity)
544 543
545 544 def connect_control(self, identity=None):
546 545 """return zmq Socket connected to the Heartbeat channel"""
547 546 return self._create_connected_socket('control', identity=identity)
548 547
549 548
550 549 __all__ = [
551 550 'write_connection_file',
552 551 'get_connection_file',
553 552 'find_connection_file',
554 553 'get_connection_info',
555 554 'connect_qtconsole',
556 555 'tunnel_to_kernel',
557 556 ]
General Comments 0
You need to be logged in to leave comments. Login now