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

sa_token/integration/fastapi_oauth2.py

@@ -0,0 +1,136 @@
+"""OAuth2 的 FastAPI HTTP 端点适配。
+
+协议引擎仍在 :mod:`sa_token.oauth2`；这里仅处理 HTTP 表单/JSON 和标准错误响应。
+授权页的 UI 与「用户是否同意」属于业务，本模块只提供授权成功后的跳转函数。
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from urllib.parse import parse_qsl, urlencode
+
+try:
+    from fastapi import APIRouter, Request
+    from fastapi.responses import JSONResponse, RedirectResponse
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        '该模块需要 fastapi，请执行：pip install "sa-token-python-core[fastapi]"'
+    ) from exc
+
+from ..oauth2 import OAuth2Error, OAuth2Server
+
+__all__ = ["create_oauth2_router", "authorization_redirect"]
+
+
+async def _read_parameters(request: Request) -> dict[str, Any]:
+    content_type = request.headers.get("content-type", "")
+    if "application/json" in content_type:
+        payload = await request.json()
+        return payload if isinstance(payload, dict) else {}
+    body = (await request.body()).decode("utf-8")
+    return dict(parse_qsl(body, keep_blank_values=True))
+
+
+def _scopes(value: Any) -> list[str]:
+    if isinstance(value, list):
+        return [str(item) for item in value]
+    return str(value or "").split()
+
+
+def _oauth_error(error: OAuth2Error) -> JSONResponse:
+    status = 401 if error.error == "invalid_client" else error.http_status
+    headers = {"WWW-Authenticate": "Basic"} if status == 401 else None
+    return JSONResponse(
+        {
+            "error": error.error,
+            "error_description": error.description,
+        },
+        status_code=status,
+        headers=headers,
+    )
+
+
+def create_oauth2_router(
+    server: OAuth2Server,
+    *,
+    prefix: str = "/oauth2",
+) -> APIRouter:
+    """创建 token / introspect / revoke 路由。
+
+    `/token` 同时接受 JSON 与 `application/x-www-form-urlencoded`，无需额外安装
+    `python-multipart`。
+    """
+    router = APIRouter(prefix=prefix, tags=["OAuth2"])
+
+    @router.post("/token")
+    async def token_endpoint(request: Request) -> JSONResponse:
+        parameters = await _read_parameters(request)
+        grant_type = str(parameters.get("grant_type", ""))
+        try:
+            if grant_type == "authorization_code":
+                response = await server.exchange_code_for_token(
+                    code=str(parameters.get("code", "")),
+                    client_id=str(parameters.get("client_id", "")),
+                    client_secret=parameters.get("client_secret"),
+                    redirect_uri=parameters.get("redirect_uri"),
+                    code_verifier=parameters.get("code_verifier"),
+                )
+            elif grant_type == "refresh_token":
+                response = await server.refresh_access_token(
+                    refresh_token=str(parameters.get("refresh_token", "")),
+                    client_id=str(parameters.get("client_id", "")),
+                    client_secret=parameters.get("client_secret"),
+                )
+            elif grant_type == "client_credentials":
+                response = await server.client_credentials_token(
+                    client_id=str(parameters.get("client_id", "")),
+                    client_secret=str(parameters.get("client_secret", "")),
+                    scopes=_scopes(parameters.get("scope")),
+                )
+            else:
+                raise OAuth2Error("unsupported_grant_type", f"不支持的 grant_type: {grant_type}")
+            return JSONResponse(response.to_dict())
+        except OAuth2Error as error:
+            return _oauth_error(error)
+
+    @router.post("/introspect")
+    async def introspect_endpoint(request: Request) -> JSONResponse:
+        parameters = await _read_parameters(request)
+        return JSONResponse(await server.introspect(str(parameters.get("token", ""))))
+
+    @router.post("/revoke")
+    async def revoke_endpoint(request: Request) -> JSONResponse:
+        parameters = await _read_parameters(request)
+        await server.revoke_token(str(parameters.get("token", "")))
+        # RFC 7009：token 不存在时同样返回 200，避免泄露有效性。
+        return JSONResponse({})
+
+    return router
+
+
+async def authorization_redirect(
+    server: OAuth2Server,
+    *,
+    login_id: str,
+    client_id: str,
+    redirect_uri: str,
+    scope: str = "",
+    state: str | None = None,
+    code_challenge: str | None = None,
+    code_challenge_method: str = "S256",
+) -> RedirectResponse:
+    """用户在业务授权页点击同意后，签发授权码并跳回客户端。"""
+    code = await server.create_authorization_code(
+        client_id=client_id,
+        login_id=login_id,
+        redirect_uri=redirect_uri,
+        scopes=_scopes(scope),
+        code_challenge=code_challenge,
+        code_challenge_method=code_challenge_method,
+        state=state,
+    )
+    parameters = {"code": code.code}
+    if state:
+        parameters["state"] = state
+    separator = "&" if "?" in redirect_uri else "?"
+    return RedirectResponse(f"{redirect_uri}{separator}{urlencode(parameters)}")

sa_token/integration/flask.py

@@ -0,0 +1,191 @@
+"""Flask 适配。
+
+Flask 是 WSGI 同步框架，所有调用经由 :mod:`sa_token.sync` 的同步桥接进入
+同一套 ``StpLogic``，不存在第二份鉴权实现。
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from typing import TYPE_CHECKING, 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
+
+if TYPE_CHECKING:  # pragma: no cover - 仅供类型检查
+    from flask import Flask
+
+__all__ = ["FlaskHttpContext", "SaTokenFlask"]
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+
+class FlaskHttpContext(HttpContext):
+    """把 Flask ``request`` 包装成框架无关的上下文。"""
+
+    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.args.get(name)
+
+    def get_path(self) -> str:
+        return self._request.path
+
+    def get_method(self) -> str:
+        return self._request.method
+
+
+class SaTokenFlask:
+    """Flask 的标准注解鉴权，与 FastAPI / Django 同一套语义。
+
+    用法::
+
+        sa = SaTokenFlask(app)
+
+        @app.get("/user")
+        @sa.check_login
+        def user_info():
+            return {"id": sa.login_id()}
+    """
+
+    def __init__(
+        self,
+        app: Flask | 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: Flask) -> None:
+        app.before_request(self._before_request)
+        app.register_error_handler(SaTokenException, self._handle_exception)
+
+    # 请求钩子 --------------------------------------------------------------
+
+    def _context(self) -> FlaskHttpContext:
+        from flask import request
+
+        return FlaskHttpContext(request)
+
+    def _before_request(self) -> Any:
+        from flask import g
+
+        ctx = self._context()
+        manager = get_manager()
+        if self.path_auth is None:
+            resolve_token(ctx, manager)
+        else:
+            run_sync(run_path_auth(ctx, manager, self.path_auth, login_type=self.login_type))
+        g.sa_token = ctx.state.get("stp_token")
+        g.sa_login_id = ctx.state.get("stp_login_id")
+        return None
+
+    def _handle_exception(self, exc: SaTokenException) -> Any:
+        from flask import jsonify
+
+        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 jsonify(payload), exc.http_status
+
+    # 取值 -----------------------------------------------------------------
+
+    def token(self) -> str | None:
+        return resolve_token(self._context(), get_manager())
+
+    def login_id(self) -> str:
+        """当前登录用户；未登录抛异常（会被错误处理器转成 401）。"""
+        from flask import g
+
+        cached = getattr(g, "sa_login_id", None)
+        if cached:
+            return cached
+        result = run_sync(run_auth_flow(self._context(), get_manager(), build_rule()))
+        assert result.login_id is not None
+        g.sa_login_id = result.login_id
+        return result.login_id
+
+    def login_id_or_none(self) -> str | None:
+        return run_sync(get_manager().stp(self.login_type).get_login_id_or_none(self.token()))
+
+    # 装饰器 ----------------------------------------------------------------
+
+    def _guard(self, rule_factory: Callable[[], Any]) -> Callable[[_F], _F]:
+        def decorator(view: _F) -> _F:
+            @functools.wraps(view)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                from flask import g
+
+                result = run_sync(
+                    run_auth_flow(
+                        self._context(),
+                        get_manager(),
+                        rule_factory(),
+                        login_type=self.login_type,
+                    )
+                )
+                g.sa_login_id = result.login_id
+                g.sa_token = result.token
+                return view(*args, **kwargs)
+
+            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]:
+        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]:
+        def decorator(view: _F) -> _F:
+            @functools.wraps(view)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                logic = get_manager().stp(self.login_type)
+                run_sync(logic.check_safe(self.token(), business))
+                return view(*args, **kwargs)
+
+            return wrapper  # type: ignore[return-value]
+
+        return decorator
+
+    def check_disable(self, service: str = "login", level: int = 1) -> Callable[[_F], _F]:
+        def decorator(view: _F) -> _F:
+            @functools.wraps(view)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                logic = get_manager().stp(self.login_type)
+                run_sync(logic.check_disable(self.login_id(), service=service, level=level))
+                return view(*args, **kwargs)
+
+            return wrapper  # type: ignore[return-value]
+
+        return decorator

sa_token/integration/starlette.py

@@ -0,0 +1,227 @@
+"""Starlette / FastAPI 共用的适配层。
+
+职责严格限定在三件事：Request → HttpContext、调用统一管道、异常翻译。
+鉴权逻辑一行都不在这里实现。
+
+注意：这里必须在**运行时**导入 Starlette 的 ``Request``，不能放进
+``TYPE_CHECKING``。FastAPI 依靠运行时类型注解判断依赖项参数，
+拿不到真实类型时会把 ``request`` 当成查询参数，返回 422。
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+try:
+    from starlette.applications import Starlette
+    from starlette.middleware.base import BaseHTTPMiddleware
+    from starlette.requests import Request
+    from starlette.responses import JSONResponse, Response
+except ImportError as exc:  # pragma: no cover - 依赖缺失路径
+    raise ImportError(
+        '该模块需要 starlette，请执行：pip install "sa-token-python-core[fastapi]"'
+    ) from exc
+
+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
+
+__all__ = [
+    "StarletteHttpContext",
+    "SaTokenMiddleware",
+    "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",
+]
+
+
+class StarletteHttpContext(HttpContext):
+    """把 Starlette ``Request`` 包装成框架无关的上下文。"""
+
+    def __init__(self, request: Request) -> 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.query_params.get(name)
+
+    def get_path(self) -> str:
+        return self._request.url.path
+
+    def get_method(self) -> str:
+        return self._request.method
+
+
+def _error_response(exc: SaTokenException) -> Response:
+    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_code=exc.http_status)
+
+
+class SaTokenMiddleware(BaseHTTPMiddleware):
+    """请求级中间件。
+
+    默认只解析 token 并绑定上下文，**不强制登录**：登录接口、健康检查这类
+    公开路由永远存在，默认拦截会逼用户到处写例外。传入 ``path_auth``
+    后才按规则表执行鉴权。
+    """
+
+    def __init__(
+        self,
+        app: Any,
+        *,
+        path_auth: PathAuthConfig | None = None,
+        login_type: str = "login",
+    ) -> None:
+        super().__init__(app)
+        self.path_auth = path_auth
+        self.login_type = login_type
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        ctx = StarletteHttpContext(request)
+        manager = get_manager()
+        try:
+            if self.path_auth is None:
+                resolve_token(ctx, manager)
+            else:
+                await run_path_auth(ctx, manager, self.path_auth, login_type=self.login_type)
+        except SaTokenException as exc:
+            return _error_response(exc)
+
+        request.state.sa_token = ctx.state.get("stp_token")
+        request.state.sa_login_id = ctx.state.get("stp_login_id")
+        return await call_next(request)
+
+
+async def sa_token_exception_handler(request: Request, exc: Exception) -> Response:
+    """把核心异常翻译成 401 / 403 / 500。"""
+    assert isinstance(exc, SaTokenException)
+    return _error_response(exc)
+
+
+def install_exception_handlers(app: Starlette) -> None:
+    """注册异常处理器，让路由里抛出的鉴权异常也能正确变成 HTTP 状态码。"""
+    app.add_exception_handler(SaTokenException, sa_token_exception_handler)
+
+
+# FastAPI / Starlette 的 Depends 工厂（扩展写法）。
+# 各框架标准鉴权是注解：FastAPI 用 @sa.check_login，见 SaTokenFastAPI。
+
+
+async def _authorize(
+    request: Request,
+    *,
+    permissions: list[str] | None = None,
+    roles: list[str] | None = None,
+    mode: MatchMode = "OR",
+) -> tuple[str, str | None]:
+    ctx = StarletteHttpContext(request)
+    rule = build_rule(permissions=permissions, roles=roles, mode=mode)
+    result = await run_auth_flow(ctx, get_manager(), rule)
+    assert result.login_id is not None
+    request.state.sa_login_id = result.login_id
+    request.state.sa_token = result.token
+    return result.login_id, result.token
+
+
+def check_login() -> Callable[[Request], Awaitable[str]]:
+    """Depends 扩展：``Depends(check_login())``。标准写法是 ``@sa.check_login``。"""
+
+    async def dependency(request: Request) -> str:
+        login_id, _ = await _authorize(request)
+        return login_id
+
+    return dependency
+
+
+def check_permission(
+    *permissions: str,
+    mode: MatchMode = "OR",
+) -> Callable[[Request], Awaitable[str]]:
+    """Depends 扩展：``Depends(check_permission("order:delete"))``。"""
+
+    async def dependency(request: Request) -> str:
+        login_id, _ = await _authorize(request, permissions=list(permissions), mode=mode)
+        return login_id
+
+    return dependency
+
+
+def check_role(*roles: str, mode: MatchMode = "OR") -> Callable[[Request], Awaitable[str]]:
+    """Depends 扩展：``Depends(check_role("admin"))``。"""
+
+    async def dependency(request: Request) -> str:
+        login_id, _ = await _authorize(request, roles=list(roles), mode=mode)
+        return login_id
+
+    return dependency
+
+
+def check_disable(
+    service: str = "login",
+    level: int = 1,
+) -> Callable[[Request], Awaitable[str]]:
+    """Depends 扩展：``Depends(check_disable("comment"))``。"""
+
+    async def dependency(request: Request) -> str:
+        login_id, _ = await _authorize(request)
+        await get_manager().stp().check_disable(login_id, service=service, level=level)
+        return login_id
+
+    return dependency
+
+
+def check_safe(business: str) -> Callable[[Request], Awaitable[str]]:
+    """Depends 扩展：``Depends(check_safe("pay"))``。"""
+
+    async def dependency(request: Request) -> str:
+        login_id, token = await _authorize(request)
+        await get_manager().stp().check_safe(token, business)
+        return login_id
+
+    return dependency
+
+
+async def current_login_id(request: Request) -> str:
+    """Depends 扩展：``Depends(current_login_id)``。标准写法是 ``@sa.check_login``。"""
+    login_id, _ = await _authorize(request)
+    return login_id
+
+
+async def current_login_id_or_none(request: Request) -> str | None:
+    """Depends 扩展：未登录时返回 None。"""
+    ctx = StarletteHttpContext(request)
+    token = resolve_token(ctx, get_manager())
+    return await get_manager().stp().get_login_id_or_none(token)
+
+
+async def current_token(request: Request) -> str | None:
+    """Depends 扩展：只取 token，不做校验。"""
+    return resolve_token(StarletteHttpContext(request), get_manager())

sa_token/listener.py

@@ -0,0 +1,100 @@
+"""事件系统：用于审计日志、指标、消息推送。
+
+监听器异常不会影响主流程——审计失败不该导致用户登录失败。
+事件载荷只带 token 指纹，避免把原始 token 写进日志文件。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import logging
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+from .model import now_ms
+
+__all__ = ["Event", "EventData", "EventBus", "Listener"]
+
+logger = logging.getLogger("sa_token.listener")
+
+
+class Event(str, Enum):
+    LOGIN = "login"
+    LOGOUT = "logout"
+    KICKOUT = "kickout"
+    REPLACED = "replaced"
+    DISABLE = "disable"
+    UNTIE = "untie"
+    RENEW = "renew"
+    PERMISSION_CHECK = "permission_check"
+    ROLE_CHECK = "role_check"
+    ALL = "*"
+
+
+@dataclass
+class EventData:
+    event: Event
+    login_id: str | None = None
+    login_type: str = "login"
+    device: str | None = None
+    token_fingerprint: str | None = None
+    timestamp: int = field(default_factory=now_ms)
+    detail: dict[str, Any] = field(default_factory=dict)
+
+
+Listener = Callable[[EventData], Any | Awaitable[Any]]
+
+
+def fingerprint(token: str | None) -> str | None:
+    """token 的短指纹：可用于关联同一条会话，但无法还原出 token。"""
+    if not token:
+        return None
+    return hashlib.sha256(token.encode("utf-8")).hexdigest()[:16]
+
+
+@dataclass(order=True)
+class _Registration:
+    # 先按优先级倒序，再按注册序号，保证同优先级下注册顺序稳定。
+    sort_key: tuple[int, int]
+    listener: Listener = field(compare=False)
+    is_async: bool = field(compare=False, default=False)
+
+
+class EventBus:
+    def __init__(self) -> None:
+        self._listeners: dict[Event, list[_Registration]] = {}
+        self._counter = 0
+
+    def on(self, event: Event, listener: Listener, *, priority: int = 0) -> None:
+        """注册监听器，``priority`` 越大越先执行。"""
+        self._counter += 1
+        registration = _Registration(
+            sort_key=(-priority, self._counter),
+            listener=listener,
+            is_async=asyncio.iscoroutinefunction(listener),
+        )
+        bucket = self._listeners.setdefault(event, [])
+        bucket.append(registration)
+        bucket.sort()
+
+    def off(self, event: Event, listener: Listener) -> None:
+        bucket = self._listeners.get(event)
+        if not bucket:
+            return
+        self._listeners[event] = [item for item in bucket if item.listener is not listener]
+
+    def clear(self) -> None:
+        self._listeners.clear()
+
+    async def emit(self, data: EventData) -> None:
+        matched = [*self._listeners.get(data.event, []), *self._listeners.get(Event.ALL, [])]
+        for registration in matched:
+            try:
+                result = registration.listener(data)
+                if registration.is_async or asyncio.iscoroutine(result):
+                    await result
+            except Exception:
+                logger.exception("sa-token 事件监听器执行失败：event=%s", data.event.value)