-
Notifications
You must be signed in to change notification settings - Fork 4.9k
/
secure_notebook.py
139 lines (104 loc) · 4.28 KB
/
secure_notebook.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
#!/usr/bin/env python
"""
script to automatically setup notebook over SSL.
Generate cert and keyfiles (rsa 1024) in ~/.ssh/, ask for a password, and add
the corresponding entries in the notebook json configuration file.
"""
import six
from notebook.auth import passwd
from traitlets.config.loader import JSONFileConfigLoader, ConfigFileNotFound
from jupyter_core.paths import jupyter_config_dir
from traitlets.config import Config
from contextlib import contextmanager
from OpenSSL import crypto
from os.path import exists, join
import io
import os
import json
import traceback
def create_self_signed_cert(cert_dir, keyfile, certfile):
"""
Create a self-signed `keyfile` and `certfile` in `cert_dir`
Abort if one of the keyfile of certfile exist.
"""
if exists(join(cert_dir, certfile)) or exists(join(cert_dir, keyfile)):
raise FileExistsError('{} or {} already exist in {}. Aborting.'.format(keyfile, certfile, cert_dir))
else:
# create a key pair
k = crypto.PKey()
k.generate_key(crypto.TYPE_RSA, 1024)
# create a self-signed cert
cert = crypto.X509()
cert.get_subject().C = "US"
cert.get_subject().ST = "Jupyter notebook self-signed certificate"
cert.get_subject().L = "Jupyter notebook self-signed certificate"
cert.get_subject().O = "Jupyter notebook self-signed certificate"
cert.get_subject().OU = "my organization"
cert.get_subject().CN = "Jupyter notebook self-signed certificate"
cert.set_serial_number(1000)
cert.gmtime_adj_notBefore(0)
cert.gmtime_adj_notAfter(365*24*60*60)
cert.set_issuer(cert.get_subject())
cert.set_pubkey(k)
cert.sign(k, 'sha256')
with io.open(join(cert_dir, certfile), "wt") as f:
f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert).decode('utf8'))
os.chmod(join(cert_dir, certfile), 0o600)
with io.open(join(cert_dir, keyfile), "wt") as f:
f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, k).decode('utf8'))
os.chmod(join(cert_dir, keyfile), 0o600)
@contextmanager
def persist_config(mode=0o600):
"""Context manager that can be use to modify a config object
On exit of the context manager, the config will be written back to disk,
by default with 600 permissions.
"""
loader = JSONFileConfigLoader('jupyter_notebook_config.json', jupyter_config_dir())
try:
config = loader.load_config()
except ConfigFileNotFound:
config = Config()
yield config
filepath = os.path.join(jupyter_config_dir(), 'jupyter_notebook_config.json')
with io.open(filepath, 'w') as f:
f.write(six.u(json.dumps(config, indent=2)))
try:
os.chmod(filepath, mode)
except Exception:
traceback.print_exc()
print("Something went wrong changing file permissions")
def set_password():
"""Ask user for password, store it in notebook json configuration file"""
print("First choose a password.")
hashedpw = passwd()
print("We will store your password encrypted in the notebook configuration file: ")
print(hashedpw)
with persist_config() as config:
config.NotebookApp.password = hashedpw
print('... done\n')
def set_certifs():
"""
Generate certificate to run notebook over ssl and set up the notebook config.
"""
print("Let's generate self-signed certificates to secure your connexion.")
print("where should the certificate live?")
location = input('path [~/.ssh]: ')
if not location.strip():
location = os.path.expanduser('~/.ssh')
keyfile = input('keyfile name [jupyter_server.key]: ')
if not keyfile.strip():
keyfile = 'jupyter_server.key'
certfile = input('certfile name [jupyter_server.crt]: ')
if not certfile.strip():
certfile = 'jupyter_server.crt'
create_self_signed_cert(location, keyfile, certfile)
fullkey = os.path.join(location, keyfile)
fullcrt = os.path.join(location, certfile)
with persist_config() as config:
config.NotebookApp.certfile = fullcrt
config.NotebookApp.keyfile = fullkey
print('done.\n')
if __name__ == '__main__':
print("This will guide you through the steps towards securing your notebook server.")
set_password()
set_certifs()