From 2cc5b2c9e13b06f4fa3ff61c443fa5f0e1005319 Mon Sep 17 00:00:00 2001
From: Micke Nordin <kano@sunet.se>
Date: Fri, 12 Jan 2024 14:04:13 +0100
Subject: [PATCH] Switch to custom class

---
 jupyter/base/values/values.yaml | 44 ++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/jupyter/base/values/values.yaml b/jupyter/base/values/values.yaml
index 5f2d6a4..b4313e4 100644
--- a/jupyter/base/values/values.yaml
+++ b/jupyter/base/values/values.yaml
@@ -12,17 +12,39 @@ hub:
   extraConfig:
     oauthCode: |
       from oauthenticator.generic import GenericOAuthenticator
-      c.JupyterHub.authenticator_class = GenericOAuthenticator
-      c.GenericOAuthenticator.client_id = os.environ['NEXTCLOUD_CLIENT_ID']
-      c.GenericOAuthenticator.client_secret = os.environ['NEXTCLOUD_CLIENT_SECRET']
-      c.GenericOAuthenticator.login_service = 'Sunet Drive'
-      c.GenericOAuthenticator.username_claim = lambda r: r.get('ocs', {}).get('data', {}).get('id')
-      c.GenericOAuthenticator.userdata_url = 'https://' + os.environ['NEXTCLOUD_HOST'] + '/ocs/v2.php/cloud/user?format=json'
-      c.GenericOAuthenticator.authorize_url = 'https://' + os.environ['NEXTCLOUD_HOST'] + '/index.php/apps/oauth2/authorize'
-      c.GenericOAuthenticator.token_url = 'https://' + os.environ['NEXTCLOUD_HOST'] + '/index.php/apps/oauth2/api/v1/token'
-      c.GenericOAuthenticator.oauth_callback_url = 'https://' + os.environ['JUPYTER_HOST'] + '/hub/oauth_callback'     
-      c.GenericOAuthenticator.allow_all = True
-      c.GenericOAuthenticator.enable_auth_state = True
+
+      def post_auth_hook(authenticator, handler, authentication):
+          user = authentication['auth_state']['oauth_user']['ocs']['data']['id']
+          auth_state = authentication['auth_state']
+          authenticator.user_dict[user] = auth_state
+          return authentication
+
+      class NextcloudOAuthenticator(GenericOAuthenticator):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+              self.user_dict = {}
+
+        def pre_spawn_start(self, user, spawner):
+              super().pre_spawn_start(user, spawner)
+              access_token = self.user_dict[user.name]['access_token']
+              # refresh_token = self.user_dict[user.name]['refresh_token']
+              spawner.environment['NEXTCLOUD_ACCESS_TOKEN'] = access_token
+
+      c.JupyterHub.authenticator_class = NextcloudOAuthenticator
+      c.NextcloudOAuthenticator.client_id = os.environ['NEXTCLOUD_CLIENT_ID']
+      c.NextcloudOAuthenticator.client_secret = os.environ['NEXTCLOUD_CLIENT_SECRET']
+      c.NextcloudOAuthenticator.login_service = 'Sunet Drive'
+      c.NextcloudOAuthenticator.username_claim = lambda r: r.get('ocs', {}).get('data', {}).get('id')
+      c.NextcloudOAuthenticator.userdata_url = 'https://' + os.environ['NEXTCLOUD_HOST'] + '/ocs/v2.php/cloud/user?format=json'
+      c.NextcloudOAuthenticator.authorize_url = 'https://' + os.environ['NEXTCLOUD_HOST'] + '/index.php/apps/oauth2/authorize'
+      c.NextcloudOAuthenticator.token_url = 'https://' + os.environ['NEXTCLOUD_HOST'] + '/index.php/apps/oauth2/api/v1/token'
+      c.NextcloudOAuthenticator.oauth_callback_url = 'https://' + os.environ['JUPYTER_HOST'] + '/hub/oauth_callback'
+      c.NextcloudOAuthenticator.allow_all = True
+      c.NextcloudOAuthenticator.refresh_pre_spawn = True
+      c.NextcloudOAuthenticator.enable_auth_state = True
+      c.NextcloudOAuthenticator.auth_refresh_age = 3600
+      c.NextcloudOAuthenticator.post_auth_hook = post_auth_hook
+
   extraEnv:
     NEXTCLOUD_HOST: sunet.drive.test.sunet.se
     JUPYTER_HOST: jupyter.drive.test.sunet.se