sa-token-python-core 0.1.1__py3-none-any.whl

sa_token/exception.py

@@ -0,0 +1,143 @@
+"""sa-token 异常体系。
+
+核心层只抛这些异常，框架适配层负责把它们翻译成各自的 HTTP 响应，
+禁止在适配层重新发明错误码。
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+__all__ = [
+    "NotLoginType",
+    "SaTokenException",
+    "SaTokenNotInitializedException",
+    "NotLoginException",
+    "NotPermissionException",
+    "NotRoleException",
+    "DisableException",
+    "NotSafeException",
+    "SecurityException",
+]
+
+
+class NotLoginType(str, Enum):
+    """未登录的细分原因，用于让调用方区分「没带 token」和「被踢下线」。"""
+
+    NOT_TOKEN = "NOT_TOKEN"
+    INVALID_TOKEN = "INVALID_TOKEN"
+    TOKEN_TIMEOUT = "TOKEN_TIMEOUT"
+    TOKEN_FREEZE = "TOKEN_FREEZE"
+    BE_REPLACED = "BE_REPLACED"
+    KICK_OUT = "KICK_OUT"
+
+
+_NOT_LOGIN_MESSAGES: dict[NotLoginType, str] = {
+    NotLoginType.NOT_TOKEN: "未提供 token",
+    NotLoginType.INVALID_TOKEN: "token 无效",
+    NotLoginType.TOKEN_TIMEOUT: "token 已过期",
+    NotLoginType.TOKEN_FREEZE: "token 已被冻结（长时间未活跃）",
+    NotLoginType.BE_REPLACED: "token 已被顶下线",
+    NotLoginType.KICK_OUT: "token 已被踢下线",
+}
+
+
+class SaTokenException(Exception):
+    """所有 sa-token 异常的基类。"""
+
+    http_status = 500
+
+    def __init__(self, message: str, *, login_type: str = "login") -> None:
+        super().__init__(message)
+        self.message = message
+        self.login_type = login_type
+
+
+class SaTokenNotInitializedException(SaTokenException):
+    """在调用 StpUtil 之前没有构建 Manager。"""
+
+    def __init__(self) -> None:
+        super().__init__(
+            "sa-token 尚未初始化，请先调用 SaToken.builder()...build() 或 set_manager()"
+        )
+
+
+class NotLoginException(SaTokenException):
+    """未登录、token 失效、被踢、被顶、被冻结。"""
+
+    http_status = 401
+
+    def __init__(
+        self,
+        not_login_type: NotLoginType,
+        *,
+        login_type: str = "login",
+        token: str | None = None,
+    ) -> None:
+        super().__init__(_NOT_LOGIN_MESSAGES[not_login_type], login_type=login_type)
+        self.type = not_login_type
+        self.token = token
+
+
+class NotPermissionException(SaTokenException):
+    """权限不足。"""
+
+    http_status = 403
+
+    def __init__(self, permission: str, *, login_type: str = "login") -> None:
+        super().__init__(f"缺少权限：{permission}", login_type=login_type)
+        self.permission = permission
+
+
+class NotRoleException(SaTokenException):
+    """角色不足。"""
+
+    http_status = 403
+
+    def __init__(self, role: str, *, login_type: str = "login") -> None:
+        super().__init__(f"缺少角色：{role}", login_type=login_type)
+        self.role = role
+
+
+class DisableException(SaTokenException):
+    """账号（或账号的某项服务）被封禁。"""
+
+    http_status = 403
+
+    def __init__(
+        self,
+        login_id: str,
+        service: str,
+        level: int,
+        remaining: int,
+        *,
+        login_type: str = "login",
+    ) -> None:
+        super().__init__(
+            f"账号 {login_id} 的服务 {service} 已被封禁（等级 {level}，剩余 {remaining} 秒）",
+            login_type=login_type,
+        )
+        self.login_id = login_id
+        self.service = service
+        self.level = level
+        self.remaining = remaining
+
+
+class NotSafeException(SaTokenException):
+    """二级认证未通过。"""
+
+    http_status = 403
+
+    def __init__(self, business: str, *, login_type: str = "login") -> None:
+        super().__init__(f"二级认证未通过：{business}", login_type=login_type)
+        self.business = business
+
+
+class SecurityException(SaTokenException):
+    """Nonce、Refresh Token、Temp Token 等安全流程错误。"""
+
+    http_status = 400
+
+    def __init__(self, code: str, message: str, *, login_type: str = "login") -> None:
+        super().__init__(message, login_type=login_type)
+        self.code = code

sa_token/integration/__init__.py

@@ -0,0 +1,10 @@
+"""框架绑定。
+
+每个子模块都只做「Request → HttpContext、调用统一管道、异常翻译」三件事，
+因此各框架的行为天然一致。这些模块依赖对应的可选 extra，
+不在这里做顶层导入，避免只装了核心包的用户 import 失败。
+"""
+
+from __future__ import annotations
+
+__all__ = ["fastapi", "starlette", "flask", "django"]

sa_token/integration/django.py

@@ -0,0 +1,131 @@
+"""Django 适配（同步中间件 + 装饰器）。
+
+Django 自带一套完整的 Session/User 体系，因此这里的定位是：
+给**纯 API 项目**（DRF / 无模板）提供与 FastAPI、Flask 一致的 Token 鉴权，
+而不是替换 ``django.contrib.auth``。
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from ..adapter.http import HttpContext
+from ..adapter.path import PathAuthConfig
+from ..adapter.pipeline import build_rule, resolve_token, run_auth_flow, run_path_auth
+from ..exception import SaTokenException
+from ..permission import MatchMode
+from ..stp_util import get_manager
+from ..sync import run_sync
+
+__all__ = [
+    "DjangoHttpContext",
+    "SaTokenDjangoMiddleware",
+    "check_login",
+    "check_permission",
+    "check_role",
+    "to_json_response",
+]
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+#: 由项目在 settings 里赋值，用于开启路径鉴权。
+PATH_AUTH: PathAuthConfig | None = None
+
+
+class DjangoHttpContext(HttpContext):
+    def __init__(self, request: Any) -> None:
+        self._request = request
+        self.state: dict[str, Any] = {}
+
+    def get_header(self, name: str) -> str | None:
+        return self._request.headers.get(name)
+
+    def get_cookie(self, name: str) -> str | None:
+        return self._request.COOKIES.get(name)
+
+    def get_query(self, name: str) -> str | None:
+        return self._request.GET.get(name)
+
+    def get_path(self) -> str:
+        return self._request.path
+
+    def get_method(self) -> str:
+        return self._request.method
+
+
+def to_json_response(exc: SaTokenException) -> Any:
+    from django.http import JsonResponse
+
+    payload: dict[str, Any] = {
+        "code": exc.http_status,
+        "message": exc.message,
+        "error": type(exc).__name__,
+    }
+    detail_type = getattr(exc, "type", None)
+    if detail_type is not None:
+        payload["type"] = detail_type.value
+    return JsonResponse(payload, status=exc.http_status)
+
+
+class SaTokenDjangoMiddleware:
+    """加入 ``MIDDLEWARE`` 即可。
+
+    默认只解析 token 并挂到 ``request.sa_token`` / ``request.sa_login_id``；
+    把 :data:`PATH_AUTH` 设为规则表后才执行强制鉴权。
+    """
+
+    def __init__(self, get_response: Callable[[Any], Any]) -> None:
+        self._get_response = get_response
+
+    def __call__(self, request: Any) -> Any:
+        ctx = DjangoHttpContext(request)
+        manager = get_manager()
+        try:
+            if PATH_AUTH is None:
+                resolve_token(ctx, manager)
+            else:
+                run_sync(run_path_auth(ctx, manager, PATH_AUTH))
+        except SaTokenException as exc:
+            return to_json_response(exc)
+
+        request.sa_token = ctx.state.get("stp_token")
+        request.sa_login_id = ctx.state.get("stp_login_id")
+        return self._get_response(request)
+
+    def process_exception(self, request: Any, exception: Exception) -> Any:
+        if isinstance(exception, SaTokenException):
+            return to_json_response(exception)
+        return None
+
+
+def _guard(rule_factory: Callable[[], Any]) -> Callable[[_F], _F]:
+    def decorator(view: _F) -> _F:
+        @functools.wraps(view)
+        def wrapper(request: Any, *args: Any, **kwargs: Any) -> Any:
+            ctx = DjangoHttpContext(request)
+            try:
+                result = run_sync(run_auth_flow(ctx, get_manager(), rule_factory()))
+            except SaTokenException as exc:
+                return to_json_response(exc)
+            request.sa_login_id = result.login_id
+            request.sa_token = result.token
+            return view(request, *args, **kwargs)
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def check_login(view: _F) -> _F:
+    """标准注解：``@check_login``。"""
+    return _guard(build_rule)(view)
+
+
+def check_permission(*permissions: str, mode: MatchMode = "OR") -> Callable[[_F], _F]:
+    return _guard(lambda: build_rule(permissions=list(permissions), mode=mode))
+
+
+def check_role(*roles: str, mode: MatchMode = "OR") -> Callable[[_F], _F]:
+    return _guard(lambda: build_rule(roles=list(roles), mode=mode))

sa_token/integration/fastapi.py

@@ -0,0 +1,315 @@
+"""FastAPI 适配。
+
+各框架的标准鉴权写法都是注解：``@sa.check_login`` / ``@sa.check_permission``。
+:class:`SaTokenFastAPI` 提供这一套，与 Flask、Django 同一套 ``StpLogic``。
+
+FastAPI 作为主要适配框架，额外提供 ``Depends`` / ``Annotated``（``LoginId``、
+``BearerLoginId`` 等），方便按 FastAPI 习惯写依赖注入。
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from collections.abc import Callable
+from typing import Annotated, Any, TypeVar
+
+try:
+    from fastapi import Depends, FastAPI, Request, Response, WebSocket
+    from fastapi.security import (  # pyright: ignore[reportMissingImports]
+        HTTPAuthorizationCredentials,
+        HTTPBearer,
+    )
+except ImportError as exc:  # pragma: no cover - 依赖缺失路径
+    raise ImportError(
+        '该模块需要 fastapi，请执行：pip install "sa-token-python-core[fastapi]"'
+    ) from exc
+
+from ..adapter.path import PathAuthConfig
+from ..adapter.pipeline import build_rule, run_auth_flow
+from ..context import get_current_login_id, get_current_token
+from ..exception import NotLoginException, NotLoginType
+from ..permission import MatchMode
+from ..stp_util import get_manager
+from .starlette import (
+    SaTokenMiddleware,
+    StarletteHttpContext,
+    check_disable,
+    check_login,
+    check_permission,
+    check_role,
+    check_safe,
+    current_login_id,
+    current_login_id_or_none,
+    current_token,
+    install_exception_handlers,
+    sa_token_exception_handler,
+)
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+__all__ = [
+    "SaTokenFastAPI",
+    "SaTokenMiddleware",
+    "StarletteHttpContext",
+    "install_exception_handlers",
+    "sa_token_exception_handler",
+    "check_login",
+    "check_permission",
+    "check_role",
+    "check_disable",
+    "check_safe",
+    "current_login_id",
+    "current_login_id_or_none",
+    "current_token",
+    "LoginId",
+    "OptionalLoginId",
+    "TokenValue",
+    "BearerLoginId",
+    "set_token_cookie",
+    "delete_token_cookie",
+    "FastAPIWebSocketContext",
+    "authenticate_websocket",
+]
+
+
+def _extract_request(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Request:
+    request = kwargs.get("request")
+    if isinstance(request, Request):
+        return request
+    for value in args:
+        if isinstance(value, Request):
+            return value
+    for value in kwargs.values():
+        if isinstance(value, Request):
+            return value
+    raise RuntimeError("未找到 Request")
+
+
+def _endpoint_signature(view: Callable[..., Any]) -> inspect.Signature:
+    signature = inspect.signature(view)
+    if "request" in signature.parameters:
+        return signature
+    request_param = inspect.Parameter(
+        "request",
+        inspect.Parameter.KEYWORD_ONLY,
+        annotation=Request,
+    )
+    return signature.replace(parameters=[*signature.parameters.values(), request_param])
+
+
+class SaTokenFastAPI:
+    """FastAPI 的标准注解鉴权，与 Flask / Django 同一套语义。
+
+    用法::
+
+        app = FastAPI()
+        sa = SaTokenFastAPI(app)
+
+        @app.get("/user")
+        @sa.check_login
+        async def user_info():
+            return {"id": sa.login_id()}
+    """
+
+    def __init__(
+        self,
+        app: FastAPI | None = None,
+        *,
+        path_auth: PathAuthConfig | None = None,
+        login_type: str = "login",
+    ) -> None:
+        self.path_auth = path_auth
+        self.login_type = login_type
+        if app is not None:
+            self.init_app(app)
+
+    def init_app(self, app: FastAPI) -> None:
+        app.add_middleware(
+            SaTokenMiddleware,
+            path_auth=self.path_auth,
+            login_type=self.login_type,
+        )
+        install_exception_handlers(app)
+
+    def login_id(self) -> str:
+        login_id = get_current_login_id()
+        if login_id:
+            return login_id
+        raise NotLoginException(NotLoginType.NOT_TOKEN, login_type=self.login_type)
+
+    def token(self) -> str | None:
+        return get_current_token()
+
+    def _guard(
+        self,
+        rule_factory: Callable[[], Any],
+        extra: Callable[[Any], Any] | None = None,
+    ) -> Callable[[_F], _F]:
+        def decorator(view: _F) -> _F:
+            had_request = "request" in inspect.signature(view).parameters
+
+            @functools.wraps(view)
+            async def wrapper(*args: Any, **kwargs: Any) -> Any:
+                request = _extract_request(args, kwargs)
+                ctx = StarletteHttpContext(request)
+                result = await run_auth_flow(
+                    ctx,
+                    get_manager(),
+                    rule_factory(),
+                    login_type=self.login_type,
+                )
+                request.state.sa_login_id = result.login_id
+                request.state.sa_token = result.token
+                if extra is not None:
+                    extra_result = extra(result)
+                    if inspect.isawaitable(extra_result):
+                        await extra_result
+                call_kwargs = kwargs
+                call_args = args
+                if not had_request:
+                    call_kwargs = dict(kwargs)
+                    call_kwargs.pop("request", None)
+                    call_args = tuple(item for item in args if not isinstance(item, Request))
+                outcome = view(*call_args, **call_kwargs)
+                if inspect.isawaitable(outcome):
+                    return await outcome
+                return outcome
+
+            wrapper.__signature__ = _endpoint_signature(view)
+            return wrapper  # type: ignore[return-value]
+
+        return decorator
+
+    @property
+    def check_login(self) -> Callable[[_F], _F]:
+        """标准注解：``@sa.check_login``。"""
+        return self._guard(build_rule)
+
+    def check_permission(self, *permissions: str, mode: MatchMode = "OR") -> Callable[[_F], _F]:
+        """标准注解：``@sa.check_permission("order:delete")``。"""
+        return self._guard(lambda: build_rule(permissions=list(permissions), mode=mode))
+
+    def check_role(self, *roles: str, mode: MatchMode = "OR") -> Callable[[_F], _F]:
+        return self._guard(lambda: build_rule(roles=list(roles), mode=mode))
+
+    def check_safe(self, business: str) -> Callable[[_F], _F]:
+        async def extra(result: Any) -> None:
+            await get_manager().stp(self.login_type).check_safe(result.token, business)
+
+        return self._guard(build_rule, extra=extra)
+
+    def check_disable(self, service: str = "login", level: int = 1) -> Callable[[_F], _F]:
+        async def extra(result: Any) -> None:
+            await get_manager().stp(self.login_type).check_disable(
+                result.login_id,
+                service=service,
+                level=level,
+            )
+
+        return self._guard(build_rule, extra=extra)
+
+
+# FastAPI 惯用扩展：Depends / Annotated。
+
+LoginId = Annotated[str, Depends(current_login_id)]
+OptionalLoginId = Annotated[str | None, Depends(current_login_id_or_none)]
+TokenValue = Annotated[str | None, Depends(current_token)]
+
+_bearer_scheme = HTTPBearer(auto_error=False)
+
+
+async def _bearer_login_id(
+    request: Request,
+    credentials: HTTPAuthorizationCredentials | None = Depends(_bearer_scheme),
+) -> str:
+    """Depends 扩展：把 Bearer 写进 OpenAPI，读取逻辑与注解鉴权相同。"""
+    del credentials
+    return await current_login_id(request)
+
+
+BearerLoginId = Annotated[str, Depends(_bearer_login_id)]
+
+
+def set_token_cookie(
+    response: Response,
+    token: str,
+    *,
+    max_age: int | None = None,
+) -> None:
+    """按全局 Cookie 配置把登录 token 写入响应。
+
+    这是显式函数而不是自动拦截 ``StpUtil.login``：核心层没有 Response，
+    自动写入会破坏框架无关边界。
+    """
+    from ..stp_util import get_manager
+
+    config = get_manager().config
+    if not config.is_write_cookie:
+        return
+    resolved_max_age = config.timeout if max_age is None else max_age
+    response.set_cookie(
+        key=config.token_name,
+        value=token,
+        max_age=None if resolved_max_age < 0 else resolved_max_age,
+        path=config.cookie_path,
+        domain=config.cookie_domain,
+        secure=config.cookie_secure,
+        httponly=config.cookie_http_only,
+        samesite=config.cookie_same_site,
+    )
+
+
+def delete_token_cookie(response: Response) -> None:
+    """删除登录 Cookie，参数与写入时保持一致。"""
+    from ..stp_util import get_manager
+
+    config = get_manager().config
+    response.delete_cookie(
+        key=config.token_name,
+        path=config.cookie_path,
+        domain=config.cookie_domain,
+        secure=config.cookie_secure,
+        httponly=config.cookie_http_only,
+        samesite=config.cookie_same_site,
+    )
+
+
+class FastAPIWebSocketContext:
+    """把 FastAPI / Starlette WebSocket 包装成核心 HttpContext。"""
+
+    def __init__(self, websocket: WebSocket) -> None:
+        self.websocket = websocket
+        self.state: dict[str, Any] = {}
+
+    def get_header(self, name: str) -> str | None:
+        return self.websocket.headers.get(name)
+
+    def get_cookie(self, name: str) -> str | None:
+        return self.websocket.cookies.get(name)
+
+    def get_query(self, name: str) -> str | None:
+        return self.websocket.query_params.get(name)
+
+    def get_path(self) -> str:
+        return self.websocket.url.path
+
+    def get_method(self) -> str:
+        return "WEBSOCKET"
+
+
+async def authenticate_websocket(
+    websocket: WebSocket,
+    *,
+    login_type: str = "login",
+) -> tuple[str, str]:
+    """在 ``accept`` 之前校验 WebSocket，返回 ``(login_id, token)``。"""
+    from ..online import WebSocketAuthenticator
+    from ..stp_util import get_manager
+
+    context = FastAPIWebSocketContext(websocket)
+    login_id = await WebSocketAuthenticator(
+        get_manager(), login_type=login_type
+    ).authenticate(context)
+    token = context.state["stp_token"]
+    return login_id, token