From 154cc00a314a0e8e2f99096007c1a3ea49c7ec19 Mon Sep 17 00:00:00 2001
From: yuvipanda <yuvipanda@gmail.com>
Date: Tue, 3 Jul 2018 16:18:32 -0700
Subject: [PATCH] Add --admin params to bootstrap / installer to set up admin
 accts

This removes a big step requiring file changes from the quickstart
path. You can specify the admin username on the commandline.
---
 bootstrap/bootstrap.py | 21 +++++++++-----
 setup.py               |  3 +-
 tljh/installer.py      | 65 ++++++++++++++++++++++++++++++++++++------
 3 files changed, 73 insertions(+), 16 deletions(-)

diff --git a/bootstrap/bootstrap.py b/bootstrap/bootstrap.py
index c3e2beb..5c2299e 100644
--- a/bootstrap/bootstrap.py
+++ b/bootstrap/bootstrap.py
@@ -17,6 +17,7 @@ import urllib.request
 import contextlib
 import hashlib
 import tempfile
+import sys
 
 
 def md5_file(fname):
@@ -118,16 +119,22 @@ def main():
     else:
         print('Upgrading TLJH installer...')
 
-    pip_install(hub_prefix, [
-        os.environ.get('TLJH_BOOTSTRAP_PIP_SPEC', 'git+https://github.com/yuvipanda/the-littlest-jupyterhub.git')
-    ], editable=os.environ.get('TLJH_BOOTSTRAP_DEV', 'no') == 'yes')
+    is_dev = os.environ.get('TLJH_BOOTSTRAP_DEV', 'no') == 'yes'
+    tljh_repo_path = os.environ.get(
+        'TLJH_BOOTSTRAP_PIP_SPEC',
+        'git+https://github.com/yuvipanda/the-littlest-jupyterhub.git'
+    )
+
+    pip_install(hub_prefix, [tljh_repo_path], editable=is_dev)
 
     print('Starting TLJH installer...')
-    os.execl(
+    os.execv(
         os.path.join(hub_prefix, 'bin', 'python3'),
-        os.path.join(hub_prefix, 'bin', 'python3'),
-        '-m',
-        'tljh.installer'
+        [
+            os.path.join(hub_prefix, 'bin', 'python3'),
+            '-m',
+            'tljh.installer',
+        ] + sys.argv[1:]
     )
 
 
diff --git a/setup.py b/setup.py
index d6c42d6..269e37b 100644
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,7 @@ setup(
     packages=find_packages(),
     include_package_data=True,
     install_requires=[
-        'pyyaml==3.*'
+        'pyyaml==3.*',
+        'ruamel.yaml==0.15.*'
     ]
 )
diff --git a/tljh/installer.py b/tljh/installer.py
index da0db7d..b141781 100644
--- a/tljh/installer.py
+++ b/tljh/installer.py
@@ -4,6 +4,8 @@ import tljh.systemd as systemd
 import tljh.conda as conda
 from tljh import user
 import secrets
+import argparse
+from ruamel.yaml import YAML
 
 INSTALL_PREFIX = os.environ.get('TLJH_INSTALL_PREFIX', '/opt/tljh')
 HUB_ENV_PREFIX = os.path.join(INSTALL_PREFIX, 'hub')
@@ -11,10 +13,13 @@ USER_ENV_PREFIX = os.path.join(INSTALL_PREFIX, 'user')
 
 HERE = os.path.abspath(os.path.dirname(__file__))
 
+rt_yaml = YAML()
+
 
 def ensure_jupyterhub_service(prefix):
     """
-    Ensure JupyterHub & CHP Services are set up properly """
+    Ensure JupyterHub & CHP Services are set up properly
+    """
     with open(os.path.join(HERE, 'systemd-units', 'jupyterhub.service')) as f:
         hub_unit_template = f.read()
 
@@ -62,16 +67,14 @@ def ensure_jupyterhub_package(prefix):
     conda.ensure_pip_packages(prefix, [
         'jupyterhub-dummyauthenticator==0.3.1',
         'jupyterhub-systemdspawner==0.9.12',
-        'jupyterhub-firstuseauthenticator==0.9'
+        'jupyterhub-firstuseauthenticator==0.10'
     ])
 
 
-def main():
-    print("Setting up JupyterHub...")
-    ensure_jupyterhub_package(HUB_ENV_PREFIX)
-    ensure_jupyterhub_service(HUB_ENV_PREFIX)
-
-    print("Setting up system user groups...")
+def ensure_usergroups():
+    """
+    Sets up user groups & sudo rules
+    """
     user.ensure_group('jupyterhub-admins')
     user.ensure_group('jupyterhub-users')
 
@@ -84,6 +87,11 @@ def main():
         # `pip` is in the $PATH we set in jupyterhub_config.py to include the user conda env.
         f.write('Defaults exempt_group = jupyterhub-admins\n')
 
+
+def ensure_user_environment():
+    """
+    Set up user conda environment with required packages
+    """
     print("Setting up user environment...")
     conda.ensure_conda_env(USER_ENV_PREFIX)
     conda.ensure_conda_packages(USER_ENV_PREFIX, [
@@ -100,6 +108,47 @@ def main():
         'nteract-on-jupyter==1.8.1'
     ])
 
+
+def ensure_admins(admins):
+    """
+    Setup given list of users as admins.
+    """
+    if not admins:
+        return
+    print("Setting up admin users")
+    config_path = os.path.join(INSTALL_PREFIX, 'config.yaml')
+    if os.path.exists(config_path):
+        with open(config_path, 'r') as f:
+            config = rt_yaml.load(f)
+    else:
+        config = {}
+
+    config['users'] = config.get('users', {})
+    config['users']['admin'] = list(admins)
+
+    with open(config_path, 'w+') as f:
+        rt_yaml.dump(config, f)
+
+
+def main():
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument(
+        '--admin',
+        nargs='*',
+        help='List of usernames set to be admin'
+    )
+
+    args = argparser.parse_args()
+
+    ensure_admins(args.admin)
+
+    ensure_usergroups()
+    ensure_user_environment()
+
+    print("Setting up JupyterHub...")
+    ensure_jupyterhub_package(HUB_ENV_PREFIX)
+    ensure_jupyterhub_service(HUB_ENV_PREFIX)
+
     print("Done!")