fix(jupyterhub): admin vault token renewal

This commit is contained in:
Masaki Yatsu
2025-09-08 18:43:08 +09:00
parent c82c6aa22b
commit 4df776c181
6 changed files with 322 additions and 127 deletions

View File

@@ -23,6 +23,11 @@ hub:
mode: 0644
stringData: |
{{ .Env.USER_POLICY_HCL | strings.Indent 8 }}
pre_spawn_hook.py:
mountPath: /srv/jupyterhub/pre_spawn_hook.py
mode: 0644
stringData: |
{{ file.Read "pre_spawn_hook.py" | strings.Indent 8 }}
# Override the default command to run our startup script first
command:
@@ -56,112 +61,10 @@ hub:
- email
extraConfig:
pre-spawn-hook: |
# Set environment variables for spawned containers
import hvac
{{- if eq .Env.JUPYTERHUB_VAULT_INTEGRATION_ENABLED "true" }}
def get_vault_token():
"""Read Vault token from file"""
import os
token_file = '/vault/secrets/vault-token'
try:
with open(token_file, 'r') as f:
token = f.read().strip()
if token:
return token
except FileNotFoundError:
print(f"Token file not found: {token_file}")
except Exception as e:
print(f"Error reading token file {token_file}: {e}")
return None
{{- end }}
async def pre_spawn_hook(spawner):
"""Set essential environment variables for spawned containers"""
# PostgreSQL configuration
spawner.environment["POSTGRES_HOST"] = "postgres-cluster-rw.postgres"
spawner.environment["POSTGRES_PORT"] = "5432"
# JupyterHub API configuration
spawner.environment["JUPYTERHUB_API_URL"] = "http://hub:8081/hub/api"
# Logging configuration
spawner.environment["BUUNSTACK_LOG_LEVEL"] = "{{ .Env.JUPYTER_BUUNSTACK_LOG_LEVEL }}"
{{- if eq .Env.JUPYTERHUB_VAULT_INTEGRATION_ENABLED "true" }}
# Create user-specific Vault token directly
try:
username = spawner.user.name
# Step 1: Initialize admin Vault client with file-based token
import os
vault_addr = os.environ.get("VAULT_ADDR", "{{ .Env.VAULT_ADDR }}")
vault_token = get_vault_token()
spawner.log.info(f"pre_spawn_hook starting for {username}")
spawner.log.info(f"Vault address: {vault_addr}")
spawner.log.info(f"Vault token source: {'file' if os.path.exists('/vault/secrets/vault-token') else 'env'}")
spawner.log.info(f"Vault token present: {bool(vault_token)}, length: {len(vault_token) if vault_token else 0}")
if not vault_token:
raise Exception("No Vault token available from file or environment")
vault_client = hvac.Client(url=vault_addr, verify=False)
vault_client.token = vault_token
if not vault_client.is_authenticated():
raise Exception("Admin token is not authenticated")
# Step 2: Create user-specific policy
user_policy_name = "jupyter-user-{}".format(username)
# Read policy template from file
import os
policy_template_path = "/srv/jupyterhub/user_policy.hcl"
with open(policy_template_path, 'r') as f:
policy_template = f.read()
# Replace {username} placeholder with actual username
user_policy = policy_template.replace("{username}", username)
# Write user-specific policy
try:
vault_client.sys.create_or_update_policy(user_policy_name, user_policy)
spawner.log.info("✅ Created policy: {}".format(user_policy_name))
except Exception as policy_e:
spawner.log.warning("Policy creation failed (may already exist): {}".format(policy_e))
# Step 3: Create user-specific token
# Get TTL settings from environment variables
user_token_ttl = os.environ.get("NOTEBOOK_VAULT_TOKEN_TTL", "24h")
user_token_max_ttl = os.environ.get("NOTEBOOK_VAULT_TOKEN_MAX_TTL", "168h")
token_response = vault_client.auth.token.create_orphan(
policies=[user_policy_name],
ttl=user_token_ttl,
renewable=True,
display_name="notebook-{}".format(username),
explicit_max_ttl=user_token_max_ttl
)
user_vault_token = token_response["auth"]["client_token"]
lease_duration = token_response["auth"].get("lease_duration", 3600)
# Set user-specific Vault token as environment variable
spawner.environment["NOTEBOOK_VAULT_TOKEN"] = user_vault_token
spawner.log.info("✅ User-specific Vault token created for {} (TTL: {}s, renewable, max TTL: {})".format(username, lease_duration, user_token_max_ttl))
except Exception as e:
spawner.log.error("Failed to create user-specific Vault token for {}: {}".format(spawner.user.name, e))
import traceback
spawner.log.error("Full traceback: {}".format(traceback.format_exc()))
{{- end }}
c.KubeSpawner.pre_spawn_hook = pre_spawn_hook
load-pre-spawn-hook: |
# Load pre_spawn_hook from external file
with open('/srv/jupyterhub/pre_spawn_hook.py', 'r') as f:
exec(f.read())
{{- if eq .Env.JUPYTERHUB_VAULT_INTEGRATION_ENABLED "true" }}
# Vault Agent sidecar configuration
@@ -186,7 +89,7 @@ hub:
extraContainers:
- name: vault-agent
image: hashicorp/vault:1.15.2
image: hashicorp/vault:1.17.5
securityContext:
runAsUser: 100
runAsGroup: 101
@@ -205,6 +108,8 @@ hub:
env:
- name: VAULT_ADDR
value: {{ .Env.VAULT_ADDR | quote }}
- name: JUPYTERHUB_VAULT_TOKEN_TTL
value: {{ .Env.JUPYTERHUB_VAULT_TOKEN_TTL | quote }}
volumeMounts:
- name: vault-secrets
mountPath: /vault/secrets
@@ -378,7 +283,12 @@ cull:
# timeout: 300 # 5 minutes idle timeout (for testing) │ │
# every: 60 # Check every 1 minute (for testing) │ │
# maxAge: 86400 # Maximum age of a server pod (1 day)
# Maximum age of a server pod before forced restart
# IMPORTANT: This must be less than NOTEBOOK_VAULT_TOKEN_MAX_TTL to prevent token expiry
# - NOTEBOOK_VAULT_TOKEN_MAX_TTL: {{ .Env.NOTEBOOK_VAULT_TOKEN_MAX_TTL }} (7 days = 604800s)
# - JUPYTERHUB_CULL_MAX_AGE: {{ .Env.JUPYTERHUB_CULL_MAX_AGE }}s (6 days = 518400s)
# Pod restart creates new user token, preventing 7-day token expiry
maxAge: {{ .Env.JUPYTERHUB_CULL_MAX_AGE }}
adminUsers: true # Also cull admin users' server pods
users: false # Don't delete user accounts, only stop server pods