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

sa_token/sso/__init__.py

@@ -0,0 +1,217 @@
+"""SSO 单点登录：ticket 换登录态 + 统一登出。
+
+Server 与 Client 都只依赖核心存储与 ``StpLogic``，HTTP 跳转由使用方实现，
+因此整条流程可以在单测里不起服务地跑通。
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import secrets
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+from urllib.parse import urlencode
+
+from ..exception import SaTokenException
+
+if TYPE_CHECKING:  # pragma: no cover - 仅供类型检查
+    from ..manager import SaTokenManager
+
+__all__ = ["SsoConfig", "SsoError", "SsoTicket", "SsoServer", "SsoClient"]
+
+
+class SsoError(SaTokenException):
+    http_status = 400
+
+
+@dataclass
+class SsoConfig:
+    server_url: str = ""
+    """认证中心的登录页地址。"""
+
+    ticket_timeout: int = 300
+    """ticket 有效期，应当很短：它只用于一次跳转。"""
+
+    allowed_services: list[str] = field(default_factory=list)
+    """允许接入的 service 白名单，空列表表示不限制（仅建议内网使用）。"""
+
+    secret_key: str | None = None
+    """跨进程 ticket 校验的 HMAC 密钥；生产环境建议配置至少 32 字节。"""
+
+    def is_allowed_service(self, service: str) -> bool:
+        return not self.allowed_services or service in self.allowed_services
+
+
+@dataclass(frozen=True)
+class SsoTicket:
+    ticket: str
+    service: str
+    signature: str | None = None
+
+
+class SsoServer:
+    """认证中心。"""
+
+    def __init__(self, manager: SaTokenManager, config: SsoConfig | None = None) -> None:
+        self._manager = manager
+        self.config = config or SsoConfig()
+
+    def _key(self, suffix: str, identifier: str) -> str:
+        return self._manager.config.make_key("sso", suffix, identifier)
+
+    @property
+    def _storage(self):
+        return self._manager.storage
+
+    async def create_ticket(self, login_id: str, service: str) -> str:
+        """用户在认证中心登录后，为目标应用签发一次性 ticket。"""
+        if not self.config.is_allowed_service(service):
+            raise SsoError(f"service 未在白名单中：{service}")
+        ticket = secrets.token_urlsafe(24)
+        # ticket 绑定 service，避免 A 应用拿到的 ticket 被拿去登录 B 应用。
+        await self._storage.set(
+            self._key("ticket", ticket),
+            f"{login_id}\n{service}",
+            self.config.ticket_timeout,
+        )
+        await self._register_service(login_id, service)
+        return ticket
+
+    async def create_signed_ticket(self, login_id: str, service: str) -> SsoTicket:
+        """签发 ticket，并在配置密钥时附带 HMAC 签名。"""
+        ticket = await self.create_ticket(login_id, service)
+        return SsoTicket(
+            ticket=ticket,
+            service=service,
+            signature=self.sign_ticket(ticket, service),
+        )
+
+    def sign_ticket(self, ticket: str, service: str) -> str | None:
+        if not self.config.secret_key:
+            return None
+        payload = f"{ticket}\n{service}".encode()
+        return hmac.new(
+            self.config.secret_key.encode(),
+            payload,
+            hashlib.sha256,
+        ).hexdigest()
+
+    def verify_ticket_signature(
+        self, ticket: str, service: str, signature: str | None
+    ) -> None:
+        expected = self.sign_ticket(ticket, service)
+        if expected is None:
+            return
+        if signature is None or not hmac.compare_digest(expected, signature):
+            raise SsoError("ticket 签名无效")
+
+    async def validate_ticket(
+        self,
+        ticket: str,
+        service: str,
+        *,
+        signature: str | None = None,
+    ) -> str:
+        """校验并消费 ticket，返回 ``login_id``。
+
+        读取即删除：ticket 必须一次性，否则中间人可以重放它换取会话。
+        """
+        self.verify_ticket_signature(ticket, service, signature)
+        key = self._key("ticket", ticket)
+        raw = await self._storage.get(key)
+        if raw is None:
+            raise SsoError("ticket 无效或已过期")
+        if not await self._storage.compare_and_delete(key, raw):
+            raise SsoError("ticket 已被使用")
+        login_id, bound_service = raw.split("\n", 1)
+        if bound_service != service:
+            raise SsoError("ticket 与 service 不匹配")
+        return login_id
+
+    async def _register_service(self, login_id: str, service: str) -> None:
+        """记录该用户在哪些应用登录过，统一登出时需要逐个通知。"""
+        key = self._key("services", login_id)
+        for _ in range(12):
+            raw = await self._storage.get(key)
+            services = raw.split("\n") if raw else []
+            if service in services:
+                return
+            services.append(service)
+            new_raw = "\n".join(services)
+            if raw is None:
+                if await self._storage.set_if_absent(key, new_raw):
+                    return
+            elif await self._storage.compare_and_set(key, raw, new_raw):
+                return
+        raise SsoError("并发登记 SSO service 失败，请稍后重试")
+
+    async def get_registered_services(self, login_id: str) -> list[str]:
+        raw = await self._storage.get(self._key("services", login_id))
+        return raw.split("\n") if raw else []
+
+    async def logout(self, login_id: str) -> list[str]:
+        """统一登出：清掉认证中心会话，返回需要通知的客户端列表。
+
+        实际的 HTTP 回调由使用方发起——通知方式（同步 HTTP、消息队列）
+        取决于部署形态，不应该由库来替用户决定。
+        """
+        services = await self.get_registered_services(login_id)
+        await self._manager.stp().logout(login_id)
+        await self._storage.delete(self._key("services", login_id))
+        return services
+
+    def build_login_url(self, service: str, *, redirect: str | None = None) -> str:
+        params = {"service": service}
+        if redirect:
+            params["redirect"] = redirect
+        separator = "&" if "?" in self.config.server_url else "?"
+        return f"{self.config.server_url}{separator}{urlencode(params)}"
+
+
+class SsoClient:
+    """接入方应用。"""
+
+    def __init__(
+        self,
+        manager: SaTokenManager,
+        *,
+        server_url: str,
+        service: str,
+        login_type: str = "login",
+    ) -> None:
+        self._manager = manager
+        self.server_url = server_url
+        self.service = service
+        self.login_type = login_type
+
+    def get_login_url(self, redirect: str | None = None) -> str:
+        params = {"service": self.service}
+        if redirect:
+            params["redirect"] = redirect
+        separator = "&" if "?" in self.server_url else "?"
+        return f"{self.server_url}{separator}{urlencode(params)}"
+
+    async def login_by_ticket(
+        self,
+        server: SsoServer,
+        ticket: str,
+        *,
+        device: str | None = None,
+        signature: str | None = None,
+    ) -> str:
+        """同进程部署时直接校验 ticket 并建立本地登录态。"""
+        login_id = await server.validate_ticket(
+            ticket,
+            self.service,
+            signature=signature,
+        )
+        return await self._manager.stp(self.login_type).login(login_id, device=device)
+
+    async def login_by_login_id(self, login_id: str, *, device: str | None = None) -> str:
+        """跨进程部署时，ticket 已由 HTTP 调用认证中心换成 login_id。"""
+        return await self._manager.stp(self.login_type).login(login_id, device=device)
+
+    async def handle_logout(self, login_id: str) -> None:
+        """收到认证中心的登出通知后，清掉本地会话。"""
+        await self._manager.stp(self.login_type).logout(login_id)

sa_token/storage/__init__.py

@@ -0,0 +1,22 @@
+"""存储层：契约 + 内置实现。
+
+``RedisStorage`` 依赖可选的 redis 包，因此这里按需惰性导入，
+保证只装了核心包的用户 ``import sa_token.storage`` 不会报错。
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import TTL_NEVER_EXPIRE, SaStorage
+from .memory import MemoryStorage
+
+__all__ = ["SaStorage", "MemoryStorage", "RedisStorage", "TTL_NEVER_EXPIRE"]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "RedisStorage":
+        from .redis import RedisStorage
+
+        return RedisStorage
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

sa_token/storage/base.py

@@ -0,0 +1,66 @@
+"""存储契约。
+
+核心层只依赖这套接口，因此内存、Redis、数据库甚至自研存储都能平替。
+所有值都是字符串（JSON），这样 Redis 里的数据可以被其它语言的 sa-token 读写。
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+__all__ = ["SaStorage", "TTL_NEVER_EXPIRE"]
+
+#: ttl 传入该值（或 None）表示永不过期。
+TTL_NEVER_EXPIRE = -1
+
+
+@runtime_checkable
+class SaStorage(Protocol):
+    """键值存储接口。
+
+    ``ttl`` 单位为秒；``None`` 与 ``-1`` 等价，均表示永不过期。
+    """
+
+    async def get(self, key: str) -> str | None: ...
+
+    async def set(self, key: str, value: str, ttl: int | None = None) -> None: ...
+
+    async def delete(self, key: str) -> None: ...
+
+    async def exists(self, key: str) -> bool: ...
+
+    async def expire(self, key: str, ttl: int | None) -> bool:
+        """更新 TTL，键不存在时返回 False。"""
+        ...
+
+    async def ttl(self, key: str) -> int:
+        """返回剩余秒数；-1 表示永不过期，-2 表示键不存在。"""
+        ...
+
+    async def set_if_absent(self, key: str, value: str, ttl: int | None = None) -> bool:
+        """键不存在时写入，返回是否写入成功。用于一次性 token 的原子占位。"""
+        ...
+
+    async def compare_and_set(
+        self,
+        key: str,
+        expected: str,
+        new_value: str,
+        ttl: int | None = None,
+    ) -> bool:
+        """当前值等于 expected 时才写入，用于并发下的安全更新。"""
+        ...
+
+    async def compare_and_delete(self, key: str, expected: str) -> bool: ...
+
+    async def scan(self, pattern: str, cursor: str | None = None, count: int = 100) -> tuple[
+        str | None, list[str]
+    ]:
+        """按 glob 模式游标扫描键，返回 ``(下一个游标, 键列表)``；游标为 None 表示结束。"""
+        ...
+
+    async def clear(self) -> None:
+        """清空全部数据，仅用于测试与开发。"""
+        ...
+
+    async def close(self) -> None: ...

sa_token/storage/memory.py

@@ -0,0 +1,154 @@
+"""内存存储：开发、测试与单机场景的默认实现。
+
+过期采用惰性清理 + 可选的后台清扫，避免为了精确过期而引入常驻线程。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import fnmatch
+import time
+from dataclasses import dataclass
+
+from .base import TTL_NEVER_EXPIRE
+
+__all__ = ["MemoryStorage"]
+
+
+@dataclass
+class _Entry:
+    value: str
+    expire_at: float | None
+
+    def is_expired(self, now: float) -> bool:
+        return self.expire_at is not None and self.expire_at <= now
+
+
+def _to_expire_at(ttl: int | None) -> float | None:
+    if ttl is None or ttl == TTL_NEVER_EXPIRE:
+        return None
+    if ttl <= 0:
+        # 传入 0 或负数（-1 以外）视为立即过期，避免写入永不清理的脏数据。
+        return time.monotonic()
+    return time.monotonic() + ttl
+
+
+class MemoryStorage:
+    """进程内存储。
+
+    注意：数据随进程退出而丢失，多进程部署（gunicorn 多 worker）下各进程互相看不见
+    对方的登录态，生产环境请改用 ``RedisStorage``。
+    """
+
+    def __init__(self, *, cleanup_interval: float = 60.0) -> None:
+        self._data: dict[str, _Entry] = {}
+        self._lock = asyncio.Lock()
+        self._cleanup_interval = cleanup_interval
+        self._last_cleanup = time.monotonic()
+
+    async def get(self, key: str) -> str | None:
+        async with self._lock:
+            return self._get_unlocked(key)
+
+    async def set(self, key: str, value: str, ttl: int | None = None) -> None:
+        async with self._lock:
+            self._data[key] = _Entry(value, _to_expire_at(ttl))
+            self._maybe_cleanup()
+
+    async def delete(self, key: str) -> None:
+        async with self._lock:
+            self._data.pop(key, None)
+
+    async def exists(self, key: str) -> bool:
+        async with self._lock:
+            return self._get_unlocked(key) is not None
+
+    async def expire(self, key: str, ttl: int | None) -> bool:
+        async with self._lock:
+            entry = self._data.get(key)
+            if entry is None or entry.is_expired(time.monotonic()):
+                self._data.pop(key, None)
+                return False
+            entry.expire_at = _to_expire_at(ttl)
+            return True
+
+    async def ttl(self, key: str) -> int:
+        async with self._lock:
+            now = time.monotonic()
+            entry = self._data.get(key)
+            if entry is None or entry.is_expired(now):
+                self._data.pop(key, None)
+                return -2
+            if entry.expire_at is None:
+                return TTL_NEVER_EXPIRE
+            return max(0, int(round(entry.expire_at - now)))
+
+    async def set_if_absent(self, key: str, value: str, ttl: int | None = None) -> bool:
+        async with self._lock:
+            if self._get_unlocked(key) is not None:
+                return False
+            self._data[key] = _Entry(value, _to_expire_at(ttl))
+            return True
+
+    async def compare_and_set(
+        self,
+        key: str,
+        expected: str,
+        new_value: str,
+        ttl: int | None = None,
+    ) -> bool:
+        async with self._lock:
+            if self._get_unlocked(key) != expected:
+                return False
+            self._data[key] = _Entry(new_value, _to_expire_at(ttl))
+            return True
+
+    async def compare_and_delete(self, key: str, expected: str) -> bool:
+        async with self._lock:
+            if self._get_unlocked(key) != expected:
+                return False
+            self._data.pop(key, None)
+            return True
+
+    async def scan(
+        self,
+        pattern: str,
+        cursor: str | None = None,
+        count: int = 100,
+    ) -> tuple[str | None, list[str]]:
+        async with self._lock:
+            now = time.monotonic()
+            keys = sorted(
+                key
+                for key, entry in self._data.items()
+                if not entry.is_expired(now) and fnmatch.fnmatchcase(key, pattern)
+            )
+        start = int(cursor) if cursor else 0
+        page = keys[start : start + count]
+        next_cursor = str(start + count) if start + count < len(keys) else None
+        return next_cursor, page
+
+    async def clear(self) -> None:
+        async with self._lock:
+            self._data.clear()
+
+    async def close(self) -> None:
+        await self.clear()
+
+    def _get_unlocked(self, key: str) -> str | None:
+        entry = self._data.get(key)
+        if entry is None:
+            return None
+        if entry.is_expired(time.monotonic()):
+            self._data.pop(key, None)
+            return None
+        return entry.value
+
+    def _maybe_cleanup(self) -> None:
+        now = time.monotonic()
+        if now - self._last_cleanup < self._cleanup_interval:
+            return
+        self._last_cleanup = now
+        expired = [key for key, entry in self._data.items() if entry.is_expired(now)]
+        for key in expired:
+            self._data.pop(key, None)

sa_token/storage/redis.py

@@ -0,0 +1,136 @@
+"""Redis 存储：生产环境实现。
+
+需要额外安装：``pip install "sa-token-python-core[redis]"``。
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .base import TTL_NEVER_EXPIRE
+
+if TYPE_CHECKING:  # pragma: no cover - 仅供类型检查
+    from redis.asyncio import Redis
+
+__all__ = ["RedisStorage"]
+
+#: 值相等才删除，避免误删被其它请求刷新过的键。
+_COMPARE_AND_DELETE = """
+if redis.call('GET', KEYS[1]) == ARGV[1] then
+    return redis.call('DEL', KEYS[1])
+end
+return 0
+"""
+
+#: 值相等才更新；ARGV[3] 为 -1 时保持 key 永不过期。
+_COMPARE_AND_SET = """
+if redis.call('GET', KEYS[1]) ~= ARGV[1] then
+    return 0
+end
+if tonumber(ARGV[3]) < 0 then
+    redis.call('SET', KEYS[1], ARGV[2])
+else
+    redis.call('SET', KEYS[1], ARGV[2], 'EX', tonumber(ARGV[3]))
+end
+return 1
+"""
+
+
+def _normalize_ttl(ttl: int | None) -> int | None:
+    """把 ``None`` / ``-1`` 统一成「不设置过期」。"""
+    if ttl is None or ttl == TTL_NEVER_EXPIRE:
+        return None
+    return max(1, ttl)
+
+
+class RedisStorage:
+    """基于 ``redis.asyncio`` 的存储实现。
+
+    要求客户端以字符串模式解码（``decode_responses=True``）；使用 :meth:`from_url`
+    构造时会自动设置。
+    """
+
+    def __init__(self, client: Redis) -> None:
+        self._redis = client
+        self._compare_and_delete = client.register_script(_COMPARE_AND_DELETE)
+        self._compare_and_set = client.register_script(_COMPARE_AND_SET)
+
+    @classmethod
+    def from_url(cls, url: str, **kwargs: Any) -> RedisStorage:
+        """``RedisStorage.from_url("redis://localhost:6379/0")``"""
+        try:
+            from redis.asyncio import Redis
+        except ImportError as exc:  # pragma: no cover - 依赖缺失路径
+            raise ImportError(
+                'RedisStorage 需要 redis 依赖，请执行：pip install "sa-token-python-core[redis]"'
+            ) from exc
+        kwargs.setdefault("decode_responses", True)
+        return cls(Redis.from_url(url, **kwargs))
+
+    async def get(self, key: str) -> str | None:
+        return await self._redis.get(key)
+
+    async def set(self, key: str, value: str, ttl: int | None = None) -> None:
+        seconds = _normalize_ttl(ttl)
+        if seconds is None:
+            await self._redis.set(key, value)
+        else:
+            await self._redis.set(key, value, ex=seconds)
+
+    async def delete(self, key: str) -> None:
+        await self._redis.delete(key)
+
+    async def exists(self, key: str) -> bool:
+        return bool(await self._redis.exists(key))
+
+    async def expire(self, key: str, ttl: int | None) -> bool:
+        seconds = _normalize_ttl(ttl)
+        if seconds is None:
+            return bool(await self._redis.persist(key))
+        return bool(await self._redis.expire(key, seconds))
+
+    async def ttl(self, key: str) -> int:
+        return int(await self._redis.ttl(key))
+
+    async def set_if_absent(self, key: str, value: str, ttl: int | None = None) -> bool:
+        seconds = _normalize_ttl(ttl)
+        if seconds is None:
+            return bool(await self._redis.set(key, value, nx=True))
+        return bool(await self._redis.set(key, value, ex=seconds, nx=True))
+
+    async def compare_and_set(
+        self,
+        key: str,
+        expected: str,
+        new_value: str,
+        ttl: int | None = None,
+    ) -> bool:
+        seconds = _normalize_ttl(ttl)
+        result = await self._compare_and_set(
+            keys=[key],
+            args=[expected, new_value, TTL_NEVER_EXPIRE if seconds is None else seconds],
+        )
+        return bool(result)
+
+    async def compare_and_delete(self, key: str, expected: str) -> bool:
+        return bool(await self._compare_and_delete(keys=[key], args=[expected]))
+
+    async def scan(
+        self,
+        pattern: str,
+        cursor: str | None = None,
+        count: int = 100,
+    ) -> tuple[str | None, list[str]]:
+        next_cursor, keys = await self._redis.scan(
+            cursor=int(cursor) if cursor else 0,
+            match=pattern,
+            count=count,
+        )
+        return (str(next_cursor) if next_cursor else None), list(keys)
+
+    async def clear(self) -> None:
+        """清空当前库，仅供测试使用。"""
+        await self._redis.flushdb()
+
+    async def close(self) -> None:
+        await self._redis.aclose()

sa_token/stp_interface.py

@@ -0,0 +1,20 @@
+"""权限数据源。
+
+默认实现从存储里读 ``set_permissions`` 写入的数据，适合中小项目与网关；
+已有 RBAC 表的项目应实现本接口，直接从数据库读，避免两份权限数据不同步。
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+__all__ = ["StpInterface"]
+
+
+@runtime_checkable
+class StpInterface(Protocol):
+    """业务侧权限 / 角色数据源。"""
+
+    async def get_permission_list(self, login_id: str, login_type: str) -> list[str]: ...
+
+    async def get_role_list(self, login_id: str, login_type: str) -> list[str]: ...