Add a default roles at initialisation

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

@@ -1,14 +1,40 @@
 import json
 import os
 import time
+import typing
 import urllib
+from collections import defaultdict
 from functools import reduce
 
 from jupyterhub import scopes
 from jupyterhub.traitlets import Callable
 from oauthenticator.generic import GenericOAuthenticator
 from traitlets import Bool, Unicode, Union
 
+# A set of roles to create automatically to help with basic permissions
+DEFAULT_ROLES = [
+    {
+        "name": "allow-app-sharing-role",
+        "description": "Allow app sharing for apps created via JupyterHub App Launcher (jhub-apps)",
+        # grants permissions to share server
+        # grants permissions to read other user's names
+        # grants permissions to read other groups' names
+        # The later two are required for sharing with a group or user
+        "scopes": ["shares,read:users:name,read:groups:name"],
+        # not attaching this to any group by default as that might not be desirable for all
+        # deployments and this gives the user (admin) a choice to attach or not attach this to any
+        # user or group based on the permission structure of the team / organization.
+    },
+    {
+        "name": "allow-read-access-to-services-role",
+        "description": "Allow read access to services, such that they are visible on the home page e.g. conda-store",
+        # grants permissions to read services
+        "scopes": ["read:services"],
+        # Adding it to analyst group such that it's applied to every user.
+        "groups": ["/analyst"],
+    },
+]
+
 
 class KeyCloakOAuthenticator(GenericOAuthenticator):
     """
@@ -44,7 +70,6 @@ async def update_auth_model(self, auth_model):
         auth_model = await super().update_auth_model(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)
         user_info = auth_model["auth_state"][self.user_auth_state_key]
         user_roles_from_claims = self._get_user_roles(user_info=user_info)
@@ -113,9 +138,19 @@ async def load_managed_roles(self):
             )
         token = await self._get_token()
         jupyterhub_client_id = await self._get_jupyterhub_client_id(token=token)
+
         client_roles_rich = await self._get_jupyterhub_client_roles(
             jupyterhub_client_id=jupyterhub_client_id, token=token
         )
+        try:
+            # Creating the default roles in keycloak instead of jupyterhub directly
+            # to keep keycloak as single source of truth.
+            await self._create_default_keycloak_client_roles(
+                DEFAULT_ROLES, client_roles_rich, jupyterhub_client_id, token
+            )
+        except Exception as e:
+            self.log.error("Unable to create default roles")
+            self.log.exception(e)
         # Includes roles like "default-roles-nebari", "offline_access", "uma_authorization"
         realm_roles = await self._fetch_api(endpoint="roles", token=token)
         roles = {
@@ -151,7 +186,6 @@ async def load_managed_roles(self):
                 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_scope_from_role(self, role):
@@ -179,17 +213,129 @@ def validate_scopes(self, role_scopes):
         return []
 
     async def _get_roles_with_attributes(self, roles: dict, client_id: str, token: str):
-        """This fetches all roles by id to fetch there attributes."""
+        """This fetches all roles by id to fetch their attributes."""
         roles_rich = []
         for role in roles:
             # If this takes too much time, which isn't the case right now, we can
-            # also do multi-threaded requests
+            # also do multithreaded requests
             role_rich = await self._fetch_api(
                 endpoint=f"roles-by-id/{role['id']}?client={client_id}", token=token
             )
             roles_rich.append(role_rich)
         return roles_rich
 
+    async def _create_default_keycloak_client_roles(
+        self,
+        roles: typing.List[dict],
+        existing_roles: typing.List[dict],
+        client_id: str,
+        token: str,
+    ):
+        """Create default roles for jupyterhub keycloak client"""
+        self.log.info("Creating default roles, which does not exists already")
+        self.log.info(
+            f"Roles to create: {roles}, existing roles: {existing_roles}, client_id: {client_id}"
+        )
+        existing_role_name_mapping = {role["name"]: role for role in existing_roles}
+
+        for role in roles:
+            self.log.info(f"Creating role: {role}")
+            if role["name"] in existing_role_name_mapping:
+                self.log.info(f"role: {role} exists skipping")
+                continue
+            await self._create_keycloak_client_role(role, client_id, token)
+
+        role_name_mapping = await self._get_keycloak_roles(
+            token=token, client_id=client_id
+        )
+        groups = await self._get_keycloak_groups(token)
+        for role in roles:
+            keycloak_roles = role_name_mapping[role["name"]]
+            if len(keycloak_roles) != 1:
+                self.log.error(
+                    f"Multiple roles with same name: {keycloak_roles}, skipping"
+                )
+                continue
+            if not role.get("groups"):
+                self.log.info(
+                    f"No groups defined for the role {keycloak_roles}, not attaching to any group"
+                )
+                continue
+            for group_path in role.get("groups"):
+                keycloak_group = groups[group_path][0]
+                keycloak_group_id = keycloak_group["id"]
+                self.log.info(
+                    f"Assigning role: {keycloak_roles[0]['name']} "
+                    f"to group: {keycloak_group['name']},"
+                    f"client: {client_id}"
+                )
+                response_content = await self._assign_keycloak_client_role(
+                    client_id=client_id,
+                    group_id=keycloak_group_id,
+                    token=token,
+                    role=keycloak_roles[0],
+                )
+                self.log.info(f"Role assignment response_content: {response_content}")
+
+    async def _create_keycloak_client_role(self, role, client_id, token):
+        self.log.info(f"Creating keycloak client role: {role}")
+        body = json.dumps(
+            {
+                "name": role.get("name"),
+                "description": role.get("description"),
+                "attributes": {
+                    "scopes": role.get("scopes"),
+                    "component": ["jupyterhub"],
+                },
+            }
+        )
+        response = await self._fetch_api(
+            endpoint=f"clients/{client_id}/roles",
+            token=token,
+            method="POST",
+            body=body,
+            extra_headers={"Content-Type": "application/json"},
+        )
+        self.log.info(f"Keycloak client role creation response: {response}")
+
+    async def _get_keycloak_groups(self, token: str) -> typing.DefaultDict[str, list]:
+        self.log.info("Getting keycloak groups")
+        response_json = await self._fetch_api(endpoint="groups", token=token)
+        self.log.info(f"Keycloak groups: {response_json}")
+        group_name_mapping = defaultdict(list)
+        for group in response_json:
+            group_name_mapping[group["path"]].append(group)
+        self.log.info(f"Keycloak groups name mapping: {group_name_mapping}")
+        return group_name_mapping
+
+    async def _get_keycloak_roles(
+        self, token: str, client_id: str
+    ) -> typing.DefaultDict[str, list]:
+        """Get keycloak roles for a client"""
+        self.log.info(f"getting keycloak roles for client: {client_id}")
+        response_json = await self._fetch_api(
+            endpoint=f"clients/{client_id}/roles", token=token
+        )
+        role_name_mapping = defaultdict(list)
+        for role in response_json:
+            role_name_mapping[role["name"]].append(role)
+        self.log.info(f"keycloak roles name mapping: {role_name_mapping}")
+        return role_name_mapping
+
+    async def _assign_keycloak_client_role(
+        self, client_id: str, group_id: str, token: str, role: dict
+    ):
+        """Given a group id and a role, assign role to the group"""
+        response_content = await self._fetch_api(
+            endpoint=f"groups/{group_id}/role-mappings/clients/{client_id}",
+            token=token,
+            method="POST",
+            body=json.dumps([role]),
+            extra_headers={"Content-Type": "application/json"},
+        )
+        self.log.info(f"role assignment response: {response_content}")
+        return response_content
+
     async def _get_client_roles_for_user(self, user_id, client_id, token):
         user_roles = await self._fetch_api(
             endpoint=f"users/{user_id}/role-mappings/clients/{client_id}/composite",
@@ -226,13 +372,25 @@ async def _get_token(self) -> str:
         data = json.loads(response.body)
         return data["access_token"]  # type: ignore[no-any-return]
 
-    async def _fetch_api(self, endpoint: str, token: str):
+    async def _fetch_api(
+        self,
+        endpoint: str,
+        token: str,
+        method: str = "GET",
+        extra_headers=None,
+        **kwargs,
+    ):
+        append_headers = extra_headers if extra_headers else {}
         response = await self.http_client.fetch(
             f"{self.realm_api_url}/{endpoint}",
-            method="GET",
-            headers={"Authorization": f"Bearer {token}"},
+            method=method,
+            headers={"Authorization": f"Bearer {token}", **append_headers},
+            **kwargs,
         )
-        return json.loads(response.body)
+        try:
+            return json.loads(response.body)
+        except json.decoder.JSONDecodeError:
+            return response.body
 
 
 c.JupyterHub.authenticator_class = KeyCloakOAuthenticator

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

@@ -286,7 +286,9 @@ module "jupyterhub-openid-client" {
   jupyterlab_profiles_mapper = true
   service-accounts-enabled   = true
   service-account-roles = [
-    "view-realm", "view-users", "view-clients"
+    # "manage-clients" is required for creating roles for the client
+    # "manage-users" is required for attaching roles to groups
+    "view-realm", "view-users", "view-clients", "manage-clients", "manage-users"
   ]
 }
 

tests/tests_deployment/keycloak_utils.py

@@ -81,6 +81,14 @@ def create_keycloak_role(client_name: str, role_name: str, scopes: str, componen
     )
 
 
+def get_keycloak_client_roles(client_name):
+    keycloak_admin = get_keycloak_admin()
+    client_details = get_keycloak_client_details_by_name(
+        client_name=client_name, keycloak_admin=keycloak_admin
+    )
+    return keycloak_admin.get_client_roles(client_id=client_details["id"])
+
+
 def delete_client_keycloak_test_roles(client_name):
     keycloak_admin = get_keycloak_admin()
     client_details = get_keycloak_client_details_by_name(

tests/tests_deployment/test_jupyterhub_api.py

@@ -4,6 +4,7 @@
 from tests.tests_deployment.keycloak_utils import (
     assign_keycloak_client_role_to_user,
     create_keycloak_role,
+    get_keycloak_client_roles,
 )
 from tests.tests_deployment.utils import create_jupyterhub_token, get_jupyterhub_session
 
@@ -30,9 +31,29 @@ def test_jupyterhub_loads_roles_from_keycloak():
         "grafana_developer",
         "manage-account-links",
         "view-profile",
+        # default roles
+        "allow-read-access-to-services-role",
     }
 
 
+@pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
+def test_default_user_role_scopes():
+    token_response = create_jupyterhub_token(note="get-default-scopes")
+    token_scopes = set(token_response.json()["scopes"])
+    assert "read:services" in token_scopes
+
+
+@pytest.mark.filterwarnings(
+    "ignore:.*auto_refresh_token is deprecated:DeprecationWarning"
+)
+@pytest.mark.filterwarnings("ignore::urllib3.exceptions.InsecureRequestWarning")
+def test_check_default_roles_added_in_keycloak():
+    client_roles = get_keycloak_client_roles(client_name="jupyterhub")
+    role_names = [role["name"] for role in client_roles]
+    assert "allow-app-sharing-role" in role_names
+    assert "allow-read-access-to-services-role" in role_names
+
+
 @pytest.mark.parametrize(
     "component,scopes,expected_scopes_difference",
     (