From e5e790be890c9188ba6fc408d50213149bda3c29 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Thu, 12 Sep 2024 12:04:47 +0200 Subject: [PATCH 01/16] check token expiration --- chart/f7t4jhub/files/jupyterhub-config.py | 38 +++++++++++++++++++---- firecrestspawner/spawner.py | 16 +++++++--- 2 files changed, 44 insertions(+), 10 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index 925ddab..a38fd23 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -64,9 +64,26 @@ def get_access_token(self): return self._access_token -class GenericOAuthenticatorCSCS(GenericOAuthenticator): +class AuthenticatorCSCS(GenericOAuthenticator): + _min_token_validity = 60 + + async def authenticate(self, handler, data=None): + """Extended to add the token expiration time to the + authentication state""" + + auth_state = await super().authenticate(handler, data) + + self.log.debug("[authenticate] Authenticating") + + token_response = auth_state['auth_state']['token_response'] + + auth_state['auth_state']['access_token_expiration_ts'] = ( + time.time() + token_response['expires_in'] - self._min_token_validity + ) + + return auth_state + async def refresh_user(self, user, handler=None): - # self.log.info('Refreshing auth state') auth_state = await user.get_auth_state() params = { @@ -80,13 +97,20 @@ async def refresh_user(self, user, handler=None): "Content-Type": "application/x-www-form-urlencoded" } + if time.time() <= auth_state["access_token_expiration_ts"]: + self.log.debug(f"[refresh_user] Reusing access token for {user.name} {auth_state['access_token'][-10:]}") + return True + response = requests.post(self.token_url, data=params, headers=headers) + + self.log.debug(f"[refresh_user] Refreshing access token for {user.name} {response.json()['access_token'][-10:]}") + if response.status_code != 200: - self.log.debug(f"[refresh_user] Request to KeyCloak: {response.status_code}") + self.log.info(f"[refresh_user] Request to KeyCloak: {response.status_code}") try: - self.log.debug(f"[refresh_user] Request to KeyCloak: {response.json()}") + self.log.info(f"[refresh_user] Request to KeyCloak: {response.json()}") except json.JSONDecodeError: - self.log.debug(f"[refresh_user] Request to KeyCloak: no json output") + self.log.info(f"[refresh_user] Request to KeyCloak: no json output") return False @@ -95,6 +119,8 @@ async def refresh_user(self, user, handler=None): auth_state['token_response'].update(token_response) auth_state['access_token'] = token_response['access_token'] auth_state['refresh_token'] = token_response['refresh_token'] + access_token_expiration_ts = token_response['expires_in'] - self._min_token_validity + auth_state['access_token_expiration_ts'] = time.time() + access_token_expiration_ts return { 'name': auth_state['oauth_user']['preferred_username'], @@ -115,7 +141,7 @@ async def refresh_user(self, user, handler=None): c.Authenticator.enable_auth_state = True c.CryptKeeper.keys = gen_hex_string("/home/juhu/hex_strings_crypt.txt") -c.JupyterHub.authenticator_class = GenericOAuthenticatorCSCS +c.JupyterHub.authenticator_class = AuthenticatorCSCS c.GenericOAuthenticator.client_id = os.environ.get('KC_CLIENT_ID', '') c.GenericOAuthenticator.client_secret = os.environ.get('KC_CLIENT_SECRET', '') c.GenericOAuthenticator.oauth_callback_url = "{{ .Values.config.auth.oauthCallbackUrl }}" diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 40c7830..3a12722 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -15,6 +15,7 @@ import hostlist import httpx import base64 +import time import firecrest as f7t from enum import Enum from jinja2 import Template @@ -188,11 +189,18 @@ def cmd_formatted_for_batch(self): return ' '.join(self.cmd) async def get_firecrest_client(self): - # TODO: Check if token is expired - # auth_state = await self.user.get_auth_state() + auth_state = await self.user.get_auth_state() - auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) # noqa E501 - access_token = auth_state_refreshed['auth_state']['access_token'] + if time.time() <= auth_state["access_token_expiration_ts"]: + self.log.debug(f"[get_firecrest_client] Reusing access token for {self.user.name}") + access_token = auth_state["access_token"] + else: + try: + auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) # noqa E501 + access_token = auth_state_refreshed['auth_state']['access_token'] + self.log.debug(f"[get_firecrest_client] Refreshing access token for {self.user.name}") + except TypeError: + access_token = auth_state["access_token"] client = f7t.AsyncFirecrest( firecrest_url=self.firecrest_url, From 3e75b847bfec16bc790801f6075b23c00989eb34 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Thu, 12 Sep 2024 15:06:18 +0200 Subject: [PATCH 02/16] remove login last chars of token --- chart/f7t4jhub/files/jupyterhub-config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index a38fd23..9ba2f9f 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -98,12 +98,12 @@ async def refresh_user(self, user, handler=None): } if time.time() <= auth_state["access_token_expiration_ts"]: - self.log.debug(f"[refresh_user] Reusing access token for {user.name} {auth_state['access_token'][-10:]}") + self.log.debug(f"[refresh_user] Reusing access token for {user.name}") return True response = requests.post(self.token_url, data=params, headers=headers) - self.log.debug(f"[refresh_user] Refreshing access token for {user.name} {response.json()['access_token'][-10:]}") + self.log.debug(f"[refresh_user] Refreshing access token for {user.name}") if response.status_code != 200: self.log.info(f"[refresh_user] Request to KeyCloak: {response.status_code}") From a86c5484d42d569f8d82320354adafe281a5904e Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Thu, 12 Sep 2024 15:20:28 +0200 Subject: [PATCH 03/16] fix unit tests --- tests/test_spawner.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_spawner.py b/tests/test_spawner.py index acefef6..628fdf3 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -29,6 +29,25 @@ async def refresh_user(self, user, handler=None): return {"auth_state": auth_state} +async def get_firecrest_client(spawner): + auth_state_refreshed = await spawner.user.authenticator.refresh_user(spawner.user) # noqa E501 + access_token = auth_state_refreshed['auth_state']['access_token'] + + client = firecrest.AsyncFirecrest( + firecrest_url=spawner.firecrest_url, + authorization=FirecrestAccessTokenAuth(access_token) + ) + + return client + + +# FIXME: Setup the auth state in the unit tests +# Since the auth state is not setup for the unit tests, +# the spawner's get_firecrest_client method will fail +# when trying to get a key from a none `auth_state` +SlurmSpawner.get_firecrest_client = get_firecrest_client + + def new_spawner(db, spawner_class=SlurmSpawner, **kwargs): user = db.query(orm.User).first() hub = Hub() From 3492c0441a3ce0f908e32567ce70e67079c9fc69 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Tue, 8 Oct 2024 17:09:41 +0200 Subject: [PATCH 04/16] add aux client for polling --- chart/f7t4jhub/files/jupyterhub-config.py | 69 +++++++++++++++-------- firecrestspawner/spawner.py | 63 +++++++++++++++------ 2 files changed, 93 insertions(+), 39 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index 9ba2f9f..0fb77e2 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -1,4 +1,5 @@ import json +import grp import os import requests import secrets @@ -65,7 +66,7 @@ def get_access_token(self): class AuthenticatorCSCS(GenericOAuthenticator): - _min_token_validity = 60 + _min_token_validity = 30 async def authenticate(self, handler, data=None): """Extended to add the token expiration time to the @@ -77,6 +78,9 @@ async def authenticate(self, handler, data=None): token_response = auth_state['auth_state']['token_response'] + with open('refresh_token.txt', 'w') as file: + print(token_response["refresh_token"], file=file) + auth_state['auth_state']['access_token_expiration_ts'] = ( time.time() + token_response['expires_in'] - self._min_token_validity ) @@ -86,6 +90,10 @@ async def authenticate(self, handler, data=None): async def refresh_user(self, user, handler=None): auth_state = await user.get_auth_state() + if time.time() <= auth_state["access_token_expiration_ts"]: + self.log.debug(f"[refresh_user] Reusing access token for {user.name}") + return True + params = { 'grant_type': 'refresh_token', 'client_id': self.client_id, @@ -97,14 +105,13 @@ async def refresh_user(self, user, handler=None): "Content-Type": "application/x-www-form-urlencoded" } - if time.time() <= auth_state["access_token_expiration_ts"]: - self.log.debug(f"[refresh_user] Reusing access token for {user.name}") - return True - response = requests.post(self.token_url, data=params, headers=headers) self.log.debug(f"[refresh_user] Refreshing access token for {user.name}") + token_response = response.json() + auth_state['token_response'].update(token_response) + if response.status_code != 200: self.log.info(f"[refresh_user] Request to KeyCloak: {response.status_code}") try: @@ -114,14 +121,17 @@ async def refresh_user(self, user, handler=None): return False - token_response = response.json() - - auth_state['token_response'].update(token_response) auth_state['access_token'] = token_response['access_token'] auth_state['refresh_token'] = token_response['refresh_token'] access_token_expiration_ts = token_response['expires_in'] - self._min_token_validity auth_state['access_token_expiration_ts'] = time.time() + access_token_expiration_ts + user._auth_refreshed = time.monotonic() + await user.save_auth_state(auth_state) + + self.log.debug(f"[refresh_user] refresh_token {handler} for {user.name} " + f"{token_response['refresh_token'][-10:]} {token_response['refresh_expires_in']}") + return { 'name': auth_state['oauth_user']['preferred_username'], 'auth_state': auth_state @@ -135,23 +145,25 @@ async def refresh_user(self, user, handler=None): c.JupyterHub.admin_access = False c.Authenticator.allow_all = True -c.Authenticator.refresh_pre_spawn = True +# c.Authenticator.refresh_pre_spawn = True c.Authenticator.auth_refresh_age = 250 c.Authenticator.enable_auth_state = True c.CryptKeeper.keys = gen_hex_string("/home/juhu/hex_strings_crypt.txt") c.JupyterHub.authenticator_class = AuthenticatorCSCS -c.GenericOAuthenticator.client_id = os.environ.get('KC_CLIENT_ID', '') -c.GenericOAuthenticator.client_secret = os.environ.get('KC_CLIENT_SECRET', '') -c.GenericOAuthenticator.oauth_callback_url = "{{ .Values.config.auth.oauthCallbackUrl }}" -c.GenericOAuthenticator.authorize_url = "{{ .Values.config.auth.authorizeUrl }}" -c.GenericOAuthenticator.token_url = "{{ .Values.config.auth.tokenUrl }}" -c.GenericOAuthenticator.userdata_url = "{{ .Values.config.auth.userDataUrl }}" -c.GenericOAuthenticator.login_service = "{{ .Values.config.auth.loginService }}" -c.GenericOAuthenticator.username_key = "{{ .Values.config.auth.userNameKey }}" -c.GenericOAuthenticator.userdata_params = {{ .Values.config.auth.userDataParams }} -c.GenericOAuthenticator.scope = {{ .Values.config.auth.scope }} +c.AuthenticatorCSCS.client_id = os.environ.get('KC_CLIENT_ID', '') +c.AuthenticatorCSCS.client_secret = os.environ.get('KC_CLIENT_SECRET', '') +c.AuthenticatorCSCS.oauth_callback_url = "{{ .Values.config.auth.oauthCallbackUrl }}" +c.AuthenticatorCSCS.authorize_url = "{{ .Values.config.auth.authorizeUrl }}" +c.AuthenticatorCSCS.token_url = "{{ .Values.config.auth.tokenUrl }}" +c.AuthenticatorCSCS.userdata_url = "{{ .Values.config.auth.userDataUrl }}" +c.AuthenticatorCSCS.login_service = "{{ .Values.config.auth.loginService }}" +c.AuthenticatorCSCS.username_key = "{{ .Values.config.auth.userNameKey }}" +c.AuthenticatorCSCS.userdata_params = {{ .Values.config.auth.userDataParams }} +c.AuthenticatorCSCS.scope = {{ .Values.config.auth.scope }} + +# c.JupyterHub.cookie_max_age_days = 0.01 c.JupyterHub.default_url = '{{ .Values.config.hubDefaultUrl }}' @@ -162,10 +174,10 @@ async def refresh_user(self, user, handler=None): c.Spawner.req_host = '{{ .Values.config.spawner.host }}' c.Spawner.node_name_template = '{{ .Values.config.spawner.nodeNameTemplate }}' c.Spawner.req_partition = '{{ .Values.config.spawner.partition }}' -c.Spawner.req_account = '{{ .Values.config.spawner.account }}' c.Spawner.req_constraint = '{{ .Values.config.spawner.constraint }}' c.Spawner.req_srun = '{{ .Values.config.spawner.srun }}' c.Spawner.batch_script = """#!/bin/bash + #SBATCH --job-name={{ .Values.config.spawner.jobName }} #SBATCH --chdir={{`{{homedir}}`}} #SBATCH --get-user-env=L @@ -177,8 +189,16 @@ async def refresh_user(self, user, handler=None): {% if gres %}#SBATCH --gres={{`{{gres}}`}}{% endif %} {% if nprocs %}#SBATCH --cpus-per-task={{`{{nprocs}}`}}{% endif %} {% if nnodes %}#SBATCH --nodes={{`{{nnodes[0]}}`}}{% endif %} -{% if reservation%}#SBATCH --reservation={{`{{reservation[0]}}`}}{% endif %} -{% if constraint %}#SBATCH --constraint={{`{{constraint[0]}}`}}{% endif %} +{% if reservation is string %} +#SBATCH --reservation={{`{{reservation}}`}} +{% else %} +#SBATCH --reservation={{`{{reservation[0]}}`}} +{% endif %} +{% if constraint is string %} +#SBATCH --constraint={{`{{constraint}}`}} +{% else %} +#SBATCH --constraint={{`{{constraint[0]}}`}} +{% endif %} {% if options %}#SBATCH {{`{{options}}`}}{% endif %} # Activate a virtual environment, load modules, etc @@ -226,4 +246,9 @@ async def refresh_user(self, user, handler=None): # This should be set to the URL which the hub uses to connect to the proxy’s API. c.ConfigurableHTTPProxy.api_url = 'http://{{ .Release.Name }}-proxy-svc:{{ .Values.network.apiPort }}' + +os.environ['FIRECREST_CLIENT_ID_AUX'] = "..." +os.environ['FIRECREST_CLIENT_SECRET_AUX'] = "..." +os.environ['AUTH_TOKEN_URL_AUX'] = "..." + {{ .Values.config.extraConfig }} diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 3a12722..38f19d5 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -189,18 +189,9 @@ def cmd_formatted_for_batch(self): return ' '.join(self.cmd) async def get_firecrest_client(self): + await self.user.authenticator.refresh_user(self.user) auth_state = await self.user.get_auth_state() - - if time.time() <= auth_state["access_token_expiration_ts"]: - self.log.debug(f"[get_firecrest_client] Reusing access token for {self.user.name}") - access_token = auth_state["access_token"] - else: - try: - auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) # noqa E501 - access_token = auth_state_refreshed['auth_state']['access_token'] - self.log.debug(f"[get_firecrest_client] Refreshing access token for {self.user.name}") - except TypeError: - access_token = auth_state["access_token"] + access_token = auth_state["access_token"] client = f7t.AsyncFirecrest( firecrest_url=self.firecrest_url, @@ -219,6 +210,30 @@ async def get_firecrest_client(self): client.timeout = 30 return client + async def get_firecrest_aux_client(self): + client_id = os.environ['FIRECREST_CLIENT_ID_AUX'] + client_secret = os.environ['FIRECREST_CLIENT_SECRET_AUX'] + token_url = os.environ['AUTH_TOKEN_URL_AUX'] + + keycloak = f7t.ClientCredentialsAuth( + client_id, client_secret, token_url + ) + + client = f7t.AsyncFirecrest(firecrest_url=self.firecrest_url, authorization=keycloak) + + client.time_between_calls = { + "compute": 0, + "reservations": 0, + "status": 0, + "storage": 0, + "tasks": 0, + "utilities": 0, + } + + client.timeout = 30 + return client + + async def firecrest_poll(self): """Helper function to poll jobs. @@ -226,12 +241,21 @@ async def firecrest_poll(self): its database which could make the result of ``client.poll`` to be an empty list """ - client = await self.get_firecrest_client() - poll_result = [] + + self.log.info(f"\n\n\n[poll]\n\n\n") + + auth_state = await self.user.get_auth_state() + auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) + + if auth_state_refreshed == False: + client = await self.get_firecrest_aux_client() + else: + client = await self.get_firecrest_client() + poll_result = [] while poll_result == []: poll_result = await client.poll(self.host, [self.job_id]) - print(poll_result) + await asyncio.sleep(1) return poll_result @@ -258,12 +282,16 @@ async def submit_batch_script(self): for v in ("JUPYTERHUB_OAUTH_ACCESS_SCOPES", "JUPYTERHUB_OAUTH_SCOPES"): job_env[v] = base64.b64encode(job_env[v].encode()).decode("utf-8") + self.host = subvars['host'] + + client = await self.get_firecrest_client() + groups = await client.groups(self.host) + subvars['account'] = groups['group']['name'] + script = await self._get_batch_script(**subvars) self.log.info('Spawner submitting job using firecREST') self.log.info('Spawner submitted script:\n' + script) - self.host = subvars['host'] - try: client = await self.get_firecrest_client() self.log.info('firecREST: Submitting job') @@ -464,7 +492,7 @@ async def progress(self): return else: await yield_({ - "message": "Unknown status...", + "message": "Waiting for job status...", }) await asyncio.sleep(1) @@ -550,6 +578,7 @@ async def state_gethost(self): class SlurmSpawner(FirecRESTSpawnerRegexStates): firecrest_url = os.environ['FIRECREST_URL'] + token_url = os.environ['AUTH_TOKEN_URL'] batch_script = Unicode("""#!/bin/bash ##SBATCH --output={{homedir}}/jupyterhub_slurmspawner_%j.log From 477b50264ab1fc77016558caa92bcbbf2fc3d5cd Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Wed, 9 Oct 2024 09:50:46 +0200 Subject: [PATCH 05/16] remove test log --- firecrestspawner/spawner.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 38f19d5..f704b4e 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -242,8 +242,6 @@ async def firecrest_poll(self): to be an empty list """ - self.log.info(f"\n\n\n[poll]\n\n\n") - auth_state = await self.user.get_auth_state() auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) From ced086bf3984d93b7ec20711f9de944df5664097 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Wed, 9 Oct 2024 10:47:19 +0200 Subject: [PATCH 06/16] fix unit tests --- firecrestspawner/spawner.py | 1 - requirements.txt | 2 +- tests/fc_handlers.py | 58 ++++++++++++++++++++++++------------- tests/test_spawner.py | 5 ++++ 4 files changed, 44 insertions(+), 22 deletions(-) diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index f704b4e..198ad49 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -576,7 +576,6 @@ async def state_gethost(self): class SlurmSpawner(FirecRESTSpawnerRegexStates): firecrest_url = os.environ['FIRECREST_URL'] - token_url = os.environ['AUTH_TOKEN_URL'] batch_script = Unicode("""#!/bin/bash ##SBATCH --output={{homedir}}/jupyterhub_slurmspawner_%j.log diff --git a/requirements.txt b/requirements.txt index fe6b854..1a4deab 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ jupyterhub==4.1.6 -pyfirecrest==2.1.0 +pyfirecrest==2.6.0 SQLAlchemy==1.4.52 oauthenticator==16.3.1 python-hostlist==1.23.0 diff --git a/tests/fc_handlers.py b/tests/fc_handlers.py index 6aa7fa0..bfc5293 100644 --- a/tests/fc_handlers.py +++ b/tests/fc_handlers.py @@ -4,6 +4,44 @@ from werkzeug.wrappers import Response +def whoami_handler(request): + if request.headers["Authorization"] != "Bearer VALID_TOKEN": + return Response( + json.dumps({"message": "Bad token; invalid JSON"}), + status=401, + content_type="application/json", + ) + + if request.headers["X-Machine-Name"] != "cluster1": + return Response( + json.dumps( + {"description": "Error on whoami operation", "error": "Machine does not exist"} + ), + status=400, + headers={"X-Machine-Does-Not-Exist": "Machine does not exist"}, + content_type="application/json", + ) + + groups = request.args.get("groups", False) + if groups: + ret = { + "description": "User information", + "output": { + "group": {"id": "1000", "name": "group1"}, + "groups": [{"id": "1000", "name": "group1"}, {"id": "1001", "name": "group2"}], + "user": {"id": "10000", "name": "test_user"}, + } + } + else: + ret = {"description": "Success on whoami operation.", "output": "username"} + + return Response( + json.dumps(ret), + status=200, + content_type="application/json", + ) + + def tasks_handler(request): taskid = request.args.get("tasks") if taskid == "acct_352_id": @@ -389,23 +427,3 @@ def cancel_handler(request): return Response( json.dumps(ret), status=status_code, content_type="application/json" ) - - -@pytest.fixture -def fc_server(httpserver): - httpserver.expect_request( - "/compute/jobs/upload", method="POST" - ).respond_with_handler(submit_upload_handler) - httpserver.expect_request( - re.compile("^/status/systems.*"), method="GET" - ).respond_with_handler(systems_handler) - httpserver.expect_request("/tasks", method="GET").respond_with_handler( - tasks_handler - ) - httpserver.expect_request("/compute/acct", method="GET").respond_with_handler( - sacct_handler - ) - httpserver.expect_request( - re.compile("^/compute/jobs.*"), method="DELETE" - ).respond_with_handler(cancel_handler) - return httpserver diff --git a/tests/test_spawner.py b/tests/test_spawner.py index 628fdf3..93eac81 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -11,6 +11,7 @@ systems_handler, sacct_handler, cancel_handler, + whoami_handler ) from jupyterhub.tests.conftest import db from jupyterhub.user import User @@ -88,6 +89,10 @@ def fc_server(httpserver): re.compile("^/compute/jobs.*"), method="DELETE" ).respond_with_handler(cancel_handler) + httpserver.expect_request( + "/utilities/whoami", method="GET" + ).respond_with_handler(whoami_handler) + return httpserver From 2e5b5464c28b774daeedd80a4c1f439f8634fc97 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Thu, 10 Oct 2024 16:17:12 +0200 Subject: [PATCH 07/16] move sa to vault --- chart/f7t4jhub/files/jupyterhub-config.py | 6 +----- chart/f7t4jhub/templates/deployment-hub.yaml | 15 +++++++++++++++ chart/f7t4jhub/templates/external-secret.yaml | 8 ++++++++ chart/f7t4jhub/templates/secret.yaml | 1 + chart/values.yaml | 11 ++++++++--- firecrestspawner/spawner.py | 18 ++++++++++++------ 6 files changed, 45 insertions(+), 14 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index 0fb77e2..c9daa5f 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -171,6 +171,7 @@ async def refresh_user(self, user, handler=None): c.JupyterHub.hub_connect_ip = socket.gethostbyname(hostname) c.JupyterHub.spawner_class = 'firecrestspawner.spawner.SlurmSpawner' +c.Spawner.enable_aux_fc_client = {{ .Values.serviceAccount.enabled | toJson | replace "true" "True" | replace "false" "False" }} c.Spawner.req_host = '{{ .Values.config.spawner.host }}' c.Spawner.node_name_template = '{{ .Values.config.spawner.nodeNameTemplate }}' c.Spawner.req_partition = '{{ .Values.config.spawner.partition }}' @@ -246,9 +247,4 @@ async def refresh_user(self, user, handler=None): # This should be set to the URL which the hub uses to connect to the proxy’s API. c.ConfigurableHTTPProxy.api_url = 'http://{{ .Release.Name }}-proxy-svc:{{ .Values.network.apiPort }}' - -os.environ['FIRECREST_CLIENT_ID_AUX'] = "..." -os.environ['FIRECREST_CLIENT_SECRET_AUX'] = "..." -os.environ['AUTH_TOKEN_URL_AUX'] = "..." - {{ .Values.config.extraConfig }} diff --git a/chart/f7t4jhub/templates/deployment-hub.yaml b/chart/f7t4jhub/templates/deployment-hub.yaml index 60218fb..7c0a2ec 100644 --- a/chart/f7t4jhub/templates/deployment-hub.yaml +++ b/chart/f7t4jhub/templates/deployment-hub.yaml @@ -63,6 +63,11 @@ spec: secretKeyRef: name: {{ .Release.Name }}-secret key: authTokenUrl + - name: SA_AUTH_TOKEN_URL + valueFrom: + secretKeyRef: + name: {{ .Release.Name }}-secret + key: serviceAccountAuthTokenUrl {{- if .Values.vault.keycloak.enabled }} - name: KC_CLIENT_ID valueFrom: @@ -74,6 +79,16 @@ spec: secretKeyRef: name: {{ .Release.Name }}-common-secrets key: kc_client_secret + - name: SA_CLIENT_ID + valueFrom: + secretKeyRef: + name: {{ .Release.Name }}-common-secrets + key: sa_client_id + - name: SA_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: {{ .Release.Name }}-common-secrets + key: sa_client_secret {{- end }} {{- if .Values.vault.configProxyAuthToken.enabled }} - name: CONFIGPROXY_AUTH_TOKEN diff --git a/chart/f7t4jhub/templates/external-secret.yaml b/chart/f7t4jhub/templates/external-secret.yaml index 9382f8e..dc7d568 100644 --- a/chart/f7t4jhub/templates/external-secret.yaml +++ b/chart/f7t4jhub/templates/external-secret.yaml @@ -24,4 +24,12 @@ spec: remoteRef: key: {{ .Values.vault.configProxyAuthToken.secretPath }} property: config_proxy_auth_token + - secretKey: sa_client_id + remoteRef: + key: {{ .Values.vault.keycloak.secretPath }} + property: sa_client_id + - secretKey: sa_client_secret + remoteRef: + key: {{ .Values.vault.keycloak.secretPath }} + property: sa_client_secret {{- end }} diff --git a/chart/f7t4jhub/templates/secret.yaml b/chart/f7t4jhub/templates/secret.yaml index 5e005f1..e3a8eb9 100644 --- a/chart/f7t4jhub/templates/secret.yaml +++ b/chart/f7t4jhub/templates/secret.yaml @@ -13,6 +13,7 @@ type: Opaque stringData: firecrestUrl: {{ .Values.setup.firecrestUrl }} authTokenUrl: {{ .Values.setup.authTokenUrl}} + serviceAccountAuthTokenUrl: {{ .Values.serviceAccount.authTokenUrl }} {{- if not .Values.vault.configProxyAuthToken.enabled }} configProxyAuthToken: {{ $token }} {{- end }} diff --git a/chart/values.yaml b/chart/values.yaml index 20b80c9..6a2e39c 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -76,6 +76,14 @@ f7t4jhub: # Secret path in Vault (replace with your own secret path) secretPath: 'secret/path/proxy' + # service account for polling jobs + serviceAccount: + # Enable or disable service account for polling jobs + enabled: true + + # URL to obtain an auth token from your identity provider (replace with the SA's token URL) + authTokenUrl: 'https://auth.cscs.ch/auth/realms/firecrest-clients/protocol/openid-connect/token' + metricbeat: # Enable or disable annotations for metric beat monitoring enabled: false @@ -140,9 +148,6 @@ f7t4jhub: # Partition for the job scheduler (customize as needed) partition: 'slurm_partition' - # Account for the job scheduler (customize as needed) - account: 'slurm_account' - # Constraint for the job scheduler (customize as needed) constraint: 'slurm_constraint' diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 198ad49..a1ba406 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -21,7 +21,7 @@ from jinja2 import Template from jupyterhub.spawner import Spawner from traitlets import ( - Any, Integer, Unicode, Float, default + Any, Bool, Integer, Unicode, Float, default ) from time import sleep @@ -168,6 +168,12 @@ def _req_username_default(self): "needs specification." ).tag(config=True) + enable_aux_fc_client = Bool( + 300, + help="If positive, use an auxiliary client to poll when client " + "credentials are expired." + ).tag(config=True) + # Raw output of job submission command unless overridden job_id = Unicode() @@ -211,9 +217,9 @@ async def get_firecrest_client(self): return client async def get_firecrest_aux_client(self): - client_id = os.environ['FIRECREST_CLIENT_ID_AUX'] - client_secret = os.environ['FIRECREST_CLIENT_SECRET_AUX'] - token_url = os.environ['AUTH_TOKEN_URL_AUX'] + client_id = os.environ['SA_CLIENT_ID'] + client_secret = os.environ['SA_CLIENT_SECRET'] + token_url = os.environ['SA_AUTH_TOKEN_URL'] keycloak = f7t.ClientCredentialsAuth( client_id, client_secret, token_url @@ -242,10 +248,10 @@ async def firecrest_poll(self): to be an empty list """ - auth_state = await self.user.get_auth_state() + # auth_state = await self.user.get_auth_state() auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) - if auth_state_refreshed == False: + if self.enable_aux_fc_client and auth_state_refreshed == False: client = await self.get_firecrest_aux_client() else: client = await self.get_firecrest_client() From 91a53d1e654168e19a97009d94a937843db46e16 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Mon, 14 Oct 2024 22:33:40 +0200 Subject: [PATCH 08/16] working refreshless --- chart/f7t4jhub/files/jupyterhub-config.py | 54 +++++++++++++---------- firecrestspawner/spawner.py | 15 ++++--- 2 files changed, 38 insertions(+), 31 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index c9daa5f..63c42db 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -68,31 +68,31 @@ def get_access_token(self): class AuthenticatorCSCS(GenericOAuthenticator): _min_token_validity = 30 - async def authenticate(self, handler, data=None): - """Extended to add the token expiration time to the - authentication state""" + # async def authenticate(self, handler, data=None): + # """Extended to add the token expiration time to the + # authentication state""" - auth_state = await super().authenticate(handler, data) + # auth_state = await super().authenticate(handler, data) - self.log.debug("[authenticate] Authenticating") + # self.log.debug("[authenticate] Authenticating") - token_response = auth_state['auth_state']['token_response'] + # token_response = auth_state['auth_state']['token_response'] - with open('refresh_token.txt', 'w') as file: - print(token_response["refresh_token"], file=file) + # with open('refresh_token.txt', 'w') as file: + # print(token_response["refresh_token"], file=file) - auth_state['auth_state']['access_token_expiration_ts'] = ( - time.time() + token_response['expires_in'] - self._min_token_validity - ) + # auth_state['auth_state']['access_token_expiration_ts'] = ( + # time.time() + token_response['expires_in'] - self._min_token_validity + # ) - return auth_state + # return auth_state - async def refresh_user(self, user, handler=None): + async def _refresh_user(self, user, handler=None): auth_state = await user.get_auth_state() - if time.time() <= auth_state["access_token_expiration_ts"]: - self.log.debug(f"[refresh_user] Reusing access token for {user.name}") - return True + # if time.time() <= auth_state["access_token_expiration_ts"]: + # self.log.debug(f"[refresh_user] Reusing access token for {user.name}") + # return True params = { 'grant_type': 'refresh_token', @@ -113,21 +113,21 @@ async def refresh_user(self, user, handler=None): auth_state['token_response'].update(token_response) if response.status_code != 200: - self.log.info(f"[refresh_user] Request to KeyCloak: {response.status_code}") - try: - self.log.info(f"[refresh_user] Request to KeyCloak: {response.json()}") - except json.JSONDecodeError: - self.log.info(f"[refresh_user] Request to KeyCloak: no json output") + # self.log.info(f"[refresh_user] Request to KeyCloak: {response.status_code}") + # try: + # self.log.info(f"[refresh_user] Request to KeyCloak: {response.json()}") + # except json.JSONDecodeError: + # self.log.info(f"[refresh_user] Request to KeyCloak: no json output") - return False + return True auth_state['access_token'] = token_response['access_token'] auth_state['refresh_token'] = token_response['refresh_token'] access_token_expiration_ts = token_response['expires_in'] - self._min_token_validity auth_state['access_token_expiration_ts'] = time.time() + access_token_expiration_ts - user._auth_refreshed = time.monotonic() - await user.save_auth_state(auth_state) + # user._auth_refreshed = time.monotonic() + # await user.save_auth_state(auth_state) self.log.debug(f"[refresh_user] refresh_token {handler} for {user.name} " f"{token_response['refresh_token'][-10:]} {token_response['refresh_expires_in']}") @@ -247,4 +247,10 @@ async def refresh_user(self, user, handler=None): # This should be set to the URL which the hub uses to connect to the proxy’s API. c.ConfigurableHTTPProxy.api_url = 'http://{{ .Release.Name }}-proxy-svc:{{ .Values.network.apiPort }}' + +def pre_spawn_hook(spawner): + spawner.user.authenticator._refresh_user(spawner.user) + +c.Spawner.pre_spawn_hook = pre_spawn_hook + {{ .Values.config.extraConfig }} diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index a1ba406..29a1e28 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -195,8 +195,8 @@ def cmd_formatted_for_batch(self): return ' '.join(self.cmd) async def get_firecrest_client(self): - await self.user.authenticator.refresh_user(self.user) - auth_state = await self.user.get_auth_state() + auth_state_refreshed = await self.user.authenticator._refresh_user(self.user) + auth_state = auth_state_refreshed['auth_state'] #await self.user.get_auth_state() access_token = auth_state["access_token"] client = f7t.AsyncFirecrest( @@ -249,12 +249,13 @@ async def firecrest_poll(self): """ # auth_state = await self.user.get_auth_state() - auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) + # auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) - if self.enable_aux_fc_client and auth_state_refreshed == False: - client = await self.get_firecrest_aux_client() - else: - client = await self.get_firecrest_client() + # if self.enable_aux_fc_client and auth_state_refreshed == False: + # client = await self.get_firecrest_aux_client() + # else: + # client = await self.get_firecrest_client() + client = await self.get_firecrest_aux_client() poll_result = [] while poll_result == []: From 99acf69ca24b08f908613beaf4c22fbfc1b40fce Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Fri, 18 Oct 2024 17:12:18 +0200 Subject: [PATCH 09/16] cleaning code --- chart/f7t4jhub/files/jupyterhub-config.py | 53 ++--------------------- firecrestspawner/spawner.py | 33 +++++++++----- 2 files changed, 26 insertions(+), 60 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index 63c42db..cb9a56d 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -1,3 +1,4 @@ +import asyncio import json import grp import os @@ -66,34 +67,9 @@ def get_access_token(self): class AuthenticatorCSCS(GenericOAuthenticator): - _min_token_validity = 30 - - # async def authenticate(self, handler, data=None): - # """Extended to add the token expiration time to the - # authentication state""" - - # auth_state = await super().authenticate(handler, data) - - # self.log.debug("[authenticate] Authenticating") - - # token_response = auth_state['auth_state']['token_response'] - - # with open('refresh_token.txt', 'w') as file: - # print(token_response["refresh_token"], file=file) - - # auth_state['auth_state']['access_token_expiration_ts'] = ( - # time.time() + token_response['expires_in'] - self._min_token_validity - # ) - - # return auth_state - async def _refresh_user(self, user, handler=None): auth_state = await user.get_auth_state() - # if time.time() <= auth_state["access_token_expiration_ts"]: - # self.log.debug(f"[refresh_user] Reusing access token for {user.name}") - # return True - params = { 'grant_type': 'refresh_token', 'client_id': self.client_id, @@ -107,30 +83,15 @@ async def _refresh_user(self, user, handler=None): response = requests.post(self.token_url, data=params, headers=headers) - self.log.debug(f"[refresh_user] Refreshing access token for {user.name}") - token_response = response.json() + auth_state['token_response'].update(token_response) if response.status_code != 200: - # self.log.info(f"[refresh_user] Request to KeyCloak: {response.status_code}") - # try: - # self.log.info(f"[refresh_user] Request to KeyCloak: {response.json()}") - # except json.JSONDecodeError: - # self.log.info(f"[refresh_user] Request to KeyCloak: no json output") - - return True + return False auth_state['access_token'] = token_response['access_token'] auth_state['refresh_token'] = token_response['refresh_token'] - access_token_expiration_ts = token_response['expires_in'] - self._min_token_validity - auth_state['access_token_expiration_ts'] = time.time() + access_token_expiration_ts - - # user._auth_refreshed = time.monotonic() - # await user.save_auth_state(auth_state) - - self.log.debug(f"[refresh_user] refresh_token {handler} for {user.name} " - f"{token_response['refresh_token'][-10:]} {token_response['refresh_expires_in']}") return { 'name': auth_state['oauth_user']['preferred_username'], @@ -163,8 +124,6 @@ async def _refresh_user(self, user, handler=None): c.AuthenticatorCSCS.userdata_params = {{ .Values.config.auth.userDataParams }} c.AuthenticatorCSCS.scope = {{ .Values.config.auth.scope }} -# c.JupyterHub.cookie_max_age_days = 0.01 - c.JupyterHub.default_url = '{{ .Values.config.hubDefaultUrl }}' hostname = socket.gethostname() @@ -247,10 +206,4 @@ async def _refresh_user(self, user, handler=None): # This should be set to the URL which the hub uses to connect to the proxy’s API. c.ConfigurableHTTPProxy.api_url = 'http://{{ .Release.Name }}-proxy-svc:{{ .Values.network.apiPort }}' - -def pre_spawn_hook(spawner): - spawner.user.authenticator._refresh_user(spawner.user) - -c.Spawner.pre_spawn_hook = pre_spawn_hook - {{ .Values.config.extraConfig }} diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 29a1e28..e654ebc 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -5,6 +5,7 @@ # Modified and redistributed under the terms of the BSD 3-Clause License. # See the accompanying LICENSE file for more details. +import inspect import asyncio from async_generator import async_generator, yield_ import pwd @@ -73,6 +74,9 @@ class FirecRESTSpawnerBase(Spawner): state_gethost """ + # define since the begining since this is used in the home.html template + access_token_is_valid = False + # override default since batch systems typically need longer start_timeout = Integer(300).tag(config=True) @@ -196,7 +200,7 @@ def cmd_formatted_for_batch(self): async def get_firecrest_client(self): auth_state_refreshed = await self.user.authenticator._refresh_user(self.user) - auth_state = auth_state_refreshed['auth_state'] #await self.user.get_auth_state() + auth_state = auth_state_refreshed['auth_state'] access_token = auth_state["access_token"] client = f7t.AsyncFirecrest( @@ -216,7 +220,7 @@ async def get_firecrest_client(self): client.timeout = 30 return client - async def get_firecrest_aux_client(self): + async def get_firecrest_client_service_account(self): client_id = os.environ['SA_CLIENT_ID'] client_secret = os.environ['SA_CLIENT_SECRET'] token_url = os.environ['SA_AUTH_TOKEN_URL'] @@ -247,15 +251,11 @@ async def firecrest_poll(self): its database which could make the result of ``client.poll`` to be an empty list """ - - # auth_state = await self.user.get_auth_state() - # auth_state_refreshed = await self.user.authenticator.refresh_user(self.user) - # if self.enable_aux_fc_client and auth_state_refreshed == False: - # client = await self.get_firecrest_aux_client() - # else: - # client = await self.get_firecrest_client() - client = await self.get_firecrest_aux_client() + if self.enable_aux_fc_client: + client = await self.get_firecrest_client_service_account() + else: + client = await self.get_firecrest_client() poll_result = [] while poll_result == []: @@ -406,6 +406,19 @@ def state_gethost(self): async def poll(self): """Poll the process""" + + # check if the refresh token is valid when + # accessing /hub/home and set `self.access_token_is_valid` + # for the template in `home.html` + stack = inspect.stack() + caller_frame = stack[3] + if 'self' in caller_frame.frame.f_locals: + class_name = caller_frame.frame.f_locals['self'].__class__.__name__ + + if class_name == "HomeHandler": + auth_state_refreshed = await self.user.authenticator._refresh_user(self.user) + self.access_token_is_valid = bool(auth_state_refreshed) + status = await self.query_job_status() if status in (JobStatus.PENDING, JobStatus.RUNNING, JobStatus.UNKNOWN): return None From 5705275804ba9b3e934847f019807143e1ffe72e Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Fri, 18 Oct 2024 17:17:51 +0200 Subject: [PATCH 10/16] fix unit tests --- firecrestspawner/spawner.py | 2 +- tests/test_spawner.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index e654ebc..f7a259e 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -173,7 +173,7 @@ def _req_username_default(self): ).tag(config=True) enable_aux_fc_client = Bool( - 300, + True, help="If positive, use an auxiliary client to poll when client " "credentials are expired." ).tag(config=True) diff --git a/tests/test_spawner.py b/tests/test_spawner.py index 93eac81..ace5505 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -63,6 +63,7 @@ def new_spawner(db, spawner_class=SlurmSpawner, **kwargs): req_host="cluster1", port=testport, node_name_template="{}.cluster1.ch", + enable_aux_fc_client=False ) return _spawner From 31097c31281496ca163131b5af266678b346b216 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Thu, 24 Oct 2024 14:21:23 +0200 Subject: [PATCH 11/16] fix values --- chart/values.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/chart/values.yaml b/chart/values.yaml index 6a2e39c..4e1a11c 100644 --- a/chart/values.yaml +++ b/chart/values.yaml @@ -82,7 +82,7 @@ f7t4jhub: enabled: true # URL to obtain an auth token from your identity provider (replace with the SA's token URL) - authTokenUrl: 'https://auth.cscs.ch/auth/realms/firecrest-clients/protocol/openid-connect/token' + authTokenUrl: 'https://auth-sa.example.com/auth/realms/yourrealm/protocol/openid-connect/token' metricbeat: # Enable or disable annotations for metric beat monitoring From 44e8aa85b5d77472b6c952c0191210428998a754 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Thu, 24 Oct 2024 14:29:18 +0200 Subject: [PATCH 12/16] fix comment --- firecrestspawner/spawner.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index f7a259e..9ac7321 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -245,18 +245,16 @@ async def get_firecrest_client_service_account(self): async def firecrest_poll(self): - """Helper function to poll jobs. - - This is needed in cases that the scheduler is slow updating - its database which could make the result of ``client.poll`` - to be an empty list - """ + """Helper function to poll jobs.""" if self.enable_aux_fc_client: client = await self.get_firecrest_client_service_account() else: client = await self.get_firecrest_client() + # This is needed in case the scheduler is slow updating + # its database which could make the result of ``client.poll`` + # to be an empty list poll_result = [] while poll_result == []: poll_result = await client.poll(self.host, [self.job_id]) From b2b9f36350b49993b830f7b97778d5c0b37aed3e Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Tue, 5 Nov 2024 10:51:15 +0100 Subject: [PATCH 13/16] update spawner --- chart/f7t4jhub/files/jupyterhub-config.py | 66 ++++---------------- firecrestspawner/spawner.py | 76 ++++++++++++++++++----- requirements-tests.txt | 2 +- tests/context.py | 3 +- tests/test_spawner.py | 16 ++++- 5 files changed, 87 insertions(+), 76 deletions(-) diff --git a/chart/f7t4jhub/files/jupyterhub-config.py b/chart/f7t4jhub/files/jupyterhub-config.py index cb9a56d..121d311 100644 --- a/chart/f7t4jhub/files/jupyterhub-config.py +++ b/chart/f7t4jhub/files/jupyterhub-config.py @@ -55,50 +55,6 @@ async def get_node_ip_from_output(spawner): time.sleep(2) -class FirecrestAccessTokenAuth: - - _access_token: str = None - - def __init__(self, access_token): - self._access_token = access_token - - def get_access_token(self): - return self._access_token - - -class AuthenticatorCSCS(GenericOAuthenticator): - async def _refresh_user(self, user, handler=None): - auth_state = await user.get_auth_state() - - params = { - 'grant_type': 'refresh_token', - 'client_id': self.client_id, - 'client_secret': self.client_secret, - 'refresh_token': auth_state['refresh_token'] - } - - headers = { - "Content-Type": "application/x-www-form-urlencoded" - } - - response = requests.post(self.token_url, data=params, headers=headers) - - token_response = response.json() - - auth_state['token_response'].update(token_response) - - if response.status_code != 200: - return False - - auth_state['access_token'] = token_response['access_token'] - auth_state['refresh_token'] = token_response['refresh_token'] - - return { - 'name': auth_state['oauth_user']['preferred_username'], - 'auth_state': auth_state - } - - c = get_config() @@ -112,17 +68,17 @@ async def _refresh_user(self, user, handler=None): c.Authenticator.enable_auth_state = True c.CryptKeeper.keys = gen_hex_string("/home/juhu/hex_strings_crypt.txt") -c.JupyterHub.authenticator_class = AuthenticatorCSCS -c.AuthenticatorCSCS.client_id = os.environ.get('KC_CLIENT_ID', '') -c.AuthenticatorCSCS.client_secret = os.environ.get('KC_CLIENT_SECRET', '') -c.AuthenticatorCSCS.oauth_callback_url = "{{ .Values.config.auth.oauthCallbackUrl }}" -c.AuthenticatorCSCS.authorize_url = "{{ .Values.config.auth.authorizeUrl }}" -c.AuthenticatorCSCS.token_url = "{{ .Values.config.auth.tokenUrl }}" -c.AuthenticatorCSCS.userdata_url = "{{ .Values.config.auth.userDataUrl }}" -c.AuthenticatorCSCS.login_service = "{{ .Values.config.auth.loginService }}" -c.AuthenticatorCSCS.username_key = "{{ .Values.config.auth.userNameKey }}" -c.AuthenticatorCSCS.userdata_params = {{ .Values.config.auth.userDataParams }} -c.AuthenticatorCSCS.scope = {{ .Values.config.auth.scope }} +c.JupyterHub.authenticator_class = GenericOAuthenticator +c.GenericOAuthenticator.client_id = os.environ.get('KC_CLIENT_ID', '') +c.GenericOAuthenticator.client_secret = os.environ.get('KC_CLIENT_SECRET', '') +c.GenericOAuthenticator.oauth_callback_url = "{{ .Values.config.auth.oauthCallbackUrl }}" +c.GenericOAuthenticator.authorize_url = "{{ .Values.config.auth.authorizeUrl }}" +c.GenericOAuthenticator.token_url = "{{ .Values.config.auth.tokenUrl }}" +c.GenericOAuthenticator.userdata_url = "{{ .Values.config.auth.userDataUrl }}" +c.GenericOAuthenticator.login_service = "{{ .Values.config.auth.loginService }}" +c.GenericOAuthenticator.username_key = "{{ .Values.config.auth.userNameKey }}" +c.GenericOAuthenticator.userdata_params = {{ .Values.config.auth.userDataParams }} +c.GenericOAuthenticator.scope = {{ .Values.config.auth.scope }} c.JupyterHub.default_url = '{{ .Values.config.hubDefaultUrl }}' diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 9ac7321..8ca3614 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -18,6 +18,7 @@ import base64 import time import firecrest as f7t +import requests from enum import Enum from jinja2 import Template from jupyterhub.spawner import Spawner @@ -48,18 +49,43 @@ class JobStatus(Enum): UNKNOWN = 3 -class FirecrestAccessTokenAuth: - """Utility class to provide an object with the - `get_access_token()` attribute needed by PyFirecREST's - authenticator""" +class AuthorizationCodeFlowAuth: - _access_token: str = None - - def __init__(self, access_token): - self._access_token = access_token + def __init__( + self, + client_id: str, + client_secret: str, + refresh_token: str, + token_url: str, + ): + self.client_id = client_id + self.client_secret = client_secret + self.token_url = token_url + self.refresh_token = refresh_token def get_access_token(self): - return self._access_token + params = { + 'grant_type': 'refresh_token', + 'client_id': self.client_id, + 'client_secret': self.client_secret, + 'refresh_token': self.refresh_token + } + headers = { + "Content-Type": "application/x-www-form-urlencoded" + } + response = requests.post( + self.token_url, + data=params, + headers=headers + ) + + if response.status_code != 200: + return False + + json_response = response.json() + self.refresh_token = json_response["refresh_token"] + + return json_response['access_token'] class FirecRESTSpawnerBase(Spawner): @@ -199,13 +225,21 @@ def cmd_formatted_for_batch(self): return ' '.join(self.cmd) async def get_firecrest_client(self): - auth_state_refreshed = await self.user.authenticator._refresh_user(self.user) - auth_state = auth_state_refreshed['auth_state'] - access_token = auth_state["access_token"] + # auth_state_refreshed = await self.user.authenticator.refresh_access_token(self.user) + # auth_state = auth_state_refreshed['auth_state'] + # access_token = auth_state["access_token"] + auth_state = await self.user.get_auth_state() + + auth = AuthorizationCodeFlowAuth( + client_id = self.user.authenticator.client_id, + client_secret = self.user.authenticator.client_secret, + refresh_token = auth_state["refresh_token"], + token_url = self.user.authenticator.token_url, + ) client = f7t.AsyncFirecrest( firecrest_url=self.firecrest_url, - authorization=FirecrestAccessTokenAuth(access_token) + authorization=auth ) client.time_between_calls = { @@ -225,11 +259,11 @@ async def get_firecrest_client_service_account(self): client_secret = os.environ['SA_CLIENT_SECRET'] token_url = os.environ['SA_AUTH_TOKEN_URL'] - keycloak = f7t.ClientCredentialsAuth( + auth = f7t.ClientCredentialsAuth( client_id, client_secret, token_url ) - client = f7t.AsyncFirecrest(firecrest_url=self.firecrest_url, authorization=keycloak) + client = f7t.AsyncFirecrest(firecrest_url=self.firecrest_url, authorization=auth) client.time_between_calls = { "compute": 0, @@ -414,8 +448,16 @@ async def poll(self): class_name = caller_frame.frame.f_locals['self'].__class__.__name__ if class_name == "HomeHandler": - auth_state_refreshed = await self.user.authenticator._refresh_user(self.user) - self.access_token_is_valid = bool(auth_state_refreshed) + # auth_state_refreshed = await self.user.authenticator.refresh_access_token(self.user) + auth_state = await self.user.get_auth_state() + + auth = AuthorizationCodeFlowAuth( + client_id = self.user.authenticator.client_id, + client_secret = self.user.authenticator.client_secret, + refresh_token = auth_state["refresh_token"], + token_url = self.user.authenticator.token_url, + ) + self.access_token_is_valid = bool(auth.get_access_token()) status = await self.query_job_status() if status in (JobStatus.PENDING, JobStatus.RUNNING, JobStatus.UNKNOWN): diff --git a/requirements-tests.txt b/requirements-tests.txt index 592c0e2..6881eb4 100644 --- a/requirements-tests.txt +++ b/requirements-tests.txt @@ -1,4 +1,4 @@ pytest==7.4.2 pytest_httpserver==1.0.10 -werkzeug==3.0.3 +werkzeug==3.0.6 pytest-asyncio==0.23.7 diff --git a/tests/context.py b/tests/context.py index 0fb3584..db6479a 100644 --- a/tests/context.py +++ b/tests/context.py @@ -8,6 +8,5 @@ from firecrestspawner.spawner import ( SlurmSpawner, - format_template, - FirecrestAccessTokenAuth, + format_template ) diff --git a/tests/test_spawner.py b/tests/test_spawner.py index ace5505..f5ede9b 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -4,7 +4,7 @@ import getpass import pytest from werkzeug.wrappers import Response -from context import FirecrestAccessTokenAuth, SlurmSpawner, format_template +from context import SlurmSpawner, format_template from fc_handlers import ( tasks_handler, submit_upload_handler, @@ -30,6 +30,20 @@ async def refresh_user(self, user, handler=None): return {"auth_state": auth_state} +class FirecrestAccessTokenAuth: + """Utility class to provide an object with the + `get_access_token()` attribute needed by PyFirecREST's + authenticator""" + + _access_token: str = None + + def __init__(self, access_token): + self._access_token = access_token + + def get_access_token(self): + return self._access_token + + async def get_firecrest_client(spawner): auth_state_refreshed = await spawner.user.authenticator.refresh_user(spawner.user) # noqa E501 access_token = auth_state_refreshed['auth_state']['access_token'] From e80798ab5aad09de0ab6dc442e8d2833ae5cd4fa Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Tue, 5 Nov 2024 10:52:37 +0100 Subject: [PATCH 14/16] update pyfirecrest version in dockerfile --- dockerfiles/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dockerfiles/Dockerfile b/dockerfiles/Dockerfile index c59141a..281da5a 100644 --- a/dockerfiles/Dockerfile +++ b/dockerfiles/Dockerfile @@ -15,7 +15,7 @@ RUN apt-get update -qq \ RUN python3 -m venv /opt/jhub-env RUN . /opt/jhub-env/bin/activate && \ - pip install --no-cache jupyterhub==4.1.6 pyfirecrest==2.1.0 SQLAlchemy==1.4.52 oauthenticator==16.3.1 python-hostlist==1.23.0 + pip install --no-cache jupyterhub==4.1.6 pyfirecrest==2.6.0 SQLAlchemy==1.4.52 oauthenticator==16.3.1 python-hostlist==1.23.0 COPY . firecrestspawner From 7d6dd275060ba135fb936749c5ca635e4851f2d0 Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Tue, 5 Nov 2024 11:22:01 +0100 Subject: [PATCH 15/16] update docs --- docs/source/deployment.rst | 4 ++-- firecrestspawner/spawner.py | 35 ++++++++++++++++++++--------------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/docs/source/deployment.rst b/docs/source/deployment.rst index 21f686c..0b9faa7 100644 --- a/docs/source/deployment.rst +++ b/docs/source/deployment.rst @@ -104,9 +104,9 @@ They can be accessed in our kubernetes deployment via a set of secrets: - A `SecretStore `_, which interacts with the ``vault-approle-secret`` secret. -- An `ExternalSecret `_ which interacts with the ``SecretStore`` allowing the deployment to access the Keycloak client's IDs and secrets. +- An `ExternalSecret `_ which interacts with the ``SecretStore`` allowing the deployment to access the client's IDs and secrets. -- An optional `ExternalSecret `_ that's used for to store credentials for a custom container registry. That's currently not in use. +- An optional `ExternalSecret to access credentials for a custom container registry `_. That's currently not in use. The section of the chart related to Vault is optional and can be disabled in the ``values.yaml``. diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 89f94a4..7613167 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -50,6 +50,14 @@ class JobStatus(Enum): class AuthorizationCodeFlowAuth: + """ + Authorization Code Flow class + + :param client_id: name of the client as registered in the authorization server + :param client_secret: secret associated to the client + :param refresh_token: refresh token for the SSO session + :param token_url: URL of the token request in the authorization server + """ def __init__( self, @@ -64,6 +72,8 @@ def __init__( self.refresh_token = refresh_token def get_access_token(self): + """Returns an access token to be used for accessing resources. + """ params = { 'grant_type': 'refresh_token', 'client_id': self.client_id, @@ -89,16 +99,7 @@ def get_access_token(self): class FirecRESTSpawnerBase(Spawner): - """Base class for spawners using PyFirecrest to submit jobs. - - At minimum, subclasses should provide reasonable defaults for the traits: - batch_script - - and must provide implementations for the methods: - state_ispending - state_isrunning - state_gethost - """ + """Base class for spawners using PyFirecrest to submit jobs""" # define since the begining since this is used in the home.html template access_token_is_valid = False @@ -210,8 +211,8 @@ def _req_username_default(self): # Will get the raw output of the job status command unless overridden job_status = Unicode() - # Prepare substitution variables for templates using req_xyz traits def get_req_subvars(self): + """Prepare substitution variables for templates using req_xyz traits.""" reqlist = [t for t in self.trait_names() if t.startswith('req_')] subvars = {} for t in reqlist: @@ -220,14 +221,12 @@ def get_req_subvars(self): return subvars def cmd_formatted_for_batch(self): - """The command which is substituted inside of the batch script""" + """The command which is substituted inside of the batch script.""" # return ' '.join([self.batchspawner_singleuser_cmd] + self.cmd + self.get_args()) # noqa E501 return ' '.join(self.cmd) async def get_firecrest_client(self): - # auth_state_refreshed = await self.user.authenticator.refresh_access_token(self.user) - # auth_state = auth_state_refreshed['auth_state'] - # access_token = auth_state["access_token"] + """Returns a firecrest client that uses the Authorization Code Flow method""" auth_state = await self.user.get_auth_state() auth = AuthorizationCodeFlowAuth( @@ -255,6 +254,7 @@ async def get_firecrest_client(self): return client async def get_firecrest_client_service_account(self): + """Returns a firecrest client that uses the Client Credentials Authorization method""" client_id = os.environ['SA_CLIENT_ID'] client_secret = os.environ['SA_CLIENT_SECRET'] token_url = os.environ['SA_AUTH_TOKEN_URL'] @@ -302,6 +302,10 @@ async def _get_batch_script(self, **subvars): return format_template(self.batch_script, **subvars) async def submit_batch_script(self): + """Submits the batch script that starts the notebook server job + + It's called by ``spawner.start``. + """ subvars = self.get_req_subvars() # `subvars['cmd']` is what is run _inside_ the batch script, # put into the template. @@ -639,6 +643,7 @@ async def state_gethost(self): class SlurmSpawner(FirecRESTSpawnerRegexStates): + """Implementation of the FirecRESTSpawner for Slurm""" firecrest_url = os.environ['FIRECREST_URL'] batch_script = Unicode("""#!/bin/bash From 04266be24b4a0f7956e8705b5af489fb7886a0ef Mon Sep 17 00:00:00 2001 From: Rafael Sarmiento Date: Tue, 5 Nov 2024 12:29:59 +0100 Subject: [PATCH 16/16] remove comments --- firecrestspawner/spawner.py | 1 - 1 file changed, 1 deletion(-) diff --git a/firecrestspawner/spawner.py b/firecrestspawner/spawner.py index 7613167..7346c94 100644 --- a/firecrestspawner/spawner.py +++ b/firecrestspawner/spawner.py @@ -452,7 +452,6 @@ async def poll(self): class_name = caller_frame.frame.f_locals['self'].__class__.__name__ if class_name == "HomeHandler": - # auth_state_refreshed = await self.user.authenticator.refresh_access_token(self.user) auth_state = await self.user.get_auth_state() auth = AuthorizationCodeFlowAuth(