add service account with allow-app-sharing-role permissions

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/conda-store/server.tf

@@ -70,9 +70,9 @@ module "conda-store-openid-client" {
     "https://${var.external-url}/conda-store/oauth_callback"
   ]
   service-accounts-enabled = true
-  service-account-roles = [
-    "view-realm", "view-users", "view-clients"
-  ]
+  service-account-roles = {
+    "realm-management" : ["view-realm", "view-users", "view-clients"]
+  }
 }
 
 

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/files/jupyterhub/02-spawner.py

@@ -2,7 +2,6 @@
 import json
 
 import kubernetes.client.models
-from tornado import gen
 
 kubernetes.client.models.V1EndpointPort = (
     kubernetes.client.models.CoreV1EndpointPort
@@ -11,9 +10,8 @@
 from kubespawner import KubeSpawner  # noqa: E402
 
 
-@gen.coroutine
-def get_username_hook(spawner):
-    auth_state = yield spawner.user.get_auth_state()
+async def get_username_hook(spawner):
+    auth_state = await spawner.user.get_auth_state()
     username = auth_state["oauth_user"]["preferred_username"]
 
     spawner.environment.update(
@@ -23,6 +21,13 @@ def get_username_hook(spawner):
     )
 
 
+async def pre_spawn_hook(spawner):
+    # if we are starting a service account pod, set/update auth_state
+    if spawner.user.name == spawner.authenticator.JHUB_SERVICE_ACCOUNT_NAME:
+        await spawner.authenticator.set_jhub_service_account_auth_state(spawner.user)
+    await get_username_hook(spawner)
+
+
 def get_conda_store_environments(user_info: dict):
     import urllib3
     import yarl
@@ -44,7 +49,7 @@ def get_conda_store_environments(user_info: dict):
     return [f"{env['namespace']['name']}-{env['name']}" for env in j.get("data", [])]
 
 
-c.Spawner.pre_spawn_hook = get_username_hook
+c.Spawner.pre_spawn_hook = pre_spawn_hook
 
 c.JupyterHub.allow_named_servers = False
 c.JupyterHub.spawner_class = KubeSpawner

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/files/jupyterhub/03-profiles.py

@@ -554,7 +554,6 @@ def render_profiles(spawner):
     # userinfo request to have the groups in the key
     # "auth_state.oauth_user.groups"
     auth_state = yield spawner.user.get_auth_state()
-
     username = auth_state["oauth_user"]["preferred_username"]
 
     # only return the lowest level group name

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/files/jupyterhub/04-auth.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import logging
 import os
 import time
 import urllib
@@ -8,7 +9,7 @@
 from jupyterhub import scopes
 from jupyterhub.traitlets import Callable
 from oauthenticator.generic import GenericOAuthenticator
-from traitlets import Bool, Unicode, Union
+from traitlets import Bool, Unicode, Union, default
 
 
 class KeyCloakOAuthenticator(GenericOAuthenticator):
@@ -18,6 +19,12 @@ class KeyCloakOAuthenticator(GenericOAuthenticator):
     feature added in JupyterHub 5.0 (https://github.com/jupyterhub/jupyterhub/pull/4748).
     """
 
+    JHUB_SERVICE_ACCOUNT_NAME = Unicode()
+
+    @default("JHUB_SERVICE_ACCOUNT_NAME")
+    def _default_jhub_service_account_name(self):
+        return f"service-account-{self.client_id}"
+
     claim_roles_key = Union(
         [Unicode(os.environ.get("OAUTH2_ROLES_KEY", "groups")), Callable()],
         config=True,
@@ -30,6 +37,36 @@ class KeyCloakOAuthenticator(GenericOAuthenticator):
 
     reset_managed_roles_on_startup = Bool(True)
 
+    async def set_jhub_service_account_auth_state(self, user):
+        if user.name != self.JHUB_SERVICE_ACCOUNT_NAME:
+            raise ValueError(
+                'User name "{user.name}" does not match service account name "{self.JHUB_SERVICE_ACCOUNT_NAME}"'
+            )
+        auth_model = await self.authenticate_service_account()
+        await user.save_auth_state(auth_model["auth_state"])
+        logging.info(f'Auth state set for service account: "{user.name}"')
+
+    async def authenticate_service_account(self):
+        token_info = await self._get_token_info()
+
+        # Get user info using the access token
+        user_info = await self.token_to_user(token_info)
+
+        # Get/set username
+        username = self.user_info_to_username(user_info)
+        username = self.normalize_username(username)
+
+        # Build auth model similar to OAuth flow
+        auth_model = {
+            "name": username,
+            "admin": True if username in self.admin_users else None,
+            "auth_state": self.build_auth_state_dict(token_info, user_info),
+        }
+
+        auth_model = await self.update_auth_model(auth_model)
+
+        return auth_model
+
     async def update_auth_model(self, auth_model):
         """Updates and returns the auth_model dict.
         This function is called every time a user authenticates with JupyterHub, as in
@@ -46,15 +83,15 @@ async def update_auth_model(self, auth_model):
         user_id = auth_model["auth_state"]["oauth_user"]["sub"]
         token = await self._get_token()
 
-        jupyterhub_client_id = await self._get_jupyterhub_client_id(token=token)
+        jupyterhub_client_uuid = await self._get_jupyterhub_client_uuid(token=token)
         user_info = auth_model["auth_state"][self.user_auth_state_key]
         user_roles_from_claims = self._get_user_roles(user_info=user_info)
         keycloak_api_call_start = time.time()
         user_roles = await self._get_client_roles_for_user(
-            user_id=user_id, client_id=jupyterhub_client_id, token=token
+            user_id=user_id, client_id=jupyterhub_client_uuid, token=token
         )
         user_roles_rich = await self._get_roles_with_attributes(
-            roles=user_roles, client_id=jupyterhub_client_id, token=token
+            roles=user_roles, client_id=jupyterhub_client_uuid, token=token
         )
 
         # Include which groups have permission to mount shared directories (user by
@@ -63,7 +100,7 @@ async def update_auth_model(self, auth_model):
             await self.get_client_groups_with_mount_permissions(
                 user_groups=auth_model["auth_state"]["oauth_user"]["groups"],
                 user_roles=user_roles_rich,
-                client_id=jupyterhub_client_id,
+                client_id=jupyterhub_client_uuid,
                 token=token,
             )
         )
@@ -114,7 +151,7 @@ async def _get_jupyterhub_client_roles(self, jupyterhub_client_id, token):
         )
         return client_roles_rich
 
-    async def _get_jupyterhub_client_id(self, token):
+    async def _get_jupyterhub_client_uuid(self, token):
         # Get the clients list to find the "id" of "jupyterhub" client.
         clients_data = await self._fetch_api(endpoint="clients/", token=token)
         jupyterhub_clients = [
@@ -131,9 +168,9 @@ async def load_managed_roles(self):
                 "Managed roles can only be loaded when `manage_roles` is True"
             )
         token = await self._get_token()
-        jupyterhub_client_id = await self._get_jupyterhub_client_id(token=token)
+        jupyterhub_client_uuid = await self._get_jupyterhub_client_uuid(token=token)
         client_roles_rich = await self._get_jupyterhub_client_roles(
-            jupyterhub_client_id=jupyterhub_client_id, token=token
+            jupyterhub_client_id=jupyterhub_client_uuid, token=token
         )
 
         # Includes roles like "default-roles-nebari", "offline_access", "uma_authorization"
@@ -168,7 +205,7 @@ async def load_managed_roles(self):
                 await self._get_users_and_groups_for_role(
                     role_name,
                     token=token,
-                    client_id=jupyterhub_client_id,
+                    client_id=jupyterhub_client_uuid,
                 )
             )
 
@@ -307,7 +344,7 @@ def _get_user_roles(self, user_info):
             )
             return set()
 
-    async def _get_token(self) -> str:
+    async def _get_token_info(self) -> str:
         http = self.http_client
 
         body = urllib.parse.urlencode(
@@ -322,8 +359,12 @@ async def _get_token(self) -> str:
             method="POST",
             body=body,
         )
-        data = json.loads(response.body)
-        return data["access_token"]  # type: ignore[no-any-return]
+        token_info = json.loads(response.body)
+        return token_info
+
+    async def _get_token(self) -> str:
+        token_info = await self._get_token_info()
+        return token_info["access_token"]  # type: ignore[no-any-return]
 
     async def _fetch_api(self, endpoint: str, token: str):
         response = await self.http_client.fetch(

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/main.tf

@@ -344,9 +344,10 @@ module "jupyterhub-openid-client" {
   ]
   jupyterlab_profiles_mapper = true
   service-accounts-enabled   = true
-  service-account-roles = [
-    "view-realm", "view-users", "view-clients"
-  ]
+  service-account-roles = {
+    "realm-management" : ["view-realm", "view-users", "view-clients"],
+    "jupyterhub" : ["allow-app-sharing-role"]
+  }
 }
 
 

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/jupyterhub/versions.tf

@@ -0,0 +1,9 @@
+terraform {
+  required_providers {
+    keycloak = {
+      source  = "mrparkers/keycloak"
+      version = "3.7.0"
+    }
+  }
+  required_version = ">= 1.0"
+}

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/keycloak-client/main.tf

@@ -67,29 +67,54 @@ data "keycloak_realm" "master" {
   realm = "nebari"
 }
 
-data "keycloak_openid_client" "realm_management" {
-  realm_id  = var.realm_id
-  client_id = "realm-management"
-}
 
-data "keycloak_role" "main-service" {
-  for_each = toset(var.service-account-roles)
+# Get client data for each service account client
+data "keycloak_openid_client" "service_clients" {
+  for_each = var.service-account-roles
 
-  realm_id  = data.keycloak_realm.master.id
-  client_id = data.keycloak_openid_client.realm_management.id
-  name      = each.key
-}
+  realm_id   = var.realm_id
+  client_id  = each.key
+  depends_on = [keycloak_openid_client.main]
+}
+
+# Get role data for each client's roles
+data "keycloak_role" "client_roles" {
+  for_each = {
+    for pair in flatten([
+      for client, roles in var.service-account-roles : [
+        for role in roles : {
+          key    = "${client}-${role}"
+          client = client
+          role   = role
+        }
+      ]
+    ]) : pair.key => pair
+  }
 
-resource "keycloak_openid_client_service_account_role" "main" {
-  for_each = toset(var.service-account-roles)
+  realm_id  = var.realm_id
+  client_id = data.keycloak_openid_client.service_clients[each.value.client].id
+  name      = each.value.role
+}
+
+resource "keycloak_openid_client_service_account_role" "client_roles" {
+  for_each = {
+    for pair in flatten([
+      for client, roles in var.service-account-roles : [
+        for role in roles : {
+          key    = "${client}-${role}"
+          client = client
+          role   = role
+        }
+      ]
+    ]) : pair.key => pair
+  }
 
   realm_id                = var.realm_id
   service_account_user_id = keycloak_openid_client.main.service_account_user_id
-  client_id               = data.keycloak_openid_client.realm_management.id
-  role                    = data.keycloak_role.main-service[each.key].name
+  client_id               = data.keycloak_openid_client.service_clients[each.value.client].id
+  role                    = data.keycloak_role.client_roles[each.key].name
 }
 
-
 resource "keycloak_role" "main" {
   for_each = toset(flatten(values(var.role_mapping)))
 

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/keycloak-client/outputs.tf

@@ -12,3 +12,10 @@ output "config" {
     callback_urls      = var.callback-url-paths
   }
 }
+
+output "client_role_ids" {
+  description = "Map of role names to their IDs"
+  value = {
+    for role_key, role in keycloak_role.default_client_roles : role_key => role.id
+  }
+}

src/_nebari/stages/kubernetes_services/template/modules/kubernetes/services/keycloak-client/variables.tf

@@ -22,10 +22,17 @@ variable "service-accounts-enabled" {
   default     = false
 }
 
+
 variable "service-account-roles" {
-  description = "Roles to be granted to the service account. Requires setting service-accounts-enabled to true."
-  type        = list(string)
-  default     = []
+  description = <<-EOT
+  Map of client to client-roles to be granted to the service account client. Requires setting service-accounts-enabled to true.
+
+  e.g. {
+    "my-client": ["my-client-role"],
+  }
+  EOT
+  type        = map(list(string))
+  default     = {}
 }