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

sa_token/manager.py

@@ -0,0 +1,244 @@
+"""SaTokenManager 与 Builder。
+
+Manager 负责组装：配置 + 存储 + Token 策略 + 事件总线 + 权限数据源，
+并按 ``login_type`` 缓存 :class:`StpLogic` 实例。
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, Any
+
+from .config import SaTokenConfig
+from .listener import Event, EventBus, Listener
+from .model import DEFAULT_LOGIN_TYPE
+from .stp_logic import StpLogic
+from .strategy import TokenStrategy, create_strategy
+
+if TYPE_CHECKING:  # pragma: no cover - 仅供类型检查
+    from .security import NonceManager, RefreshTokenManager, TempTokenManager
+    from .storage.base import SaStorage
+    from .stp_interface import StpInterface
+
+__all__ = ["SaTokenManager", "SaTokenBuilder", "SaToken", "__version__"]
+
+__version__ = "0.1.1"
+
+_BANNER = r"""
+   _____         ______      __
+  / ___/____ _  /_  __/___  / /_____  ____
+  \__ \/ __ `/   / / / __ \/ //_/ _ \/ __ \
+ ___/ / /_/ /   / / / /_/ / ,< /  __/ / / /
+/____/\__,_/   /_/  \____/_/|_|\___/_/ /_/
+"""
+
+
+class SaTokenManager:
+    """框架无关的运行时核心。"""
+
+    def __init__(
+        self,
+        config: SaTokenConfig,
+        storage: SaStorage,
+        *,
+        strategy: TokenStrategy | None = None,
+        stp_interface: StpInterface | None = None,
+        events: EventBus | None = None,
+    ) -> None:
+        self.config = config
+        self.storage = storage
+        self.strategy = strategy or create_strategy(config)
+        self.stp_interface = stp_interface
+        self.events = events or EventBus()
+        self._logics: dict[str, StpLogic] = {}
+        self._nonce_manager: NonceManager | None = None
+        self._refresh_token_manager: RefreshTokenManager | None = None
+        self._temp_token_manager: TempTokenManager | None = None
+
+    def stp(self, login_type: str = DEFAULT_LOGIN_TYPE) -> StpLogic:
+        """取得指定账号体系的逻辑实例，不同 ``login_type`` 之间存储完全隔离。"""
+        logic = self._logics.get(login_type)
+        if logic is None:
+            logic = StpLogic(self, login_type)
+            self._logics[login_type] = logic
+        return logic
+
+    def on(self, event: Event, listener: Listener, *, priority: int = 0) -> None:
+        self.events.on(event, listener, priority=priority)
+
+    @property
+    def nonces(self) -> NonceManager:
+        """服务端 Nonce 管理器，按需创建且复用当前 Storage。"""
+        if self._nonce_manager is None:
+            from .security import NonceManager
+
+            self._nonce_manager = NonceManager(
+                self.storage,
+                key_prefix=self.config.storage_key_prefix,
+                timeout=self.config.nonce_timeout,
+            )
+        return self._nonce_manager
+
+    @property
+    def refresh_tokens(self) -> RefreshTokenManager:
+        """登录态 Refresh Token 管理器（不是 OAuth2 Refresh Token）。"""
+        if self._refresh_token_manager is None:
+            from .security import RefreshTokenManager
+
+            self._refresh_token_manager = RefreshTokenManager(self)
+        return self._refresh_token_manager
+
+    @property
+    def temp_tokens(self) -> TempTokenManager:
+        """邀请、重置密码等短时业务 Token 管理器。"""
+        if self._temp_token_manager is None:
+            from .security import TempTokenManager
+
+            self._temp_token_manager = TempTokenManager(
+                self.storage, key_prefix=self.config.storage_key_prefix
+            )
+        return self._temp_token_manager
+
+    def print_banner(self) -> None:
+        lines = [
+            _BANNER,
+            f":: sa-token-python ::   (v{__version__})",
+            f":: Python ::            {sys.version.split()[0]}",
+            f":: Token Style ::       {self.config.token_style}",
+            f":: Token Timeout ::     {self.config.timeout} 秒",
+            f":: Storage ::           {type(self.storage).__name__}",
+            "",
+        ]
+        print("\n".join(lines))
+
+
+class SaTokenBuilder:
+    """链式配置构建器。
+
+    每个方法只改配置，``build()` 之前不会产生任何副作用，便于测试。
+    """
+
+    def __init__(self) -> None:
+        self._config = SaTokenConfig()
+        self._storage: SaStorage | None = None
+        self._strategy: TokenStrategy | None = None
+        self._stp_interface: StpInterface | None = None
+        self._events = EventBus()
+        self._set_global = True
+
+    # 配置项 ---------------------------------------------------------------
+
+    def config(self, config: SaTokenConfig) -> SaTokenBuilder:
+        """整体替换配置，之后仍可用其它方法微调。"""
+        self._config = config
+        return self
+
+    def storage(self, storage: SaStorage) -> SaTokenBuilder:
+        self._storage = storage
+        return self
+
+    def strategy(self, strategy: TokenStrategy) -> SaTokenBuilder:
+        self._strategy = strategy
+        return self
+
+    def stp_interface(self, stp_interface: StpInterface) -> SaTokenBuilder:
+        self._stp_interface = stp_interface
+        return self
+
+    def token_name(self, name: str) -> SaTokenBuilder:
+        self._config.token_name = name
+        return self
+
+    def timeout(self, seconds: int) -> SaTokenBuilder:
+        self._config.timeout = seconds
+        return self
+
+    def active_timeout(self, seconds: int) -> SaTokenBuilder:
+        self._config.active_timeout = seconds
+        return self
+
+    def token_style(self, style: str) -> SaTokenBuilder:
+        self._config.token_style = style
+        return self
+
+    def token_prefix(self, prefix: str) -> SaTokenBuilder:
+        self._config.token_prefix = prefix
+        return self
+
+    def is_concurrent(self, value: bool) -> SaTokenBuilder:
+        self._config.is_concurrent = value
+        return self
+
+    def is_share(self, value: bool) -> SaTokenBuilder:
+        self._config.is_share = value
+        return self
+
+    def max_login_count(self, count: int) -> SaTokenBuilder:
+        self._config.max_login_count = count
+        return self
+
+    def auto_renew(self, value: bool) -> SaTokenBuilder:
+        self._config.auto_renew = value
+        return self
+
+    def storage_key_prefix(self, prefix: str) -> SaTokenBuilder:
+        self._config.storage_key_prefix = prefix
+        return self
+
+    def jwt_secret_key(self, secret: str) -> SaTokenBuilder:
+        self._config.jwt_secret_key = secret
+        return self
+
+    def print_banner(self, value: bool) -> SaTokenBuilder:
+        self._config.is_print_banner = value
+        return self
+
+    def set_option(self, **options: Any) -> SaTokenBuilder:
+        """批量设置任意配置项，便于从配置文件加载。"""
+        for key, value in options.items():
+            if not hasattr(self._config, key):
+                raise ValueError(f"未知配置项：{key}")
+            setattr(self._config, key, value)
+        return self
+
+    def on(self, event: Event, listener: Listener, *, priority: int = 0) -> SaTokenBuilder:
+        self._events.on(event, listener, priority=priority)
+        return self
+
+    def as_global(self, value: bool) -> SaTokenBuilder:
+        """是否把构建结果设为 ``StpUtil`` 使用的全局实例，默认是。"""
+        self._set_global = value
+        return self
+
+    # 构建 -----------------------------------------------------------------
+
+    def build(self) -> SaTokenManager:
+        storage = self._storage
+        if storage is None:
+            # 不静默降级到内存：多进程部署下会出现「登录了但下个请求说没登录」。
+            from .storage.memory import MemoryStorage
+
+            storage = MemoryStorage()
+        self._config.__post_init__()
+        manager = SaTokenManager(
+            self._config,
+            storage,
+            strategy=self._strategy,
+            stp_interface=self._stp_interface,
+            events=self._events,
+        )
+        if self._config.is_print_banner:
+            manager.print_banner()
+        if self._set_global:
+            from .stp_util import set_manager
+
+            set_manager(manager)
+        return manager
+
+
+class SaToken:
+    """入口门面：``SaToken.builder()...build()``。"""
+
+    @staticmethod
+    def builder() -> SaTokenBuilder:
+        return SaTokenBuilder()

sa_token/model.py

@@ -0,0 +1,145 @@
+"""核心数据模型。
+
+所有模型都能无损地 JSON 序列化：存储层只接受字符串，这样才能跨进程、
+跨语言共享同一份 Redis 数据，也避免 pickle 带来的反序列化风险。
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from .exception import NotLoginType
+
+__all__ = [
+    "DEFAULT_DEVICE",
+    "DEFAULT_LOGIN_TYPE",
+    "TokenInfo",
+    "TerminalInfo",
+    "SessionData",
+    "now_ms",
+]
+
+DEFAULT_DEVICE = "default"
+DEFAULT_LOGIN_TYPE = "login"
+
+#: token 记录里允许出现的下线状态，与 NotLoginType 中的子集一一对应。
+_OFFLINE_STATES = {
+    NotLoginType.KICK_OUT.value,
+    NotLoginType.BE_REPLACED.value,
+    NotLoginType.TOKEN_FREEZE.value,
+}
+
+
+def now_ms() -> int:
+    return int(time.time() * 1000)
+
+
+@dataclass
+class TokenInfo:
+    """``token -> 身份`` 的存储记录。
+
+    被踢 / 被顶时不删除该记录，而是把 ``state`` 置为下线原因并缩短 TTL，
+    这样下一次请求能返回明确原因，而不是笼统的「未登录」。
+    """
+
+    login_id: str
+    device: str = DEFAULT_DEVICE
+    login_type: str = DEFAULT_LOGIN_TYPE
+    create_time: int = field(default_factory=now_ms)
+    active_time: int = field(default_factory=now_ms)
+    timeout: int | None = None
+    active_timeout: int | None = None
+    tag: str | None = None
+    state: str | None = None
+    offline_time: int | None = None
+
+    @property
+    def is_offline(self) -> bool:
+        return self.state in _OFFLINE_STATES
+
+    @property
+    def offline_type(self) -> NotLoginType | None:
+        if not self.is_offline:
+            return None
+        return NotLoginType(self.state)
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), ensure_ascii=False, separators=(",", ":"))
+
+    @classmethod
+    def from_json(cls, raw: str) -> TokenInfo | None:
+        try:
+            payload = json.loads(raw)
+        except (TypeError, ValueError):
+            return None
+        if not isinstance(payload, dict) or "login_id" not in payload:
+            return None
+        allowed = {f for f in cls.__dataclass_fields__}
+        return cls(**{k: v for k, v in payload.items() if k in allowed})
+
+
+@dataclass
+class TerminalInfo:
+    """Account-Session 中的一条在线终端记录。"""
+
+    token: str
+    device: str = DEFAULT_DEVICE
+    login_time: int = field(default_factory=now_ms)
+    index: int = 0
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> TerminalInfo:
+        allowed = {f for f in cls.__dataclass_fields__}
+        return cls(**{k: v for k, v in payload.items() if k in allowed})
+
+
+@dataclass
+class SessionData:
+    """Session 的可序列化载荷。
+
+    Account-Session 额外承担「该账号当前有哪些终端在线」的索引职责，
+    避免为多端登录再维护一份容易和 Session 失配的索引键。
+    """
+
+    id: str
+    create_time: int = field(default_factory=now_ms)
+    data: dict[str, Any] = field(default_factory=dict)
+    terminal_list: list[TerminalInfo] = field(default_factory=list)
+    history_terminal_count: int = 0
+
+    def to_json(self) -> str:
+        payload = {
+            "id": self.id,
+            "create_time": self.create_time,
+            "data": self.data,
+            "terminal_list": [terminal.to_dict() for terminal in self.terminal_list],
+            "history_terminal_count": self.history_terminal_count,
+        }
+        return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+
+    @classmethod
+    def from_json(cls, raw: str) -> SessionData | None:
+        try:
+            payload = json.loads(raw)
+        except (TypeError, ValueError):
+            return None
+        if not isinstance(payload, dict) or "id" not in payload:
+            return None
+        terminals = [
+            TerminalInfo.from_dict(item)
+            for item in payload.get("terminal_list", [])
+            if isinstance(item, dict)
+        ]
+        return cls(
+            id=payload["id"],
+            create_time=payload.get("create_time", now_ms()),
+            data=payload.get("data", {}),
+            terminal_list=terminals,
+            history_terminal_count=payload.get("history_terminal_count", len(terminals)),
+        )

sa_token/oauth2/__init__.py

@@ -0,0 +1,19 @@
+"""OAuth2 授权码流程（含 PKCE）。
+
+框架无关：协议引擎只依赖核心存储，HTTP 端点由使用方挂载。
+"""
+
+from __future__ import annotations
+
+from .model import AccessTokenInfo, AuthorizationCode, OAuth2Client, TokenResponse
+from .server import OAuth2Error, OAuth2Server, generate_pkce_pair
+
+__all__ = [
+    "OAuth2Client",
+    "AuthorizationCode",
+    "AccessTokenInfo",
+    "TokenResponse",
+    "OAuth2Server",
+    "OAuth2Error",
+    "generate_pkce_pair",
+]

sa_token/oauth2/model.py

@@ -0,0 +1,122 @@
+"""OAuth2 数据模型。"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from ..model import now_ms
+
+__all__ = ["OAuth2Client", "AuthorizationCode", "AccessTokenInfo", "TokenResponse"]
+
+
+@dataclass
+class OAuth2Client:
+    """已注册的客户端。
+
+    ``client_secret`` 为空表示公开客户端（SPA / 移动端），此类客户端
+    **必须**使用 PKCE，否则授权码可能被同设备上的恶意应用截获。
+    """
+
+    client_id: str
+    client_secret: str | None = None
+    redirect_uris: list[str] = field(default_factory=list)
+    grant_types: list[str] = field(default_factory=lambda: ["authorization_code"])
+    scopes: list[str] = field(default_factory=list)
+    name: str | None = None
+
+    @property
+    def is_public(self) -> bool:
+        return not self.client_secret
+
+    def allows_redirect(self, redirect_uri: str) -> bool:
+        # 精确匹配，不做前缀匹配：前缀匹配是经典的开放重定向漏洞来源。
+        return redirect_uri in self.redirect_uris
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), ensure_ascii=False, separators=(",", ":"))
+
+    @classmethod
+    def from_json(cls, raw: str) -> OAuth2Client | None:
+        try:
+            payload = json.loads(raw)
+        except ValueError:
+            return None
+        if not isinstance(payload, dict) or "client_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})
+
+
+@dataclass
+class AuthorizationCode:
+    code: str
+    client_id: str
+    login_id: str
+    redirect_uri: str
+    scopes: list[str] = field(default_factory=list)
+    create_time: int = field(default_factory=now_ms)
+    code_challenge: str | None = None
+    code_challenge_method: str = "S256"
+    state: str | None = None
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), ensure_ascii=False, separators=(",", ":"))
+
+    @classmethod
+    def from_json(cls, raw: str) -> AuthorizationCode | None:
+        try:
+            payload = json.loads(raw)
+        except ValueError:
+            return None
+        if not isinstance(payload, dict) or "code" not in payload:
+            return None
+        allowed = set(cls.__dataclass_fields__)
+        return cls(**{k: v for k, v in payload.items() if k in allowed})
+
+
+@dataclass
+class AccessTokenInfo:
+    access_token: str
+    client_id: str
+    login_id: str
+    scopes: list[str] = field(default_factory=list)
+    create_time: int = field(default_factory=now_ms)
+    expires_in: int = 7200
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), ensure_ascii=False, separators=(",", ":"))
+
+    @classmethod
+    def from_json(cls, raw: str) -> AccessTokenInfo | None:
+        try:
+            payload = json.loads(raw)
+        except ValueError:
+            return None
+        if not isinstance(payload, dict) or "access_token" not in payload:
+            return None
+        allowed = set(cls.__dataclass_fields__)
+        return cls(**{k: v for k, v in payload.items() if k in allowed})
+
+
+@dataclass
+class TokenResponse:
+    """``/oauth2/token`` 端点的标准响应体。"""
+
+    access_token: str
+    refresh_token: str | None = None
+    token_type: str = "Bearer"
+    expires_in: int = 7200
+    scope: str = ""
+
+    def to_dict(self) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "access_token": self.access_token,
+            "token_type": self.token_type,
+            "expires_in": self.expires_in,
+            "scope": self.scope,
+        }
+        if self.refresh_token:
+            payload["refresh_token"] = self.refresh_token
+        return payload