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

sa_token/oauth2/server.py

@@ -0,0 +1,361 @@
+"""OAuth2 授权码流程。
+
+只依赖核心的存储抽象，不依赖任何 Web 框架，因此协议逻辑可以在单测里
+纯函数式地跑通；HTTP 端点由使用方按自己的框架挂载。
+
+签发出来的 access_token 是**有状态**的：存储里删掉即刻失效，
+与本项目「服务端必须能即时作废凭证」的整体语义一致。
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import secrets
+from typing import TYPE_CHECKING
+
+from ..exception import SaTokenException
+from .model import AccessTokenInfo, AuthorizationCode, OAuth2Client, TokenResponse
+
+if TYPE_CHECKING:  # pragma: no cover - 仅供类型检查
+    from ..manager import SaTokenManager
+
+__all__ = ["OAuth2Error", "OAuth2Server", "generate_pkce_pair"]
+
+
+class OAuth2Error(SaTokenException):
+    """OAuth2 协议错误，``error`` 为规范定义的错误码。"""
+
+    http_status = 400
+
+    def __init__(self, error: str, description: str) -> None:
+        super().__init__(f"{error}: {description}")
+        self.error = error
+        self.description = description
+
+
+def _s256(verifier: str) -> str:
+    digest = hashlib.sha256(verifier.encode("ascii")).digest()
+    return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
+
+
+class OAuth2Server:
+    """授权服务端。
+
+    支持 ``authorization_code``（含 PKCE）与 ``refresh_token`` 两种授权类型。
+    """
+
+    def __init__(
+        self,
+        manager: SaTokenManager,
+        *,
+        code_timeout: int = 300,
+        access_token_timeout: int = 7200,
+        refresh_token_timeout: int = 2592000,
+    ) -> None:
+        self._manager = manager
+        self.code_timeout = code_timeout
+        self.access_token_timeout = access_token_timeout
+        self.refresh_token_timeout = refresh_token_timeout
+
+    # 存储键 ---------------------------------------------------------------
+
+    def _key(self, suffix: str, identifier: str) -> str:
+        return self._manager.config.make_key("oauth2", suffix, identifier)
+
+    @property
+    def _storage(self):
+        return self._manager.storage
+
+    # 客户端 ---------------------------------------------------------------
+
+    async def register_client(self, client: OAuth2Client) -> None:
+        await self._storage.set(self._key("client", client.client_id), client.to_json())
+
+    async def get_client(self, client_id: str) -> OAuth2Client | None:
+        raw = await self._storage.get(self._key("client", client_id))
+        return OAuth2Client.from_json(raw) if raw else None
+
+    async def remove_client(self, client_id: str) -> None:
+        await self._storage.delete(self._key("client", client_id))
+
+    async def _require_client(self, client_id: str) -> OAuth2Client:
+        client = await self.get_client(client_id)
+        if client is None:
+            raise OAuth2Error("invalid_client", f"未注册的 client_id: {client_id}")
+        return client
+
+    # 授权码 ---------------------------------------------------------------
+
+    async def create_authorization_code(
+        self,
+        *,
+        client_id: str,
+        login_id: str,
+        redirect_uri: str,
+        scopes: list[str] | None = None,
+        code_challenge: str | None = None,
+        code_challenge_method: str = "S256",
+        state: str | None = None,
+    ) -> AuthorizationCode:
+        """用户在授权页点击「同意」之后调用。"""
+        client = await self._require_client(client_id)
+        if not client.allows_redirect(redirect_uri):
+            raise OAuth2Error("invalid_request", "redirect_uri 未在客户端注册列表中")
+        if "authorization_code" not in client.grant_types:
+            raise OAuth2Error("unauthorized_client", "客户端不支持 authorization_code")
+        if client.is_public and not code_challenge:
+            raise OAuth2Error("invalid_request", "公开客户端必须使用 PKCE")
+
+        requested = list(scopes or client.scopes)
+        invalid = [scope for scope in requested if scope not in client.scopes]
+        if invalid:
+            raise OAuth2Error("invalid_scope", f"未授权的 scope: {', '.join(invalid)}")
+
+        code = AuthorizationCode(
+            code=secrets.token_urlsafe(32),
+            client_id=client_id,
+            login_id=login_id,
+            redirect_uri=redirect_uri,
+            scopes=requested,
+            code_challenge=code_challenge,
+            code_challenge_method=code_challenge_method,
+            state=state,
+        )
+        await self._storage.set(self._key("code", code.code), code.to_json(), self.code_timeout)
+        return code
+
+    async def _consume_code(self, code: str) -> AuthorizationCode:
+        """原子消费授权码：读取即删除，防止重放。"""
+        key = self._key("code", code)
+        raw = await self._storage.get(key)
+        if raw is None:
+            raise OAuth2Error("invalid_grant", "授权码无效或已过期")
+        if not await self._storage.compare_and_delete(key, raw):
+            # 删除失败说明已被并发请求取走，同样按重放处理。
+            raise OAuth2Error("invalid_grant", "授权码已被使用")
+        parsed = AuthorizationCode.from_json(raw)
+        if parsed is None:
+            raise OAuth2Error("invalid_grant", "授权码数据损坏")
+        return parsed
+
+    # 令牌 -----------------------------------------------------------------
+
+    async def exchange_code_for_token(
+        self,
+        *,
+        code: str,
+        client_id: str,
+        client_secret: str | None = None,
+        redirect_uri: str | None = None,
+        code_verifier: str | None = None,
+    ) -> TokenResponse:
+        client = await self._require_client(client_id)
+        self._verify_client_secret(client, client_secret)
+
+        authorization_code = await self._consume_code(code)
+        if authorization_code.client_id != client_id:
+            raise OAuth2Error("invalid_grant", "授权码不属于该客户端")
+        if redirect_uri is not None and authorization_code.redirect_uri != redirect_uri:
+            raise OAuth2Error("invalid_grant", "redirect_uri 与申请时不一致")
+        self._verify_pkce(authorization_code, code_verifier)
+
+        return await self._issue_tokens(
+            client_id=client_id,
+            login_id=authorization_code.login_id,
+            scopes=authorization_code.scopes,
+        )
+
+    async def refresh_access_token(
+        self,
+        *,
+        refresh_token: str,
+        client_id: str,
+        client_secret: str | None = None,
+    ) -> TokenResponse:
+        client = await self._require_client(client_id)
+        self._verify_client_secret(client, client_secret)
+
+        key = self._key("refresh", refresh_token)
+        raw = await self._storage.get(key)
+        if raw is None:
+            raise OAuth2Error("invalid_grant", "refresh_token 无效或已过期")
+        info = AccessTokenInfo.from_json(raw)
+        if info is None or info.client_id != client_id:
+            raise OAuth2Error("invalid_grant", "refresh_token 不属于该客户端")
+
+        # 原子消费并轮转：并发刷新时只能有一个请求成功。
+        if not await self._storage.compare_and_delete(key, raw):
+            raise OAuth2Error("invalid_grant", "refresh_token 已被使用")
+        await self._storage.delete(self._key("access", info.access_token))
+        return await self._issue_tokens(
+            client_id=client_id, login_id=info.login_id, scopes=info.scopes
+        )
+
+    async def client_credentials_token(
+        self,
+        *,
+        client_id: str,
+        client_secret: str,
+        scopes: list[str] | None = None,
+    ) -> TokenResponse:
+        """机器到机器的 ``client_credentials`` 授权。
+
+        公开客户端不能使用该模式；返回的 access token 不签发 refresh token，
+        因为客户端可随时用自己的凭证重新申请。
+        """
+        client = await self._require_client(client_id)
+        if client.is_public:
+            raise OAuth2Error("unauthorized_client", "公开客户端不能使用 client_credentials")
+        self._verify_client_secret(client, client_secret)
+        if "client_credentials" not in client.grant_types:
+            raise OAuth2Error("unauthorized_client", "客户端未启用 client_credentials")
+        requested = list(scopes or client.scopes)
+        invalid = [scope for scope in requested if scope not in client.scopes]
+        if invalid:
+            raise OAuth2Error("invalid_scope", f"未授权的 scope: {', '.join(invalid)}")
+
+        access_token = secrets.token_urlsafe(32)
+        info = AccessTokenInfo(
+            access_token=access_token,
+            client_id=client_id,
+            login_id=f"client:{client_id}",
+            scopes=requested,
+            expires_in=self.access_token_timeout,
+        )
+        await self._storage.set(
+            self._key("access", access_token),
+            info.to_json(),
+            self.access_token_timeout,
+        )
+        return TokenResponse(
+            access_token=access_token,
+            expires_in=self.access_token_timeout,
+            scope=" ".join(requested),
+        )
+
+    async def introspect(self, access_token: str) -> dict[str, object]:
+        """返回适合 RFC 7662 introspection 端点的结果。"""
+        raw = await self._storage.get(self._key("access", access_token))
+        info = AccessTokenInfo.from_json(raw) if raw else None
+        if info is None:
+            return {"active": False}
+        return {
+            "active": True,
+            "client_id": info.client_id,
+            "sub": info.login_id,
+            "scope": " ".join(info.scopes),
+            "token_type": "Bearer",
+        }
+
+    async def _issue_tokens(
+        self,
+        *,
+        client_id: str,
+        login_id: str,
+        scopes: list[str],
+    ) -> TokenResponse:
+        access_token = secrets.token_urlsafe(32)
+        refresh_token = secrets.token_urlsafe(32)
+        info = AccessTokenInfo(
+            access_token=access_token,
+            client_id=client_id,
+            login_id=login_id,
+            scopes=scopes,
+            expires_in=self.access_token_timeout,
+        )
+        await self._storage.set(
+            self._key("access", access_token), info.to_json(), self.access_token_timeout
+        )
+        await self._storage.set(
+            self._key("refresh", refresh_token), info.to_json(), self.refresh_token_timeout
+        )
+        return TokenResponse(
+            access_token=access_token,
+            refresh_token=refresh_token,
+            expires_in=self.access_token_timeout,
+            scope=" ".join(scopes),
+        )
+
+    async def verify_access_token(self, access_token: str) -> AccessTokenInfo:
+        raw = await self._storage.get(self._key("access", access_token))
+        if raw is None:
+            raise OAuth2Error("invalid_token", "access_token 无效或已过期")
+        info = AccessTokenInfo.from_json(raw)
+        if info is None:
+            raise OAuth2Error("invalid_token", "access_token 数据损坏")
+        return info
+
+    async def check_scope(self, access_token: str, scope: str) -> AccessTokenInfo:
+        info = await self.verify_access_token(access_token)
+        if scope not in info.scopes:
+            raise OAuth2Error("insufficient_scope", f"缺少 scope: {scope}")
+        return info
+
+    async def revoke_token(self, token: str) -> bool:
+        """吊销 access_token 或 refresh_token，两者都尝试。"""
+        revoked = False
+        for suffix in ("access", "refresh"):
+            key = self._key(suffix, token)
+            if await self._storage.exists(key):
+                await self._storage.delete(key)
+                revoked = True
+        return revoked
+
+    def build_authorize_url(
+        self,
+        *,
+        authorize_endpoint: str,
+        client_id: str,
+        redirect_uri: str,
+        scopes: list[str] | None = None,
+        state: str | None = None,
+        code_challenge: str | None = None,
+        code_challenge_method: str = "S256",
+    ) -> str:
+        """给客户端生成跳转到授权页的 URL。"""
+        from urllib.parse import urlencode
+
+        params = {
+            "response_type": "code",
+            "client_id": client_id,
+            "redirect_uri": redirect_uri,
+        }
+        if scopes:
+            params["scope"] = " ".join(scopes)
+        if state:
+            params["state"] = state
+        if code_challenge:
+            params["code_challenge"] = code_challenge
+            params["code_challenge_method"] = code_challenge_method
+        separator = "&" if "?" in authorize_endpoint else "?"
+        return f"{authorize_endpoint}{separator}{urlencode(params)}"
+
+    # 校验 -----------------------------------------------------------------
+
+    @staticmethod
+    def _verify_client_secret(client: OAuth2Client, client_secret: str | None) -> None:
+        if client.is_public:
+            return
+        if not client_secret or not secrets.compare_digest(
+            client.client_secret or "", client_secret
+        ):
+            raise OAuth2Error("invalid_client", "client_secret 不正确")
+
+    @staticmethod
+    def _verify_pkce(code: AuthorizationCode, code_verifier: str | None) -> None:
+        if not code.code_challenge:
+            return
+        if not code_verifier:
+            raise OAuth2Error("invalid_grant", "缺少 code_verifier")
+        expected = (
+            code_verifier if code.code_challenge_method == "plain" else _s256(code_verifier)
+        )
+        if not secrets.compare_digest(expected, code.code_challenge):
+            raise OAuth2Error("invalid_grant", "code_verifier 校验失败")
+
+
+def generate_pkce_pair() -> tuple[str, str]:
+    """生成 ``(code_verifier, code_challenge)``，供客户端使用。"""
+    verifier = secrets.token_urlsafe(48)
+    return verifier, _s256(verifier)

sa_token/online/__init__.py

@@ -0,0 +1,292 @@
+"""在线用户管理与 WebSocket 鉴权。
+
+连接鉴权复用 ``token_io`` + ``StpLogic.check_login``，不另起一套校验；
+在线状态分两层：
+
+- **存储层**：跨进程可见的在线记录，用于「谁在线」这类查询
+- **进程内**：真实的连接对象，用于推送与断开
+
+之所以拆两层：连接对象本身没法序列化进 Redis，跨进程踢人只能靠存储层
+标记 + 各进程自己断开本地连接。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Awaitable, Callable
+from dataclasses import asdict, dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from ..adapter.http import HttpContext
+from ..exception import NotLoginException, NotLoginType
+from ..listener import Event, EventData
+from ..model import now_ms
+from ..token_io import read_token
+
+if TYPE_CHECKING:  # pragma: no cover - 仅供类型检查
+    from ..manager import SaTokenManager
+
+__all__ = ["OnlineUser", "OnlineManager", "WebSocketAuthenticator"]
+
+#: 推送回调：接收 (连接对象, 消息文本)。
+Sender = Callable[[Any, str], Awaitable[None]]
+
+
+@dataclass
+class OnlineUser:
+    login_id: str
+    device: str = "default"
+    connection_id: str = ""
+    connect_time: int = field(default_factory=now_ms)
+    last_heartbeat: int = field(default_factory=now_ms)
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), ensure_ascii=False, separators=(",", ":"))
+
+    @classmethod
+    def from_json(cls, raw: str) -> OnlineUser | None:
+        try:
+            payload = json.loads(raw)
+        except ValueError:
+            return None
+        if not isinstance(payload, dict) or "login_id" not in payload:
+            return None
+        allowed = set(cls.__dataclass_fields__)
+        return cls(**{k: v for k, v in payload.items() if k in allowed})
+
+
+class OnlineManager:
+    """在线用户注册表。"""
+
+    def __init__(
+        self,
+        manager: SaTokenManager,
+        *,
+        heartbeat_timeout: int = 300,
+        sender: Sender | None = None,
+        closer: Callable[[Any], Awaitable[None]] | None = None,
+    ) -> None:
+        self._manager = manager
+        self.heartbeat_timeout = heartbeat_timeout
+        self._sender = sender
+        self._closer = closer
+        # login_id -> {connection_id: 连接对象}
+        self._connections: dict[str, dict[str, Any]] = {}
+        self._lock = asyncio.Lock()
+        manager.on(Event.KICKOUT, self._on_offline_event)
+        manager.on(Event.REPLACED, self._on_offline_event)
+
+    def _key(self, login_id: str, connection_id: str) -> str:
+        return self._manager.config.make_key("online", "user", f"{login_id}:{connection_id}")
+
+    @property
+    def _storage(self):
+        return self._manager.storage
+
+    async def register(
+        self,
+        login_id: str,
+        connection: Any,
+        *,
+        connection_id: str | None = None,
+        device: str = "default",
+    ) -> OnlineUser:
+        """登记一条在线连接。"""
+        resolved_id = connection_id or f"{id(connection):x}"
+        user = OnlineUser(
+            login_id=login_id,
+            device=device,
+            connection_id=resolved_id,
+        )
+        await self._storage.set(
+            self._key(login_id, resolved_id), user.to_json(), self.heartbeat_timeout
+        )
+        async with self._lock:
+            self._connections.setdefault(login_id, {})[resolved_id] = connection
+        return user
+
+    async def unregister(self, login_id: str, connection_id: str) -> None:
+        await self._storage.delete(self._key(login_id, connection_id))
+        async with self._lock:
+            bucket = self._connections.get(login_id)
+            if bucket is not None:
+                bucket.pop(connection_id, None)
+                if not bucket:
+                    self._connections.pop(login_id, None)
+
+    async def heartbeat(self, login_id: str, connection_id: str) -> bool:
+        """刷新心跳，连接记录已过期时返回 False。"""
+        key = self._key(login_id, connection_id)
+        raw = await self._storage.get(key)
+        if raw is None:
+            return False
+        user = OnlineUser.from_json(raw)
+        if user is None:
+            return False
+        user.last_heartbeat = now_ms()
+        await self._storage.set(key, user.to_json(), self.heartbeat_timeout)
+        return True
+
+    async def is_online(self, login_id: str) -> bool:
+        cursor: str | None = None
+        pattern = self._key(login_id, "*")
+        while True:
+            cursor, keys = await self._storage.scan(pattern, cursor, 50)
+            if keys:
+                return True
+            if cursor is None:
+                return False
+
+    async def get_online_users(self, login_id: str) -> list[OnlineUser]:
+        users: list[OnlineUser] = []
+        cursor: str | None = None
+        pattern = self._key(login_id, "*")
+        while True:
+            cursor, keys = await self._storage.scan(pattern, cursor, 100)
+            for key in keys:
+                raw = await self._storage.get(key)
+                user = OnlineUser.from_json(raw) if raw else None
+                if user is not None:
+                    users.append(user)
+            if cursor is None:
+                return users
+
+    def local_connections(self, login_id: str) -> list[Any]:
+        """本进程持有的连接对象。"""
+        return list(self._connections.get(login_id, {}).values())
+
+    async def send_to_user(self, login_id: str, message: str) -> int:
+        """向该用户在本进程上的所有连接推送，返回成功条数。"""
+        if self._sender is None:
+            raise RuntimeError("未配置 sender，无法推送消息")
+        sent = 0
+        for connection in self.local_connections(login_id):
+            try:
+                await self._sender(connection, message)
+                sent += 1
+            except Exception:
+                # 推送失败通常意味着连接已断，交给连接自身的清理流程处理。
+                continue
+        return sent
+
+    async def send_to_device(self, login_id: str, device: str, message: str) -> int:
+        """只向指定设备类型推送。"""
+        if self._sender is None:
+            raise RuntimeError("未配置 sender，无法推送消息")
+        online_users = await self.get_online_users(login_id)
+        allowed_ids = {
+            user.connection_id for user in online_users if user.device == device
+        }
+        sent = 0
+        local = self._connections.get(login_id, {})
+        for connection_id, connection in local.items():
+            if connection_id not in allowed_ids:
+                continue
+            try:
+                await self._sender(connection, message)
+                sent += 1
+            except Exception:
+                continue
+        return sent
+
+    async def kickout(self, login_id: str) -> None:
+        """踢人并断开该用户的连接。
+
+        先走核心的 ``kickout`` 让 token 失效，再断开连接：顺序反了的话，
+        客户端可能在断线重连的瞬间用旧 token 重新连上。
+        """
+        await self._manager.stp().kickout(login_id)
+
+    async def disconnect_user(self, login_id: str) -> None:
+        connections = self.local_connections(login_id)
+        async with self._lock:
+            connection_ids = list(self._connections.get(login_id, {}).keys())
+            self._connections.pop(login_id, None)
+        for connection_id in connection_ids:
+            await self._storage.delete(self._key(login_id, connection_id))
+        if self._closer is None:
+            return
+        for connection in connections:
+            try:
+                await self._closer(connection)
+            except Exception:
+                continue
+
+    async def disconnect_device(self, login_id: str, device: str) -> None:
+        """断开本进程中该用户指定设备的连接。"""
+        online_users = await self.get_online_users(login_id)
+        connection_ids = {
+            user.connection_id for user in online_users if user.device == device
+        }
+        async with self._lock:
+            bucket = self._connections.get(login_id, {})
+            connections = [
+                bucket.pop(connection_id)
+                for connection_id in list(connection_ids)
+                if connection_id in bucket
+            ]
+            if not bucket:
+                self._connections.pop(login_id, None)
+        for connection_id in connection_ids:
+            await self._storage.delete(self._key(login_id, connection_id))
+        if self._closer is not None:
+            for connection in connections:
+                try:
+                    await self._closer(connection)
+                except Exception:
+                    continue
+
+    async def cleanup_stale_connections(self) -> int:
+        """清理 Storage TTL 已过期、但本进程仍残留的连接对象。"""
+        stale: list[tuple[str, str, Any]] = []
+        async with self._lock:
+            for login_id, bucket in list(self._connections.items()):
+                for connection_id, connection in list(bucket.items()):
+                    if not await self._storage.exists(self._key(login_id, connection_id)):
+                        stale.append((login_id, connection_id, connection))
+                        bucket.pop(connection_id, None)
+                if not bucket:
+                    self._connections.pop(login_id, None)
+        if self._closer is not None:
+            for _, _, connection in stale:
+                try:
+                    await self._closer(connection)
+                except Exception:
+                    continue
+        return len(stale)
+
+    async def _on_offline_event(self, data: EventData) -> None:
+        """StpUtil.kickout/replaced 也能自动断开本进程连接。"""
+        if data.login_id is None:
+            return
+        if data.device:
+            await self.disconnect_device(data.login_id, data.device)
+        else:
+            await self.disconnect_user(data.login_id)
+
+
+class WebSocketAuthenticator:
+    """WebSocket 连接鉴权。
+
+    token 来源与 HTTP 完全一致（Header → Cookie → Query），因为浏览器
+    WebSocket API 不能自定义 Header，Query 往往是唯一可用的通道。
+    """
+
+    def __init__(self, manager: SaTokenManager, *, login_type: str = "login") -> None:
+        self._manager = manager
+        self.login_type = login_type
+
+    async def authenticate(self, ctx: HttpContext) -> str:
+        """握手阶段鉴权，返回 ``login_id``，失败抛 :class:`NotLoginException`。"""
+        token = read_token(ctx, self._manager.config)
+        login_id = await self._manager.stp(self.login_type).check_login(token)
+        ctx.state["stp_token"] = token
+        ctx.state["stp_login_id"] = login_id
+        return login_id
+
+    async def authenticate_token(self, token: str | None) -> str:
+        """已经自行取到 token（例如从首条消息里读）时使用。"""
+        if not token:
+            raise NotLoginException(NotLoginType.NOT_TOKEN, login_type=self.login_type)
+        return await self._manager.stp(self.login_type).check_login(token)