JupyterHub 进阶实践：集成 Azure AD 实现企业级单点登录 (二期)

import os
import pwd
import grp
import re
import subprocess

c = get_config()  #noqa

# ═══════════════════════════════════════════════════════
# 全局变量
# ═══════════════════════════════════════════════════════
JUPYTER_USER_DIR = '/data/jupyter-users'
JUPYTER_GROUP = 'jupyterhub-users'

# ═══════════════════════════════════════════════════════
# 网络配置
# ═══════════════════════════════════════════════════════
c.JupyterHub.ip = '0.0.0.0'
c.JupyterHub.port = 8000
c.JupyterHub.base_url = '/'

# CRYPT_KEY：使用 openssl rand -hex 32 生成
os.environ.setdefault('JUPYTERHUB_CRYPT_KEY', 'xxxx')

# ═══════════════════════════════════════════════════════
# Azure AD OAuth 认证与用户名规范化
# ═══════════════════════════════════════════════════════
from oauthenticator.azuread import AzureAdOAuthenticator

class NormalizedAzureAdOAuthenticator(AzureAdOAuthenticator):
    """
    覆写 normalize_username，确保 Azure AD 传回的邮箱格式能够无缝转换为合法的 Linux 用户名。
    转换示例：zhang.san@x.com → zhang-san
    """
    def normalize_username(self, username):
        # 取 @ 前面的邮箱前缀
        name = username.split('@')[0]
        # 将点、下划线、空格等统一转换为短横线，并转为小写
        name = name.replace('.', '-').replace('_', '-').replace(' ', '-').lower()
        # 移除非法字符，限制仅能包含小写字母、数字和短横线
        name = re.sub(r'[^a-z0-9-]', '', name)
        # 截断长度（Linux 用户名限制最长 32 字符）
        name = name[:32]
        return super().normalize_username(name)

c.JupyterHub.authenticator_class = NormalizedAzureAdOAuthenticator

# Azure AD 凭证配置
c.AzureAdOAuthenticator.tenant_id = 'xxxx'
c.AzureAdOAuthenticator.client_id = 'xxxx'
c.AzureAdOAuthenticator.client_secret = 'xxxx'
c.AzureAdOAuthenticator.oauth_callback_url = 'https://jupyter-callback.x.com/hub/oauth_callback'

# 关键：指定 username_claim 为 upn（用户主体名），确保获取到合法的邮箱格式
c.AzureAdOAuthenticator.username_claim = 'upn'

# 允许所有通过当前 Azure AD 租户认证的用户登录
c.AzureAdOAuthenticator.allow_all = True

# 管理员用户配置（配置为规范化后的 Linux 用户名）
c.Authenticator.admin_users = {'admin'}

# ═══════════════════════════════════════════════════════
# SystemdSpawner 配置
# ═══════════════════════════════════════════════════════
c.JupyterHub.spawner_class = 'systemdspawner.SystemdSpawner'
c.Spawner.default_url = '/lab'
c.SystemdSpawner.unit_name_template = 'jupyter-{USERNAME}'
c.Spawner.notebook_dir = JUPYTER_USER_DIR + '/{username}'
c.SystemdSpawner.default_shell = '/bin/bash'
c.SystemdSpawner.isolate_tmp = True
c.SystemdSpawner.isolate_devices = True
c.SystemdSpawner.disable_user_sudo = True
c.SystemdSpawner.slice = 'jupyter-users.slice'

# ═══════════════════════════════════════════════════════
# 自动化用户供给：pre_spawn_hook
# ═══════════════════════════════════════════════════════
CONDA_BIN = '/opt/miniconda3/bin/conda'
USERADD_BIN = '/usr/sbin/useradd'
GROUPADD_BIN = '/usr/sbin/groupadd'

def ensure_system_user(spawner):
    """
    当用户首次通过 Azure AD 登录时，自动执行以下供给逻辑：
    1. 创建 JupyterHub 共享用户组
    2. 创建对应的 Linux 系统用户（禁止 SSH 登录但允许 Terminal 内执行 cron）
    3. 自动初始化 700 隔离工作目录
    4. 写入 Conda 环境变量至其 .bashrc
    5. 构建专属 Conda 隔离沙箱并注册 Jupyter 内核
    """
    username = spawner.user.name

    # 严格校验：确保用户名必须符合规范化正则，拒绝非法输入直接写入系统命令
    if not re.fullmatch(r'[a-z0-9-]{1,32}', username):
        raise ValueError(
            f'Refusing to spawn: username {username!r} is not normalized. '
            f'请检查配置的 username_claim 或清理 jupyterhub.sqlite 中的脏数据。'
        )

    user_dir = f'{JUPYTER_USER_DIR}/{username}'
    conda_env_prefix = f'{user_dir}/.conda/envs/{username}'

    # 1. 确保 jupyterhub-users 组存在
    try:
        grp.getgrnam(JUPYTER_GROUP)
    except KeyError:
        subprocess.run([GROUPADD_BIN, JUPYTER_GROUP], check=True)

    # 2. 如果用户不存在则创建 Linux 用户
    is_new_user = False
    try:
        pwd.getpwnam(username)
    except KeyError:
        subprocess.run([
            USERADD_BIN,
            '-m',
            '-d', user_dir,
            '-s', '/bin/bash',
            '-g', JUPYTER_GROUP,
            '-G', 'crontab',
            '--no-user-group',
            username
        ], check=True)
        spawner.log.info(f'Created system user: {username}')
        is_new_user = True

    # 3. 确保目录存在并强制设定权限
    os.makedirs(user_dir, exist_ok=True)
    user_info = pwd.getpwnam(username)
    os.chown(user_dir, user_info.pw_uid, user_info.pw_gid)
    os.chmod(user_dir, 0o700)

    # 4. 配置 .bashrc 中的 Conda PATH
    if is_new_user:
        bashrc = f'{user_dir}/.bashrc'
        with open(bashrc, 'a') as f:
            f.write('\nexport PATH="/opt/miniconda3/bin:$PATH"\n')
        os.chown(bashrc, user_info.pw_uid, user_info.pw_gid)

    # 5. 构建专属 Conda 沙箱及注册内核（仅在新环境未被初始化时执行）
    if not os.path.exists(conda_env_prefix):
        spawner.log.info(f'Creating conda env: {conda_env_prefix}')
        subprocess.run([
            CONDA_BIN, 'create', '-y',
            '-c', 'conda-forge', '--override-channels',
            '--prefix', conda_env_prefix,
            'python=3.12'
        ], check=True)
        
        subprocess.run([
            CONDA_BIN, 'run', '--prefix', conda_env_prefix,
            'pip', 'install', 'ipykernel'
        ], check=True)

        kernel_prefix = f'{user_dir}/.local'
        subprocess.run([
            CONDA_BIN, 'run', '--prefix', conda_env_prefix,
            'python', '-m', 'ipykernel', 'install',
            '--prefix', kernel_prefix,
            '--name', username,
            '--display-name', f'Python ({username})'
        ], check=True)

        # 递归修正 conda 和 kernel specs 目录权限，以防以 root 身份创建时发生权限溢出
        conda_dir = f'{user_dir}/.conda'
        for root, dirs, files in os.walk(conda_dir):
            for d in dirs:
                os.chown(os.path.join(root, d), user_info.pw_uid, user_info.pw_gid)
            for f in files:
                os.chown(os.path.join(root, f), user_info.pw_uid, user_info.pw_gid)
        os.chown(conda_dir, user_info.pw_uid, user_info.pw_gid)

        for root, dirs, files in os.walk(kernel_prefix):
            for d in dirs:
                os.chown(os.path.join(root, d), user_info.pw_uid, user_info.pw_gid)
            for f in files:
                os.chown(os.path.join(root, f), user_info.pw_uid, user_info.pw_gid)
        os.chown(kernel_prefix, user_info.pw_uid, user_info.pw_gid)
        spawner.log.info(f'Conda env and kernel ready for user: {username}')

    spawner.log.info(f'User environment setup completed: {user_dir}')

c.Spawner.pre_spawn_hook = ensure_system_user

# ═══════════════════════════════════════════════════════
# 反向代理集成与 Cookie 域配置
# ═══════════════════════════════════════════════════════
c.JupyterHub.tornado_settings = {
    'cookie_options': {'domain': '.x.com'},
    'headers': {
        'Content-Security-Policy': "frame-ancestors 'self'"
    }
}