WonderQ-Project/WonderQ-Admin/app/rbac.py

from __future__ import annotations

from collections.abc import Iterable
from typing import Any

from sqlalchemy import false, or_, select
from sqlalchemy.orm import Session

from .models import AdminDepartment, AdminMenu, AdminRole, AdminRoleDepartment, AdminRoleMenu, AdminUser, AdminUserDepartment, AdminUserRole
from .redis_session import AdminSessionStore, RedisUnavailableError


DATA_SCOPE_VALUES = {"all", "dept", "dept_and_children", "custom_dept", "self"}


def build_menu_tree(menus: Iterable[AdminMenu], visible_only: bool = True) -> list[dict[str, Any]]:
    active_menus = [
        menu
        for menu in menus
        if not visible_only or (menu.isActive and (menu.type == "button" or menu.isVisible))
    ]
    nodes: dict[str, dict[str, Any]] = {
        menu.id: {
            "id": menu.id,
            "parentId": menu.parentId,
            "name": menu.name,
            "type": menu.type,
            "path": menu.path,
            "component": menu.componentKey,
            "permission": menu.permissionCode,
            "icon": menu.icon,
            "children": [],
            "sortOrder": menu.sortOrder,
            "isVisible": menu.isVisible,
            "isActive": menu.isActive,
        }
        for menu in active_menus
    }
    roots: list[dict[str, Any]] = []
    for menu in sorted(active_menus, key=lambda item: (item.sortOrder, item.id)):
        node = nodes[menu.id]
        if menu.parentId and menu.parentId in nodes:
            nodes[menu.parentId]["children"].append(node)
        else:
            roots.append(node)
    return roots


def flatten_permission_codes(menus: Iterable[AdminMenu]) -> set[str]:
    return {menu.permissionCode for menu in menus if menu.isActive and menu.permissionCode}


def get_admin_roles(user: AdminUser, db: Session) -> list[AdminRole]:
    if user.role in {"admin", "super_admin"}:
        return list(db.scalars(select(AdminRole).where(AdminRole.isActive).order_by(AdminRole.code.asc())).all())
    role_ids = db.scalars(select(AdminUserRole.roleId).where(AdminUserRole.userId == user.id)).all()
    if not role_ids:
        return []
    return list(db.scalars(select(AdminRole).where(AdminRole.id.in_(role_ids), AdminRole.isActive)).all())


def get_admin_menus(user: AdminUser, db: Session) -> list[AdminMenu]:
    if user.role in {"admin", "super_admin"}:
        return list(db.scalars(select(AdminMenu).where(AdminMenu.isActive).order_by(AdminMenu.sortOrder.asc(), AdminMenu.id.asc())).all())
    role_ids = db.scalars(select(AdminUserRole.roleId).where(AdminUserRole.userId == user.id)).all()
    if not role_ids:
        return []
    return list(
        db.scalars(
            select(AdminMenu)
            .join(AdminRoleMenu, AdminRoleMenu.menuId == AdminMenu.id)
            .where(AdminRoleMenu.roleId.in_(role_ids), AdminMenu.isActive)
            .distinct()
            .order_by(AdminMenu.sortOrder.asc(), AdminMenu.id.asc())
        ).all()
    )


def _build_permission_context(user: AdminUser, db: Session) -> dict[str, Any]:
    roles = get_admin_roles(user, db)
    menus = get_admin_menus(user, db)
    permissions = flatten_permission_codes(menus)
    if user.role in {"admin", "super_admin"}:
        permissions.add("*")
    user_dept_ids = set(db.scalars(select(AdminUserDepartment.deptId).where(AdminUserDepartment.userId == user.id)).all())
    return {
        "roles": [{"id": role.id, "name": role.name, "code": role.code, "dataScope": role.dataScope} for role in roles],
        "permissions": sorted(permissions),
        "menus": build_menu_tree(menus),
        "dataScopes": sorted({role.dataScope for role in roles if role.dataScope in DATA_SCOPE_VALUES} or {"all" if user.role in {"admin", "super_admin"} else "self"}),
        "deptIds": sorted(user_dept_ids),
    }


def build_admin_permission_context(user: AdminUser, db: Session, store: AdminSessionStore | None = None) -> dict[str, Any]:
    if store:
        try:
            cached = store.get_permission_context(user.id)
            if cached is not None:
                return cached
        except RedisUnavailableError:
            # 权限缓存不可用时重新读取数据库；绝不将缓存故障降级为放行权限。
            pass
    context = _build_permission_context(user, db)
    if store:
        try:
            from .config import get_settings

            store.set_permission_context(user.id, context, get_settings().admin_permission_cache_seconds)
        except RedisUnavailableError:
            pass
    return context


def user_has_permission(user: AdminUser, db: Session, permission: str) -> bool:
    if user.role in {"admin", "super_admin"}:
        return True
    return permission in set(build_admin_permission_context(user, db)["permissions"])


def require_admin_permission(permission: str):
    from fastapi import Depends, HTTPException

    from .auth import require_admin
    from .database import get_db

    def checker(user: AdminUser = Depends(require_admin), db: Session = Depends(get_db)) -> AdminUser:
        if not user_has_permission(user, db, permission):
            raise HTTPException(status_code=403, detail="没有执行该操作的权限")
        return user

    return checker


def _descendant_department_ids(db: Session, roots: set[str]) -> set[str]:
    if not roots:
        return set()
    rows = db.execute(select(AdminDepartment.id, AdminDepartment.parentId)).all()
    children: dict[str | None, set[str]] = {}
    for department_id, parent_id in rows:
        children.setdefault(parent_id, set()).add(department_id)
    result = set(roots)
    pending = list(roots)
    while pending:
        parent_id = pending.pop()
        for child_id in children.get(parent_id, set()):
            if child_id not in result:
                result.add(child_id)
                pending.append(child_id)
    return result


def accessible_scope(user: AdminUser, db: Session) -> tuple[set[str] | None, bool]:
    """Return (department ids, include self-created rows). None department ids means all."""
    if user.role in {"admin", "super_admin"}:
        return None, True
    roles = get_admin_roles(user, db)
    if any(role.dataScope == "all" for role in roles):
        return None, True
    user_dept_ids = set(db.scalars(select(AdminUserDepartment.deptId).where(AdminUserDepartment.userId == user.id)).all())
    role_ids = {role.id for role in roles}
    custom_dept_ids = set(
        db.scalars(select(AdminRoleDepartment.deptId).where(AdminRoleDepartment.roleId.in_(role_ids))).all()
    ) if role_ids else set()
    dept_ids: set[str] = set()
    include_self = False
    for role in roles:
        if role.dataScope == "dept":
            dept_ids.update(user_dept_ids)
        elif role.dataScope == "dept_and_children":
            dept_ids.update(_descendant_department_ids(db, user_dept_ids))
        elif role.dataScope == "custom_dept":
            dept_ids.update(custom_dept_ids)
        elif role.dataScope == "self":
            include_self = True
    return dept_ids, include_self


def apply_data_scope(stmt, model, user: AdminUser, db: Session):
    """Apply the five RBAC data scopes to an owned business model query."""
    dept_ids, include_self = accessible_scope(user, db)
    if dept_ids is None:
        return stmt
    clauses = []
    dept_column = getattr(model, "deptId", None)
    creator_column = getattr(model, "createdById", None)
    if dept_column is not None and dept_ids:
        clauses.append(dept_column.in_(dept_ids))
    if include_self and creator_column is not None:
        clauses.append(creator_column == user.id)
    return stmt.where(or_(*clauses) if clauses else false())


def is_within_data_scope(item, user: AdminUser, db: Session) -> bool:
    dept_ids, include_self = accessible_scope(user, db)
    if dept_ids is None:
        return True
    item_dept_id = getattr(item, "deptId", None)
    item_creator_id = getattr(item, "createdById", None)
    return (item_dept_id is not None and item_dept_id in dept_ids) or (include_self and item_creator_id == user.id)