Files
jupyterhub/multiuser_notebook/multiuser.py
2014-06-19 16:51:55 -07:00

356 lines
10 KiB
Python

#!/usr/bin/env python
import os
import re
import socket
import sys
import uuid
from subprocess import Popen
import json
import requests
import time
import tornado.httpserver
import tornado.ioloop
import tornado.options
from tornado.log import app_log
from tornado.escape import url_escape, xhtml_escape
from tornado.httputil import url_concat
from tornado.web import RequestHandler, Application
from tornado import web
from tornado.options import define, options
from IPython.utils.traitlets import HasTraits, Any, Unicode, Integer, Dict
from IPython.html import utils
from .headers import HeadersHandler
def random_port():
"""get a single random port"""
sock = socket.socket()
sock.bind(('', 0))
port = sock.getsockname()[1]
sock.close()
return port
auth_header_pat = re.compile(r'^token\s+([^\s]+)$')
here = os.path.dirname(__file__)
def token_authorized(method):
def check_token(self, *args, **kwargs):
auth_header = self.request.headers.get('Authorization', '')
match = auth_header_pat.match(auth_header)
if not match:
raise web.HTTPError(403)
token = match.group(1)
app_log.info("api token: %r", token)
session = self.user_manager.user_for_api_token(token)
if session is None:
raise web.HTTPError(403)
self.request_session = session
return method(self, *args, **kwargs)
return check_token
class UserSession(HasTraits):
env_prefix = Unicode('IPY_')
process = Any()
port = Integer()
user = Unicode()
cookie_secret = Unicode()
cookie_name = Unicode()
def _cookie_name_default(self):
return 'ipy-multiuser-%s' % self.user
multiuser_prefix = Unicode()
multiuser_api_url = Unicode()
url_prefix = Unicode()
def _url_prefix_default(self):
return '/user/%s/' % self.user
api_token = Unicode()
def _api_token_default(self):
return str(uuid.uuid4())
cookie_token = Unicode()
def _cookie_token_default(self):
return str(uuid.uuid4())
def _env_key(self, d, key, value):
d['%s%s' % (self.env_prefix, key)] = value
env = Dict()
def _env_default(self):
env = os.environ.copy()
self._env_key(env, 'COOKIE_SECRET', self.cookie_secret)
self._env_key(env, 'API_TOKEN', self.api_token)
return env
@property
def auth_data(self):
return dict(
user=self.user,
)
def start(self):
assert self.process is None or self.process.poll() is not None
cmd = [sys.executable, '-m', 'multiuser_notebook.singleuser',
'--user=%s' % self.user, '--port=%i' % self.port,
'--cookie-name=%s' % self.cookie_name,
'--multiuser-prefix=%s' % self.multiuser_prefix,
'--multiuser-api-url=%s' % self.multiuser_api_url,
'--base-url=%s' % self.url_prefix,
]
app_log.info("Spawning: %s" % cmd)
self.process = Popen(cmd, env=self.env)
def running(self):
if self.process is None:
return False
if self.process.poll() is not None:
self.process = None
return False
return True
def stop(self):
if self.process is None:
return
if self.process.poll() is None:
self.process.terminate()
self.process = None
class SingleUserManager(HasTraits):
users = Dict()
routes_t = Unicode('http://localhost:8000/api/routes{uri}')
single_user_t = Unicode('http://localhost:{port}')
def _wait_for_port(self, port, timeout=2):
tic = time.time()
while time.time() - tic < timeout:
try:
socket.create_connection(('localhost', port))
except socket.error:
time.sleep(0.1)
else:
break
def get_session(self, user, **kwargs):
if user not in self.users:
kwargs['user'] = user
self.users[user] = UserSession(**kwargs)
return self.users[user]
def spawn(self, user):
session = self.get_session(user)
if session.running():
app_log.warn("User session %s already running", user)
return
session.port = port = random_port()
session.start()
r = requests.post(
self.routes_t.format(uri=session.url_prefix),
data=json.dumps(dict(
target=self.single_user_t.format(port=port),
user=user,
)),
)
self._wait_for_port(port)
r.raise_for_status()
def user_for_api_token(self, token):
"""Get the user session object for a given API token"""
for session in self.users.values():
if session.api_token == token:
return session
def user_for_cookie_token(self, token):
"""Get the user session object for a given cookie token"""
for session in self.users.values():
if session.cookie_token == token:
return session
def shutdown(self, user):
assert user in self.users
session = self.users.pop(user)
session.stop()
r = requests.delete(self.routes_url,
data=json.dumps(user=user, port=session.port),
)
r.raise_for_status()
class BaseHandler(RequestHandler):
@property
def cookie_name(self):
return self.settings.get('cookie_name', 'cookie')
@property
def multiuser_url(self):
return self.settings.get('multiuser_url', '')
@property
def multiuser_prefix(self):
return self.settings.get('multiuser_prefix', '/multiuser/')
def get_current_user(self):
token = self.get_cookie(self.cookie_name, '')
if token:
session = self.user_manager.user_for_cookie_token(token)
if session:
return session.user
@property
def base_url(self):
return self.settings.setdefault('base_url', '/')
@property
def user_manager(self):
return self.settings['user_manager']
def clear_login_cookie(self):
self.clear_cookie(self.cookie_name)
def spawn_single_user(self, user):
session = self.user_manager.get_session(user,
cookie_secret=self.settings['cookie_secret'],
multiuser_api_url=self.settings['multiuser_api_url'],
multiuser_prefix=self.settings['multiuser_prefix'],
)
self.user_manager.spawn(user)
return session
class MainHandler(BaseHandler):
@web.authenticated
def get(self):
self.redirect("/user/%s/" % self.get_current_user())
class UserHandler(BaseHandler):
@web.authenticated
def get(self, user):
self.log.info("multi-user at single-user url: %s", user)
if self.get_current_user() == user:
self.spawn_single_user(user)
self.redirect('')
else:
self.clear_login_cookie()
self.redirect(url_concat(self.settings['login_url'], {
'next' : '/user/%s/' % user
}))
class AuthorizationsHandler(BaseHandler):
@token_authorized
def get(self, token):
app_log.info('cookie token: %r', token)
session = self.user_manager.user_for_cookie_token(token)
if session is None:
app_log.info('cookie tokens: %r',
{ user:s.cookie_token for user,s in self.user_manager.users.items() }
)
raise web.HTTPError(404)
self.write(json.dumps({
'user' : session.user,
}))
class LogoutHandler(BaseHandler):
def get(self):
self.clear_login_cookie()
self.write("logged out")
class LoginHandler(BaseHandler):
def _render(self, message=None, user=None):
self.render('login.html',
next=url_escape(self.get_argument('next', default='')),
user=user,
message=message,
)
def get(self):
if False and self.get_current_user():
self.redirect(self.get_argument('next', default='/'))
else:
user = self.get_argument('user', default='')
self._render(user=user)
def post(self):
user = self.get_argument('user', default='')
pwd = self.get_argument('password', default=u'')
next_url = self.get_argument('next', default='') or '/user/%s/' % user
if user and pwd == 'password':
if user not in self.user_manager.users:
session = self.spawn_single_user(user)
else:
session = self.user_manager.users[user]
cookie_token = session.cookie_token
self.set_cookie(session.cookie_name, cookie_token, path=session.url_prefix)
self.set_cookie(self.cookie_name, cookie_token, path=self.base_url)
else:
self._render(
message={'error': 'Invalid username or password'},
user=user,
)
return
self.redirect(next_url)
def main():
define("port", default=8001, help="run on the given port", type=int)
tornado.options.parse_command_line()
handlers = [
(r"/", MainHandler),
(r"/login", LoginHandler),
(r"/logout", LogoutHandler),
(r"/headers", HeadersHandler),
(r"/api/authorizations/([^/]+)", AuthorizationsHandler),
]
# add base_url to handlers
base_url = "/multiuser/"
for i, tup in enumerate(handlers):
lis = list(tup)
lis[0] = utils.url_path_join(base_url, tup[0])
handlers[i] = tuple(lis)
handlers.extend([
(r"/user/([^/]+)/?.*", UserHandler),
(r"/", web.RedirectHandler, {"url" : base_url}),
])
application = Application(handlers,
base_url=base_url,
user_manager=SingleUserManager(),
cookie_secret='super secret',
cookie_name='multiusertest',
multiuser_prefix=base_url,
multiuser_api_url=utils.url_path_join(
'http://localhost:%i' % options.port,
base_url,
'api',
),
login_url=utils.url_path_join(base_url, 'login'),
template_path=os.path.join(here, 'templates'),
)
http_server = tornado.httpserver.HTTPServer(application)
http_server.listen(options.port)
proxy = Popen(["node", os.path.join(here, 'js', 'main.js')])
try:
tornado.ioloop.IOLoop.instance().start()
finally:
proxy.terminate()
if __name__ == "__main__":
main()