Merge pull request #517 from minrk/load-tokens

allow pre-loading API tokens from config

jupyterhub/app.py

@@ -353,6 +353,13 @@ class JupyterHub(Application):
         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,37 @@ 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)
+                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)
+                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()
 
     @gen.coroutine
     def init_spawners(self):
@@ -1055,6 +1093,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()

jupyterhub/orm.py

@@ -303,11 +303,21 @@ 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:
+            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])
         orm_token = APIToken(user_id=self.id)
         orm_token.token = token
         db.add(orm_token)

jupyterhub/tests/test_app.py

@@ -5,7 +5,11 @@ 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
 
 def test_help_all():
     out = check_output([sys.executable, '-m', 'jupyterhub', '--help-all']).decode('utf8', 'replace')
@@ -48,3 +52,38 @@ 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
+        
+        # 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

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')