Fetch JupyterHub roles from Keycloak (#2447)

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

@@ -0,0 +1,133 @@
+import json
+import os
+import urllib
+from functools import reduce
+
+from jupyterhub.traitlets import Callable
+from oauthenticator.generic import GenericOAuthenticator
+from traitlets import Bool, Unicode, Union
+
+
+class KeyCloakOAuthenticator(GenericOAuthenticator):
+    """
+    Since `oauthenticator` 16.3 `GenericOAuthenticator` supports group management.
+    This subclass adds role management on top of it, building on the new `manage_roles`
+    feature added in JupyterHub 5.0 (https://github.com/jupyterhub/jupyterhub/pull/4748).
+    """
+
+    claim_roles_key = Union(
+        [Unicode(os.environ.get("OAUTH2_ROLES_KEY", "groups")), Callable()],
+        config=True,
+        help="""As `claim_groups_key` but for roles.""",
+    )
+
+    realm_api_url = Unicode(
+        config=True, help="""The keycloak REST API URL for the realm."""
+    )
+
+    reset_managed_roles_on_startup = Bool(True)
+
+    async def update_auth_model(self, auth_model):
+        auth_model = await super().update_auth_model(auth_model)
+        user_info = auth_model["auth_state"][self.user_auth_state_key]
+        user_roles = self._get_user_roles(user_info)
+        auth_model["roles"] = [{"name": role_name} for role_name in user_roles]
+        # note: because the roles check is comprehensive, we need to re-add the admin and user roles
+        if auth_model["admin"]:
+            auth_model["roles"].append({"name": "admin"})
+        if self.check_allowed(auth_model["name"], auth_model):
+            auth_model["roles"].append({"name": "user"})
+        return auth_model
+
+    async def load_managed_roles(self):
+        if not self.manage_roles:
+            raise ValueError(
+                "Managed roles can only be loaded when `manage_roles` is True"
+            )
+        token = await self._get_token()
+
+        # Get the clients list to find the "id" of "jupyterhub" client.
+        clients_data = await self._fetch_api(endpoint="clients/", token=token)
+        jupyterhub_clients = [
+            client for client in clients_data if client["clientId"] == "jupyterhub"
+        ]
+        assert len(jupyterhub_clients) == 1
+        jupyterhub_client_id = jupyterhub_clients[0]["id"]
+
+        # Includes roles like "jupyterhub_admin", "jupyterhub_developer", "dask_gateway_developer"
+        client_roles = await self._fetch_api(
+            endpoint=f"clients/{jupyterhub_client_id}/roles", token=token
+        )
+        # Includes roles like "default-roles-nebari", "offline_access", "uma_authorization"
+        realm_roles = await self._fetch_api(endpoint="roles", token=token)
+        roles = {
+            role["name"]: {"name": role["name"], "description": role["description"]}
+            for role in [*realm_roles, *client_roles]
+        }
+        # we could use either `name` (e.g. "developer") or `path` ("/developer");
+        # since the default claim key returns `path`, it seems preferable.
+        group_name_key = "path"
+        for realm_role in realm_roles:
+            role_name = realm_role["name"]
+            role = roles[role_name]
+            # fetch role assignments to groups
+            groups = await self._fetch_api(f"roles/{role_name}/groups", token=token)
+            role["groups"] = [group[group_name_key] for group in groups]
+            # fetch role assignments to users
+            users = await self._fetch_api(f"roles/{role_name}/users", token=token)
+            role["users"] = [user["username"] for user in users]
+        for client_role in client_roles:
+            role_name = client_role["name"]
+            role = roles[role_name]
+            # fetch role assignments to groups
+            groups = await self._fetch_api(
+                f"clients/{jupyterhub_client_id}/roles/{role_name}/groups", token=token
+            )
+            role["groups"] = [group[group_name_key] for group in groups]
+            # fetch role assignments to users
+            users = await self._fetch_api(
+                f"clients/{jupyterhub_client_id}/roles/{role_name}/users", token=token
+            )
+            role["users"] = [user["username"] for user in users]
+
+        return list(roles.values())
+
+    def _get_user_roles(self, user_info):
+        if callable(self.claim_roles_key):
+            return set(self.claim_roles_key(user_info))
+        try:
+            return set(reduce(dict.get, self.claim_roles_key.split("."), user_info))
+        except TypeError:
+            self.log.error(
+                f"The claim_roles_key {self.claim_roles_key} does not exist in the user token"
+            )
+            return set()
+
+    async def _get_token(self) -> str:
+        http = self.http_client
+
+        body = urllib.parse.urlencode(
+            {
+                "client_id": self.client_id,
+                "client_secret": self.client_secret,
+                "grant_type": "client_credentials",
+            }
+        )
+        response = await http.fetch(
+            self.token_url,
+            method="POST",
+            body=body,
+        )
+        data = json.loads(response.body)
+        return data["access_token"]  # type: ignore[no-any-return]
+
+    async def _fetch_api(self, endpoint: str, token: str):
+        response = await self.http_client.fetch(
+            f"{self.realm_api_url}/{endpoint}",
+            method="GET",
+            headers={"Authorization": f"Bearer {token}"},
+        )
+        return json.loads(response.body)
+
+
+c.JupyterHub.authenticator_class = KeyCloakOAuthenticator

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

@@ -130,6 +130,7 @@ resource "helm_release" "jupyterhub" {
           "01-theme.py"    = file("${path.module}/files/jupyterhub/01-theme.py")
           "02-spawner.py"  = file("${path.module}/files/jupyterhub/02-spawner.py")
           "03-profiles.py" = file("${path.module}/files/jupyterhub/03-profiles.py")
+          "04-auth.py"     = file("${path.module}/files/jupyterhub/04-auth.py")
         }
 
         services = {
@@ -143,25 +144,25 @@ resource "helm_release" "jupyterhub" {
         # for simple key value configuration with jupyterhub traitlets
         # this hub.config property should be used
         config = {
-          JupyterHub = {
-            authenticator_class = "generic-oauth"
-          }
           Authenticator = {
             enable_auth_state = true
           }
-          GenericOAuthenticator = {
+          KeyCloakOAuthenticator = {
             client_id            = module.jupyterhub-openid-client.config.client_id
             client_secret        = module.jupyterhub-openid-client.config.client_secret
             oauth_callback_url   = "https://${var.external-url}/hub/oauth_callback"
             authorize_url        = module.jupyterhub-openid-client.config.authentication_url
             token_url            = module.jupyterhub-openid-client.config.token_url
             userdata_url         = module.jupyterhub-openid-client.config.userinfo_url
+            realm_api_url        = module.jupyterhub-openid-client.config.realm_api_url
             login_service        = "Keycloak"
             username_claim       = "preferred_username"
             claim_groups_key     = "groups"
+            claim_roles_key      = "roles"
             allowed_groups       = ["/analyst", "/developer", "/admin", "jupyterhub_admin", "jupyterhub_developer"]
             admin_groups         = ["/admin", "jupyterhub_admin"]
             manage_groups        = true
+            manage_roles         = true
             refresh_pre_spawn    = true
             validate_server_cert = false
 
@@ -283,6 +284,10 @@ module "jupyterhub-openid-client" {
     var.jupyterhub-logout-redirect-url
   ]
   jupyterlab_profiles_mapper = true
+  service-accounts-enabled   = true
+  service-account-roles = [
+    "view-realm", "view-users", "view-clients"
+  ]
 }
 
 

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

@@ -15,7 +15,8 @@ resource "keycloak_openid_client" "main" {
   access_type           = "CONFIDENTIAL"
   standard_flow_enabled = true
 
-  valid_redirect_uris = var.callback-url-paths
+  valid_redirect_uris      = var.callback-url-paths
+  service_accounts_enabled = var.service-accounts-enabled
 }
 
 
@@ -62,6 +63,33 @@ resource "keycloak_openid_user_attribute_protocol_mapper" "jupyterlab_profiles"
   aggregate_attributes = true
 }
 
+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)
+
+  realm_id  = data.keycloak_realm.master.id
+  client_id = data.keycloak_openid_client.realm_management.id
+  name      = each.key
+}
+
+resource "keycloak_openid_client_service_account_role" "main" {
+  for_each = toset(var.service-account-roles)
+
+  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
+}
+
+
 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

@@ -1,12 +1,14 @@
 output "config" {
   description = "configuration credentials for connecting to openid client"
   value = {
-    client_id     = keycloak_openid_client.main.client_id
-    client_secret = keycloak_openid_client.main.client_secret
+    client_id               = keycloak_openid_client.main.client_id
+    client_secret           = keycloak_openid_client.main.client_secret
+    service_account_user_id = keycloak_openid_client.main.service_account_user_id
 
     authentication_url = "https://${var.external-url}/auth/realms/${var.realm_id}/protocol/openid-connect/auth"
     token_url          = "https://${var.external-url}/auth/realms/${var.realm_id}/protocol/openid-connect/token"
     userinfo_url       = "https://${var.external-url}/auth/realms/${var.realm_id}/protocol/openid-connect/userinfo"
+    realm_api_url      = "https://${var.external-url}/auth/admin/realms/${var.realm_id}"
     callback_urls      = var.callback-url-paths
   }
 }

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

@@ -16,6 +16,19 @@ variable "external-url" {
 }
 
 
+variable "service-accounts-enabled" {
+  description = "Whether the client should have a service account created"
+  type        = bool
+  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     = []
+}
+
+
 variable "role_mapping" {
   description = "Group to role mapping to establish for client"
   type        = map(list(string))

tests/tests_deployment/test_jupyterhub_api.py

@@ -0,0 +1,42 @@
+import pytest
+
+from tests.tests_deployment import constants
+from tests.tests_deployment.utils import get_jupyterhub_session
+
+
+@pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
+def test_jupyterhub_loads_roles_from_keycloak():
+    session = get_jupyterhub_session()
+    xsrf_token = session.cookies.get("_xsrf")
+    response = session.get(
+        f"https://{constants.NEBARI_HOSTNAME}/hub/api/users/{constants.KEYCLOAK_USERNAME}",
+        headers={"X-XSRFToken": xsrf_token},
+        verify=False,
+    )
+    user = response.json()
+    assert set(user["roles"]) == {
+        "user",
+        "manage-account",
+        "jupyterhub_developer",
+        "argo-developer",
+        "dask_gateway_developer",
+        "grafana_viewer",
+        "conda_store_developer",
+        "argo-viewer",
+        "grafana_developer",
+        "manage-account-links",
+        "view-profile",
+    }
+
+
+@pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
+def test_jupyterhub_loads_groups_from_keycloak():
+    session = get_jupyterhub_session()
+    xsrf_token = session.cookies.get("_xsrf")
+    response = session.get(
+        f"https://{constants.NEBARI_HOSTNAME}/hub/api/users/{constants.KEYCLOAK_USERNAME}",
+        headers={"X-XSRFToken": xsrf_token},
+        verify=False,
+    )
+    user = response.json()
+    assert set(user["groups"]) == {"/analyst", "/developer", "/users"}