99 lines
3.7 KiB
Python
99 lines
3.7 KiB
Python
# myapp/auth/decorators.py
|
||
from functools import wraps
|
||
from typing import Callable, Optional, Union, Iterable
|
||
from flask import session, redirect, url_for, current_app
|
||
|
||
RoleType = Union[str, Iterable[str]]
|
||
RoleChecker = Callable[[str, RoleType], bool] # (username, required_roles) -> bool
|
||
|
||
def _get_deps():
|
||
"""从 app 注入点拿依赖，避免循环导入。"""
|
||
app = current_app
|
||
uuid2username = (
|
||
app.config.get("UUID2USERNAME")
|
||
or app.extensions.get("uuid2username")
|
||
)
|
||
users_list = (
|
||
app.config.get("USERS_LIST")
|
||
or app.extensions.get("users_list")
|
||
)
|
||
role_checker: Optional[RoleChecker] = (
|
||
app.config.get("ROLE_CHECKER")
|
||
or app.extensions.get("role_checker")
|
||
)
|
||
return uuid2username, users_list, role_checker
|
||
|
||
def require_role(
|
||
view=None,
|
||
*,
|
||
# 未登录时跳转到哪个 endpoint
|
||
login_endpoint: str = "auth.login",
|
||
# 角色要求：None/空 -> 仅需登录；"teacher" 或 ["teacher", "admin"] -> 需具备其中之一
|
||
roles: Optional[RoleType] = None,
|
||
# 可选：自定义角色校验函数（优先级高于内置 users_list 适配）
|
||
checker: Optional[RoleChecker] = None,
|
||
):
|
||
"""
|
||
用法：
|
||
1) 仅需登录（等价于 require_user）：
|
||
@require_role
|
||
2) 需要 teacher 角色（等价于 require_teacher）：
|
||
@require_role(roles="teacher")
|
||
3) 需要多个角色之一：
|
||
@require_role(roles=["teacher", "admin"])
|
||
4) 自定义校验器（例如 RBAC/权限码）：
|
||
@require_role(roles="teacher", checker=my_checker)
|
||
|
||
也支持指定登录端点：
|
||
@require_role(login_endpoint="login")
|
||
"""
|
||
def decorator(func):
|
||
@wraps(func)
|
||
def wrapper(*args, **kwargs):
|
||
uuid2username, users_list, injected_checker = _get_deps()
|
||
|
||
user_id = session.get("user_id")
|
||
if not user_id or uuid2username is None or user_id not in uuid2username:
|
||
return redirect(url_for(login_endpoint))
|
||
|
||
# 仅需登录
|
||
if not roles:
|
||
return func(*args, **kwargs)
|
||
|
||
username = uuid2username[user_id]
|
||
|
||
# 选择校验器优先级：参数 checker > app 注入的 role_checker > 内置 users_list 适配
|
||
effective_checker = checker or injected_checker
|
||
|
||
# 内置对 users_list 的向后兼容（你的现有接口）
|
||
if effective_checker is None:
|
||
def builtin_checker(u: str, required: RoleType) -> bool:
|
||
# 支持 teacher 场景（原先的 get_user_is_teacher）
|
||
def has_teacher(u_: str) -> bool:
|
||
if users_list is None:
|
||
return False
|
||
getter = getattr(users_list, "get_user_is_teacher", None)
|
||
return bool(getter and getter(u_) is True)
|
||
|
||
def match_one(required_role: str) -> bool:
|
||
if required_role == "teacher":
|
||
return has_teacher(u)
|
||
# 你也可以在此扩展更多内置角色判断
|
||
# 比如 get_user_is_admin / get_user_roles 等
|
||
# 未知角色默认 False
|
||
return False
|
||
|
||
if isinstance(required, str):
|
||
return match_one(required)
|
||
return any(match_one(r) for r in required)
|
||
|
||
effective_checker = builtin_checker
|
||
|
||
ok = effective_checker(username, roles)
|
||
if not ok:
|
||
return redirect(url_for(login_endpoint))
|
||
|
||
return func(*args, **kwargs)
|
||
return wrapper
|
||
return decorator(view) if callable(view) else decorator
|