Skip to content

Commit

Permalink
Fetch JupyterHub roles from Keycloak
Browse files Browse the repository at this point in the history
  • Loading branch information
krassowski committed May 4, 2024
1 parent 3ab3cf9 commit 62e527d
Show file tree
Hide file tree
Showing 5 changed files with 198 additions and 21 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
import json
import os
import urllib
from functools import reduce

from jupyterhub.traitlets import Callable
from oauthenticator.generic import GenericOAuthenticator
from traitlets import 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."""
)

service_account_user_id = Unicode(config=True, help="""service_account_user_id.""")

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"] = user_roles
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
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand All @@ -143,27 +144,28 @@ 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 = {
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
login_service = "Keycloak"
username_claim = "preferred_username"
claim_groups_key = "groups"
allowed_groups = ["/analyst", "/developer", "/admin"]
admin_groups = ["/admin"]
manage_groups = true
refresh_pre_spawn = true
validate_server_cert = false
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
realm_api_url = module.jupyterhub-openid-client.config.realm_api_url
service_account_user_id = module.jupyterhub-openid-client.config.service_account_user_id
login_service = "Keycloak"
username_claim = "preferred_username"
claim_groups_key = "groups"
claim_roles_key = "roles"
allowed_groups = ["/analyst", "/developer", "/admin"]
admin_groups = ["/admin"]
manage_groups = true
manage_roles = true
reset_managed_roles_on_startup = true
refresh_pre_spawn = true
validate_server_cert = false

# deprecated, to be removed (replaced by validate_server_cert)
tls_verify = false
Expand Down Expand Up @@ -283,6 +285,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"
]
}


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
}


Expand Down Expand Up @@ -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)))

Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down

0 comments on commit 62e527d

Please sign in to comment.