88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
import os, grp, pwd, shutil, subprocess, configparser
|
||
from pathlib import Path
|
||
from jupyterhub.spawner import Spawner
|
||
|
||
# host端的仓库目录
|
||
SHARED_ROOT = Path("/home/zhangyi/jupyter-collection")
|
||
# jyputerhub用户home目录
|
||
USER_HOME = Path("/home/jupyter-{username}")
|
||
# 分组信息路径
|
||
CTRL_FILE = SHARED_ROOT / ".hub" / "resource_map.ini"
|
||
|
||
# ========== 2. 读控制表 ==========
|
||
def load_ctrl():
|
||
"""返回 组->资源 和 组->用户 两个字典"""
|
||
cfg = configparser.ConfigParser()
|
||
cfg.read(CTRL_FILE)
|
||
res_map = {g: [r.strip() for r in v.split(",")]
|
||
for g, v in cfg.items("map")}
|
||
usr_map = {g: [u.strip() for u in v.split(",")]
|
||
for g, v in cfg.items("users")}
|
||
return res_map, usr_map
|
||
|
||
# ========== 3. post_auth_hook：自动建用户+加组 ==========
|
||
def ensure_user_and_groups(authenticator, handler, authentication):
|
||
username = authentication['name']
|
||
res_map, usr_map = load_ctrl()
|
||
|
||
# 该用户应属的组
|
||
wanted_groups = [g for g, users in usr_map.items() if username in users]
|
||
if not wanted_groups: # 控制文件里没写就什么都不做
|
||
return authentication
|
||
|
||
# 确保组存在
|
||
for g in wanted_groups:
|
||
subprocess.run(["groupadd", "-f", g], check=False)
|
||
|
||
# 若系统账号不存在则创建（Hub 自己会创 home，这里确保 shell 账号）
|
||
#if subprocess.run(["id", "-u", username], capture_output=True).returncode != 0:
|
||
# subprocess.run([
|
||
# "useradd", "-m", "-s", "/bin/bash",
|
||
# "-d", str(USER_HOME).format(username=username), username
|
||
# ], check=False)
|
||
|
||
# 加组（幂等）
|
||
for g in wanted_groups:
|
||
subprocess.run(["usermod", "-aG", g, username], check=False)
|
||
|
||
return authentication
|
||
|
||
c.Authenticator.post_auth_hook = ensure_user_and_groups
|
||
|
||
# ========== 4. pre_spawn_hook：重建干净 home + 只读挂 Git 资源 ==========
|
||
def prepare_user(spawner: Spawner):
|
||
username = spawner.user.name
|
||
user_dir = Path(str(USER_HOME).format(username=username))
|
||
|
||
# 1. 重建空 home（保证无残留）
|
||
if user_dir.exists():
|
||
shutil.rmtree(user_dir)
|
||
user_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||
|
||
# 2. 读控制表
|
||
res_map, _ = load_ctrl()
|
||
|
||
# 3. 取系统组
|
||
try:
|
||
pwnam = pwd.getpwnam(username)
|
||
gid_list = os.getgrouplist(username, pwnam.pw_gid)
|
||
except (KeyError, OSError):
|
||
return
|
||
user_groups = {grp.getgrgid(g).gr_name for g in gid_list}
|
||
|
||
# 4. 挂载该用户该看的资源（只读）
|
||
shared = user_dir / "shared"
|
||
shared.mkdir(exist_ok=True)
|
||
bind_list = [] # 先攒列表
|
||
for grp_name in user_groups:
|
||
for res in res_map.get(grp_name, []):
|
||
res_path = SHARED_ROOT / res
|
||
if res_path.is_dir():
|
||
bind_list.append({
|
||
'source': str(res_path),
|
||
'target': str(shared / res),
|
||
'read-only': True,
|
||
})
|
||
spawner.bind_mounts = bind_list
|
||
|
||
c.Spawner.pre_spawn_hook = prepare_user |