models_user.py

"""
models_user.py — 多用户数据模型
支持 OIDC JIT 自动注册、角色映射、YAML 持久化
"""
import json
import time
import threading
from dataclasses import dataclass, field, asdict
from typing import Optional, List, Dict, Any


# ── 角色常量 ──────────────────────────────────────────
ROLE_ADMIN = "admin"
ROLE_MEMBER = "member"
ROLE_VIEWER = "viewer"
VALID_ROLES = {ROLE_ADMIN, ROLE_MEMBER, ROLE_VIEWER}


@dataclass
class User:
    """单个用户"""
    sub: str                    # OIDC subject（唯一标识）
    email: str = ""
    name: str = ""
    groups: List[str] = field(default_factory=list)
    role: str = ROLE_MEMBER
    is_active: bool = True
    created_at: float = field(default_factory=time.time)
    last_login: float = 0

    def to_dict(self) -> dict:
        return asdict(self)

    def to_public_dict(self) -> dict:
        """对外暴露（不含敏感字段）"""
        return {
            "sub": self.sub,
            "email": self.email,
            "name": self.name,
            "groups": self.groups,
            "role": self.role,
            "is_active": self.is_active,
            "created_at": self.created_at,
            "last_login": self.last_login,
        }


class UserStore:
    """
    用户存储（内存 + YAML 持久化）
    线程安全
    """

    def __init__(self, path: str = "_users.yaml"):
        self.path = path
        self._users: Dict[str, User] = {}
        self._lock = threading.RLock()

    # ── 持久化 ──────────────────────────────────────
    def load(self):
        import os
        if not os.path.exists(self.path):
            return
        try:
            import yaml
            with open(self.path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f) or {}
            with self._lock:
                self._users.clear()
                for item in data.get("users", []):
                    u = User(
                        sub=item["sub"],
                        email=item.get("email", ""),
                        name=item.get("name", ""),
                        groups=item.get("groups", []),
                        role=item.get("role", ROLE_MEMBER),
                        is_active=item.get("is_active", True),
                        created_at=item.get("created_at", time.time()),
                        last_login=item.get("last_login", 0),
                    )
                    self._users[u.sub] = u
        except Exception as e:
            print(f"[UserStore] Load failed: {e}")

    def save(self):
        import yaml
        with self._lock:
            data = {
                "users": [u.to_dict() for u in self._users.values()]
            }
        try:
            with open(self.path, "w", encoding="utf-8") as f:
                yaml.safe_dump(data, f, allow_unicode=True, sort_keys=False)
        except Exception as e:
            print(f"[UserStore] Save failed: {e}")

    # ── CRUD ────────────────────────────────────────
    def get_or_create(
        self,
        sub: str,
        email: str = "",
        name: str = "",
        groups: Optional[List[str]] = None,
        role_mapping: Optional[dict] = None,
    ) -> User:
        """JIT 自动注册：有就更新，没有就创建"""
        with self._lock:
            now = time.time()
            if sub in self._users:
                u = self._users[sub]
                if email:   u.email = email
                if name:    u.name = name
                if groups is not None: u.groups = groups
                u.last_login = now
                # 如果角色是 member 且 groups 变了，重新映射
                if role_mapping and groups is not None:
                    mapped = _role_from_groups(groups, role_mapping)
                    if mapped != ROLE_VIEWER:  # viewer 是默认降级，不覆盖
                        u.role = mapped
                self.save()
                return u

            # 新建
            groups = groups or []
            role = ROLE_MEMBER
            if role_mapping:
                role = _role_from_groups(groups, role_mapping)

            u = User(
                sub=sub,
                email=email,
                name=name,
                groups=groups,
                role=role,
                is_active=True,
                created_at=now,
                last_login=now,
            )
            self._users[sub] = u
            self.save()
            return u

    def find_by_sub(self, sub: str) -> Optional[User]:
        with self._lock:
            return self._users.get(sub)

    def list_users(self) -> List[User]:
        with self._lock:
            return list(self._users.values())

    def set_role(self, sub: str, role: str) -> bool:
        if role not in VALID_ROLES:
            return False
        with self._lock:
            u = self._users.get(sub)
            if not u:
                return False
            u.role = role
            self.save()
            return True

    def set_active(self, sub: str, active: bool) -> bool:
        with self._lock:
            u = self._users.get(sub)
            if not u:
                return False
            u.is_active = active
            self.save()
            return True

    def count(self) -> int:
        with self._lock:
            return len(self._users)


# ── 角色映射 ──────────────────────────────────────────
def _role_from_groups(groups: List[str], mapping: dict) -> str:
    """根据 Authentik groups claim 映射本地角色"""
    admin_groups = set(mapping.get("admin_groups", ["fw-admin", "admin"]))
    viewer_groups = set(mapping.get("viewer_groups", ["fw-viewer"]))
    default_role = mapping.get("default_role", ROLE_MEMBER)

    gset = set(groups)
    if gset & admin_groups:
        return ROLE_ADMIN
    if gset & viewer_groups:
        return ROLE_VIEWER
    return default_role