add RucioAuthenticator for the jhub deployment

vre-hub · Apr 25, 2024 · 09d8bf9 · 09d8bf9
1 parent b2397e5
commit 09d8bf9
Showing 1 changed file with 43 additions and 44 deletions.
diff --git a/infrastructure/cluster/flux/jhub-dev/jhub-dev-release.yaml b/infrastructure/cluster/flux/jhub-dev/jhub-dev-release.yaml
@@ -63,57 +63,56 @@ spec:
             - profile
             - email
 
-      # extraConfig:
+      extraConfig:
+        token-exchange: |
+          import pprint
+          import os
+          import warnings
+          import requests
+          from oauthenticator.generic import GenericOAuthenticator
 
-      #   token-exchange: |
-      #     import pprint
-      #     import os
-      #     import warnings
-      #     import requests
-      #     from oauthenticator.generic import GenericOAuthenticator
+          # custom authenticator to enable auth_state and get access token to set as env var for rucio extension
+          class RucioAuthenticator(GenericOAuthenticator):
+              def __init__(self, **kwargs):
+                  super().__init__(**kwargs)
+                  self.enable_auth_state = True
 
-      #     # custom authenticator to enable auth_state and get access token to set as env var for rucio extension
-      #     class RucioAuthenticator(GenericOAuthenticator):
-      #         def __init__(self, **kwargs):
-      #             super().__init__(**kwargs)
-      #             self.enable_auth_state = True
-
-      #         def exchange_token(self, token):
-      #             params = {
-      #                 'client_id': self.client_id,
-      #                 'client_secret': self.client_secret,
-      #                 'grant_type': 'urn:ietf:params:oauth:grant-type:token-exchange',
-      #                 'subject_token': token,
-      #                 'scope': 'openid profile',
-      #                 'audience': 'rucio'
-      #             }
-      #             response = requests.post(self.token_url, data=params)
-      #             rucio_token = response.json()['access_token']
-      #             return rucio_token
+              def exchange_token(self, token):
+                  params = {
+                      'client_id': self.client_id,
+                      'client_secret': self.client_secret,
+                      'grant_type': 'urn:ietf:params:oauth:grant-type:token-exchange',
+                      'subject_token': token,
+                      'scope': 'openid profile',
+                      'audience': 'rucio'
+                  }
+                  response = requests.post(self.token_url, data=params)
+                  rucio_token = response.json()['access_token']
+                  return rucio_token
                   
-      #         async def pre_spawn_start(self, user, spawner):
-      #             auth_state = await user.get_auth_state()
-      #             pprint.pprint(auth_state)
-      #             if not auth_state:
-      #                 # user has no auth state
-      #                 return
+              async def pre_spawn_start(self, user, spawner):
+                  auth_state = await user.get_auth_state()
+                  pprint.pprint(auth_state)
+                  if not auth_state:
+                      # user has no auth state
+                      return
                   
-      #             # define token environment variable from auth_state
-      #             spawner.environment['RUCIO_ACCESS_TOKEN'] = self.exchange_token(auth_state['access_token'])
-      #             spawner.environment['EOS_ACCESS_TOKEN'] = auth_state['access_token']
+                  # define token environment variable from auth_state
+                  spawner.environment['RUCIO_ACCESS_TOKEN'] = self.exchange_token(auth_state['access_token'])
+                  spawner.environment['EOS_ACCESS_TOKEN'] = auth_state['access_token']
           
-      #     # set the above authenticator as the default
-      #     c.JupyterHub.authenticator_class = RucioAuthenticator
+          # set the above authenticator as the default
+          c.JupyterHub.authenticator_class = RucioAuthenticator
 
-      #     # enable authentication state
-      #     c.GenericOAuthenticator.enable_auth_state = True
+          # enable authentication state
+          c.GenericOAuthenticator.enable_auth_state = True
 
-      #     if 'JUPYTERHUB_CRYPT_KEY' not in os.environ:
-      #         warnings.warn(
-      #             "Need JUPYTERHUB_CRYPT_KEY env for persistent auth_state.\n"
-      #             "    export JUPYTERHUB_CRYPT_KEY=$(openssl rand -hex 32)"
-      #         )
-      #         c.CryptKeeper.keys = [os.urandom(32)]
+          if 'JUPYTERHUB_CRYPT_KEY' not in os.environ:
+              warnings.warn(
+                  "Need JUPYTERHUB_CRYPT_KEY env for persistent auth_state.\n"
+                  "    export JUPYTERHUB_CRYPT_KEY=$(openssl rand -hex 32)"
+              )
+              c.CryptKeeper.keys = [os.urandom(32)]
         
     singleuser:
       defaultUrl: "/lab"