"""Tests for kernel connection utilities""" # Copyright (c) Jupyter Development Team. # Distributed under the terms of the Modified BSD License. import json import os import re import stat import tempfile import shutil from traitlets.config import Config from jupyter_core.application import JupyterApp from jupyter_core.paths import jupyter_runtime_dir from ipython_genutils.tempdir import TemporaryDirectory, TemporaryWorkingDirectory from ipython_genutils.py3compat import str_to_bytes from jupyter_client import connect, KernelClient from jupyter_client.consoleapp import JupyterConsoleApp from jupyter_client.session import Session from jupyter_client.connect import secure_write class DummyConsoleApp(JupyterApp, JupyterConsoleApp): def initialize(self, argv=[]): JupyterApp.initialize(self, argv=argv) self.init_connection_file() class DummyConfigurable(connect.ConnectionFileMixin): def initialize(self): pass sample_info = dict(ip='1.2.3.4', transport='ipc', shell_port=1, hb_port=2, iopub_port=3, stdin_port=4, control_port=5, key=b'abc123', signature_scheme='hmac-md5', kernel_name='python' ) sample_info_kn = dict(ip='1.2.3.4', transport='ipc', shell_port=1, hb_port=2, iopub_port=3, stdin_port=4, control_port=5, key=b'abc123', signature_scheme='hmac-md5', kernel_name='test' ) def test_write_connection_file(): with TemporaryDirectory() as d: cf = os.path.join(d, 'kernel.json') connect.write_connection_file(cf, **sample_info) assert os.path.exists(cf) with open(cf, 'r') as f: info = json.load(f) info['key'] = str_to_bytes(info['key']) assert info == sample_info def test_load_connection_file_session(): """test load_connection_file() after """ session = Session() app = DummyConsoleApp(session=Session()) app.initialize(argv=[]) session = app.session with TemporaryDirectory() as d: cf = os.path.join(d, 'kernel.json') connect.write_connection_file(cf, **sample_info) app.connection_file = cf app.load_connection_file() assert session.key == sample_info['key'] assert session.signature_scheme == sample_info['signature_scheme'] def test_load_connection_file_session_with_kn(): """test load_connection_file() after """ session = Session() app = DummyConsoleApp(session=Session()) app.initialize(argv=[]) session = app.session with TemporaryDirectory() as d: cf = os.path.join(d, 'kernel.json') connect.write_connection_file(cf, **sample_info_kn) app.connection_file = cf app.load_connection_file() assert session.key == sample_info_kn['key'] assert session.signature_scheme == sample_info_kn['signature_scheme'] def test_app_load_connection_file(): """test `ipython console --existing` loads a connection file""" with TemporaryDirectory() as d: cf = os.path.join(d, 'kernel.json') connect.write_connection_file(cf, **sample_info) app = DummyConsoleApp(connection_file=cf) app.initialize(argv=[]) for attr, expected in sample_info.items(): if attr in ('key', 'signature_scheme'): continue value = getattr(app, attr) assert value == expected, "app.%s = %s != %s" % (attr, value, expected) def test_load_connection_info(): client = KernelClient() info = { 'control_port': 53702, 'hb_port': 53705, 'iopub_port': 53703, 'ip': '0.0.0.0', 'key': 'secret', 'shell_port': 53700, 'signature_scheme': 'hmac-sha256', 'stdin_port': 53701, 'transport': 'tcp', } client.load_connection_info(info) assert client.control_port == info['control_port'] assert client.session.key.decode('ascii') == info['key'] assert client.ip == info['ip'] def test_find_connection_file(): with TemporaryDirectory() as d: cf = 'kernel.json' app = DummyConsoleApp(runtime_dir=d, connection_file=cf) app.initialize() security_dir = app.runtime_dir profile_cf = os.path.join(security_dir, cf) with open(profile_cf, 'w') as f: f.write("{}") for query in ( 'kernel.json', 'kern*', '*ernel*', 'k*', ): assert connect.find_connection_file(query, path=security_dir) == profile_cf def test_find_connection_file_local(): with TemporaryWorkingDirectory() as d: cf = 'test.json' abs_cf = os.path.abspath(cf) with open(cf, 'w') as f: f.write('{}') for query in ( 'test.json', 'test', abs_cf, os.path.join('.', 'test.json'), ): assert connect.find_connection_file(query, path=['.', jupyter_runtime_dir()]) == abs_cf def test_find_connection_file_relative(): with TemporaryWorkingDirectory() as d: cf = 'test.json' os.mkdir('subdir') cf = os.path.join('subdir', 'test.json') abs_cf = os.path.abspath(cf) with open(cf, 'w') as f: f.write('{}') for query in ( os.path.join('.', 'subdir', 'test.json'), os.path.join('subdir', 'test.json'), abs_cf, ): assert connect.find_connection_file(query, path=['.', jupyter_runtime_dir()]) == abs_cf def test_find_connection_file_abspath(): with TemporaryDirectory() as d: cf = 'absolute.json' abs_cf = os.path.abspath(cf) with open(cf, 'w') as f: f.write('{}') assert connect.find_connection_file(abs_cf, path=jupyter_runtime_dir()) == abs_cf os.remove(abs_cf) def test_mixin_record_random_ports(): with TemporaryDirectory() as d: dc = DummyConfigurable(data_dir=d, kernel_name='via-tcp', transport='tcp') dc.write_connection_file() assert dc._connection_file_written assert os.path.exists(dc.connection_file) assert dc._random_port_names == connect.port_names def test_mixin_cleanup_random_ports(): with TemporaryDirectory() as d: dc = DummyConfigurable(data_dir=d, kernel_name='via-tcp', transport='tcp') dc.write_connection_file() filename = dc.connection_file dc.cleanup_random_ports() assert not os.path.exists(filename) for name in dc._random_port_names: assert getattr(dc, name) == 0 def test_secure_write(): def fetch_win32_permissions(filename): '''Extracts file permissions on windows using icacls''' role_permissions = {} for index, line in enumerate(os.popen("icacls %s" % filename).read().splitlines()): if index == 0: line = line.split(filename)[-1].strip().lower() match = re.match(r'\s*([^:]+):\(([^\)]*)\)', line) if match: usergroup, permissions = match.groups() usergroup = usergroup.lower().split('\\')[-1] permissions = set(p.lower() for p in permissions.split(',')) role_permissions[usergroup] = permissions elif not line.strip(): break return role_permissions def check_user_only_permissions(fname): # Windows has it's own permissions ACL patterns if os.name == 'nt': import win32api username = win32api.GetUserName().lower() permissions = fetch_win32_permissions(fname) assert username in permissions assert permissions[username] == set(['r', 'w']) assert 'administrators' in permissions assert permissions['administrators'] == set(['f']) assert 'everyone' not in permissions assert len(permissions) == 2 else: mode = os.stat(fname).st_mode assert '0600' == oct(stat.S_IMODE(mode)).replace('0o', '0') directory = tempfile.mkdtemp() fname = os.path.join(directory, 'check_perms') try: with secure_write(fname) as f: f.write('test 1') check_user_only_permissions(fname) with open(fname, 'r') as f: assert f.read() == 'test 1' # Try changing file permissions ahead of time os.chmod(fname, 0o755) with secure_write(fname) as f: f.write('test 2') check_user_only_permissions(fname) with open(fname, 'r') as f: assert f.read() == 'test 2' finally: shutil.rmtree(directory)