From fa4b66669313b0bb9cc60ab75f0d8dc7fc6eb0b3 Mon Sep 17 00:00:00 2001 From: Min RK Date: Thu, 14 Apr 2016 15:31:36 +0200 Subject: [PATCH 1/2] allow pre-loading API tokens from config This is the first small part of easing the pain of services, which is generating the API tokens, and used to require initializing the JupyterHub database. --- jupyterhub/app.py | 30 ++++++++++++++++++++++++++++++ jupyterhub/orm.py | 14 +++++++++++--- jupyterhub/tests/test_app.py | 28 ++++++++++++++++++++++++++++ jupyterhub/tests/test_orm.py | 10 ++++++++++ 4 files changed, 79 insertions(+), 3 deletions(-) diff --git a/jupyterhub/app.py b/jupyterhub/app.py index 2c975234..cdc8e13f 100644 --- a/jupyterhub/app.py +++ b/jupyterhub/app.py @@ -352,6 +352,13 @@ class JupyterHub(Application): cookie_secret_file = Unicode('jupyterhub_cookie_secret', help="""File in which to store the cookie secret.""" ).tag(config=True) + + api_tokens = Dict(Unicode(), + help="""Dict of token:username to be loaded into the database. + + Allows ahead-of-time generation of API tokens for use by services. + """ + ).tag(config=True) authenticator_class = Type(PAMAuthenticator, Authenticator, help="""Class for authenticating users. @@ -794,6 +801,28 @@ class JupyterHub(Application): # From this point on, any user changes should be done simultaneously # to the whitelist set and user db, unless the whitelist is empty (all users allowed). + def init_api_tokens(self): + """Load predefined API tokens (for services) into database""" + db = self.db + for token, username in self.api_tokens.items(): + username = self.authenticator.normalize_username(username) + if not self.authenticator.check_whitelist(username): + raise ValueError("Token username %r is not in whitelist" % username) + if not self.authenticator.validate_username(username): + raise ValueError("Token username %r is not valid" % username) + orm_token = orm.APIToken.find(db, token) + if orm_token is None: + user = orm.User.find(db, username) + if user is None: + self.log.debug("Adding user %r to database", username) + user = orm.User(name=username) + db.add(user) + db.commit() + self.log.info("Adding API token for %s", username) + user.new_api_token(token) + else: + self.log.debug("Not duplicating token %s", orm_token) + db.commit() @gen.coroutine def init_spawners(self): @@ -1055,6 +1084,7 @@ class JupyterHub(Application): self.init_hub() self.init_proxy() yield self.init_users() + self.init_api_tokens() self.init_tornado_settings() yield self.init_spawners() self.init_handlers() diff --git a/jupyterhub/orm.py b/jupyterhub/orm.py index a0e0cfb0..1e7b7390 100644 --- a/jupyterhub/orm.py +++ b/jupyterhub/orm.py @@ -303,11 +303,19 @@ class User(Base): name=self.name, ) - def new_api_token(self): - """Create a new API token""" + def new_api_token(self, token=None): + """Create a new API token + + If `token` is given, load that token. + """ assert self.id is not None db = inspect(self).session - token = new_token() + if token is None: + token = new_token() + else: + found = APIToken.find(db, token) + if found: + raise ValueError("Collision on token: %s..." % token[:4]) orm_token = APIToken(user_id=self.id) orm_token.token = token db.add(orm_token) diff --git a/jupyterhub/tests/test_app.py b/jupyterhub/tests/test_app.py index 280c2e4b..832b91ac 100644 --- a/jupyterhub/tests/test_app.py +++ b/jupyterhub/tests/test_app.py @@ -6,6 +6,7 @@ import sys from subprocess import check_output, Popen, PIPE from tempfile import NamedTemporaryFile, TemporaryDirectory from .mocking import MockHub +from .. import orm def test_help_all(): out = check_output([sys.executable, '-m', 'jupyterhub', '--help-all']).decode('utf8', 'replace') @@ -48,3 +49,30 @@ def test_generate_config(): assert cfg_file in out assert 'Spawner.cmd' in cfg_text assert 'Authenticator.whitelist' in cfg_text + +def test_init_tokens(): + with TemporaryDirectory() as td: + db_file = os.path.join(td, 'jupyterhub.sqlite') + tokens = { + 'super-secret-token': 'alyx', + 'also-super-secret': 'gordon', + 'boagasdfasdf': 'chell', + } + app = MockHub(db_file=db_file, api_tokens=tokens) + app.initialize([]) + db = app.db + for token, username in tokens.items(): + api_token = orm.APIToken.find(db, token) + assert api_token is not None + user = api_token.user + assert user.name == username + + # simulate second startup, reloading same tokens: + app = MockHub(db_file=db_file, api_tokens=tokens) + app.initialize([]) + db = app.db + for token, username in tokens.items(): + api_token = orm.APIToken.find(db, token) + assert api_token is not None + user = api_token.user + assert user.name == username diff --git a/jupyterhub/tests/test_orm.py b/jupyterhub/tests/test_orm.py index 983c4d12..1243f622 100644 --- a/jupyterhub/tests/test_orm.py +++ b/jupyterhub/tests/test_orm.py @@ -93,6 +93,16 @@ def test_tokens(db): found = orm.APIToken.find(db, 'something else') assert found is None + secret = 'super-secret-preload-token' + token = user.new_api_token(secret) + assert token == secret + assert len(user.api_tokens) == 3 + + # raise ValueError on collision + with pytest.raises(ValueError): + user.new_api_token(token) + assert len(user.api_tokens) == 3 + def test_spawn_fails(db, io_loop): orm_user = orm.User(name='aeofel') From 094ac451c7ceefd888a81dace5e07aa600ffc334 Mon Sep 17 00:00:00 2001 From: Min RK Date: Fri, 15 Apr 2016 12:42:52 +0200 Subject: [PATCH 2/2] Don't allow bad tokens to create tokens in the db --- jupyterhub/app.py | 11 ++++++++++- jupyterhub/orm.py | 2 ++ jupyterhub/tests/test_app.py | 11 +++++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/jupyterhub/app.py b/jupyterhub/app.py index cdc8e13f..c2abbab5 100644 --- a/jupyterhub/app.py +++ b/jupyterhub/app.py @@ -813,13 +813,22 @@ class JupyterHub(Application): orm_token = orm.APIToken.find(db, token) if orm_token is None: user = orm.User.find(db, username) + user_created = False if user is None: + user_created = True self.log.debug("Adding user %r to database", username) user = orm.User(name=username) db.add(user) db.commit() self.log.info("Adding API token for %s", username) - user.new_api_token(token) + try: + user.new_api_token(token) + except Exception: + if user_created: + # don't allow bad tokens to create users + db.delete(user) + db.commit() + raise else: self.log.debug("Not duplicating token %s", orm_token) db.commit() diff --git a/jupyterhub/orm.py b/jupyterhub/orm.py index 1e7b7390..2f780459 100644 --- a/jupyterhub/orm.py +++ b/jupyterhub/orm.py @@ -313,6 +313,8 @@ class User(Base): if token is None: token = new_token() else: + if len(token) < 8: + raise ValueError("Tokens must be at least 8 characters, got %r" % token) found = APIToken.find(db, token) if found: raise ValueError("Collision on token: %s..." % token[:4]) diff --git a/jupyterhub/tests/test_app.py b/jupyterhub/tests/test_app.py index 832b91ac..e0e9dcd1 100644 --- a/jupyterhub/tests/test_app.py +++ b/jupyterhub/tests/test_app.py @@ -5,6 +5,9 @@ import re import sys from subprocess import check_output, Popen, PIPE from tempfile import NamedTemporaryFile, TemporaryDirectory + +import pytest + from .mocking import MockHub from .. import orm @@ -50,6 +53,7 @@ def test_generate_config(): assert 'Spawner.cmd' in cfg_text assert 'Authenticator.whitelist' in cfg_text + def test_init_tokens(): with TemporaryDirectory() as td: db_file = os.path.join(td, 'jupyterhub.sqlite') @@ -76,3 +80,10 @@ def test_init_tokens(): assert api_token is not None user = api_token.user assert user.name == username + + # don't allow failed token insertion to create users: + tokens['short'] = 'gman' + app = MockHub(db_file=db_file, api_tokens=tokens) + # with pytest.raises(ValueError): + app.initialize([]) + assert orm.User.find(app.db, 'gman') is None