feat: add admin session rbac foundation
This commit is contained in:
duanshuwen committed
1 parent
8f000811a8
commit
2cdf7f9d7a
22 files changed
+1945
-68
No files matched your search
@@ -0,0 +1,342 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import secrets
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Protocol
|
||||
|
||||
from .config import get_settings
|
||||
|
||||
|
||||
class RedisUnavailableError(RuntimeError):
|
||||
"""Raised when the administrator session store cannot be reached."""
|
||||
|
||||
|
||||
def hash_refresh_token(token: str) -> str:
|
||||
return hashlib.sha256(token.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def new_refresh_token() -> str:
|
||||
return secrets.token_urlsafe(48)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionRecord:
|
||||
session_id: str
|
||||
user_id: str
|
||||
access_jti: str
|
||||
refresh_hash: str
|
||||
access_expires_at: datetime
|
||||
refresh_expires_at: datetime
|
||||
|
||||
|
||||
class AdminSessionStore(Protocol):
|
||||
def create(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
user_id: str,
|
||||
access_jti: str,
|
||||
refresh_hash: str,
|
||||
access_expires_at: datetime,
|
||||
refresh_expires_at: datetime,
|
||||
) -> None: ...
|
||||
|
||||
def is_access_active(self, session_id: str, access_jti: str) -> bool: ...
|
||||
|
||||
def get_by_refresh_hash(self, refresh_hash: str) -> SessionRecord | None: ...
|
||||
|
||||
def rotate(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
refresh_hash: str,
|
||||
access_jti: str,
|
||||
refresh_hash_next: str,
|
||||
access_expires_at: datetime,
|
||||
refresh_expires_at: datetime,
|
||||
) -> bool: ...
|
||||
|
||||
def revoke(self, session_id: str, refresh_hash: str | None = None) -> None: ...
|
||||
|
||||
def get_permission_context(self, user_id: str) -> dict | None: ...
|
||||
|
||||
def set_permission_context(self, user_id: str, context: dict, ttl_seconds: int) -> None: ...
|
||||
|
||||
def invalidate_permission_cache(self, user_ids: list[str] | None = None) -> None: ...
|
||||
|
||||
def allow_login_attempt(self, identity: str, limit: int, window_seconds: int) -> bool: ...
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
class InMemoryAdminSessionStore:
|
||||
def __init__(self) -> None:
|
||||
self._sessions: dict[str, SessionRecord] = {}
|
||||
self._refresh_index: dict[str, str] = {}
|
||||
self._permission_cache: dict[str, tuple[dict, datetime]] = {}
|
||||
self._login_attempts: dict[str, tuple[int, datetime]] = {}
|
||||
|
||||
def create(self, **kwargs) -> None:
|
||||
record = SessionRecord(**kwargs)
|
||||
self._sessions[record.session_id] = record
|
||||
self._refresh_index[record.refresh_hash] = record.session_id
|
||||
|
||||
def is_access_active(self, session_id: str, access_jti: str) -> bool:
|
||||
record = self._sessions.get(session_id)
|
||||
return bool(record and record.access_jti == access_jti and record.access_expires_at > _now())
|
||||
|
||||
def get_by_refresh_hash(self, refresh_hash: str) -> SessionRecord | None:
|
||||
session_id = self._refresh_index.get(refresh_hash)
|
||||
record = self._sessions.get(session_id) if session_id else None
|
||||
if not record or record.refresh_expires_at <= _now():
|
||||
return None
|
||||
return record
|
||||
|
||||
def rotate(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
refresh_hash: str,
|
||||
access_jti: str,
|
||||
refresh_hash_next: str,
|
||||
access_expires_at: datetime,
|
||||
refresh_expires_at: datetime,
|
||||
) -> bool:
|
||||
record = self._sessions.get(session_id)
|
||||
if not record or record.refresh_hash != refresh_hash or record.refresh_expires_at <= _now():
|
||||
return False
|
||||
self._refresh_index.pop(refresh_hash, None)
|
||||
next_record = SessionRecord(
|
||||
session_id=session_id,
|
||||
user_id=record.user_id,
|
||||
access_jti=access_jti,
|
||||
refresh_hash=refresh_hash_next,
|
||||
access_expires_at=access_expires_at,
|
||||
refresh_expires_at=refresh_expires_at,
|
||||
)
|
||||
self._sessions[session_id] = next_record
|
||||
self._refresh_index[refresh_hash_next] = session_id
|
||||
return True
|
||||
|
||||
def revoke(self, session_id: str, refresh_hash: str | None = None) -> None:
|
||||
record = self._sessions.pop(session_id, None)
|
||||
if record:
|
||||
self._refresh_index.pop(record.refresh_hash, None)
|
||||
if refresh_hash:
|
||||
self._refresh_index.pop(refresh_hash, None)
|
||||
|
||||
def get_permission_context(self, user_id: str) -> dict | None:
|
||||
cached = self._permission_cache.get(user_id)
|
||||
if not cached or cached[1] <= _now():
|
||||
self._permission_cache.pop(user_id, None)
|
||||
return None
|
||||
return cached[0]
|
||||
|
||||
def set_permission_context(self, user_id: str, context: dict, ttl_seconds: int) -> None:
|
||||
self._permission_cache[user_id] = (context, _now() + timedelta(seconds=ttl_seconds))
|
||||
|
||||
def invalidate_permission_cache(self, user_ids: list[str] | None = None) -> None:
|
||||
if user_ids is None:
|
||||
self._permission_cache.clear()
|
||||
return
|
||||
for user_id in user_ids:
|
||||
self._permission_cache.pop(user_id, None)
|
||||
|
||||
def allow_login_attempt(self, identity: str, limit: int, window_seconds: int) -> bool:
|
||||
now = _now()
|
||||
attempts, expires_at = self._login_attempts.get(identity, (0, now))
|
||||
if expires_at <= now:
|
||||
attempts = 0
|
||||
expires_at = now + timedelta(seconds=window_seconds)
|
||||
attempts += 1
|
||||
self._login_attempts[identity] = (attempts, expires_at)
|
||||
return attempts <= limit
|
||||
|
||||
|
||||
class RedisAdminSessionStore:
|
||||
prefix = "wonderq:admin"
|
||||
|
||||
def __init__(self) -> None:
|
||||
try:
|
||||
import redis
|
||||
except ImportError as exc:
|
||||
raise RedisUnavailableError("Redis 客户端未安装") from exc
|
||||
try:
|
||||
self.client = redis.Redis.from_url(get_settings().redis_url, decode_responses=True)
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 会话存储不可用") from exc
|
||||
|
||||
def _session_key(self, session_id: str) -> str:
|
||||
return f"{self.prefix}:session:{session_id}"
|
||||
|
||||
def _refresh_key(self, refresh_hash: str) -> str:
|
||||
return f"{self.prefix}:refresh:{refresh_hash}"
|
||||
|
||||
def _permission_key(self, user_id: str) -> str:
|
||||
return f"{self.prefix}:permission:{user_id}"
|
||||
|
||||
def _login_limit_key(self, identity: str) -> str:
|
||||
return f"{self.prefix}:login-limit:{hashlib.sha256(identity.encode('utf-8')).hexdigest()}"
|
||||
|
||||
@staticmethod
|
||||
def _serialize(record: SessionRecord) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"sessionId": record.session_id,
|
||||
"userId": record.user_id,
|
||||
"accessJti": record.access_jti,
|
||||
"refreshHash": record.refresh_hash,
|
||||
"accessExpiresAt": record.access_expires_at.isoformat(),
|
||||
"refreshExpiresAt": record.refresh_expires_at.isoformat(),
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _deserialize(value: str | None) -> SessionRecord | None:
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(value)
|
||||
return SessionRecord(
|
||||
session_id=payload["sessionId"],
|
||||
user_id=payload["userId"],
|
||||
access_jti=payload["accessJti"],
|
||||
refresh_hash=payload["refreshHash"],
|
||||
access_expires_at=datetime.fromisoformat(payload["accessExpiresAt"]),
|
||||
refresh_expires_at=datetime.fromisoformat(payload["refreshExpiresAt"]),
|
||||
)
|
||||
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise RedisUnavailableError("Redis 会话数据无效") from exc
|
||||
|
||||
@staticmethod
|
||||
def _ttl(expires_at: datetime) -> int:
|
||||
return max(1, int((expires_at - _now()).total_seconds()))
|
||||
|
||||
def _ensure_available(self) -> None:
|
||||
try:
|
||||
self.client.ping()
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 会话存储不可用") from exc
|
||||
|
||||
def create(self, **kwargs) -> None:
|
||||
record = SessionRecord(**kwargs)
|
||||
self._ensure_available()
|
||||
pipe = self.client.pipeline(transaction=True)
|
||||
pipe.set(self._session_key(record.session_id), self._serialize(record), ex=self._ttl(record.refresh_expires_at))
|
||||
pipe.set(self._refresh_key(record.refresh_hash), record.session_id, ex=self._ttl(record.refresh_expires_at))
|
||||
try:
|
||||
pipe.execute()
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 会话创建失败") from exc
|
||||
|
||||
def is_access_active(self, session_id: str, access_jti: str) -> bool:
|
||||
self._ensure_available()
|
||||
record = self._deserialize(self.client.get(self._session_key(session_id)))
|
||||
return bool(record and record.access_jti == access_jti and record.access_expires_at > _now())
|
||||
|
||||
def get_by_refresh_hash(self, refresh_hash: str) -> SessionRecord | None:
|
||||
self._ensure_available()
|
||||
session_id = self.client.get(self._refresh_key(refresh_hash))
|
||||
if not session_id:
|
||||
return None
|
||||
return self._deserialize(self.client.get(self._session_key(session_id)))
|
||||
|
||||
def rotate(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
refresh_hash: str,
|
||||
access_jti: str,
|
||||
refresh_hash_next: str,
|
||||
access_expires_at: datetime,
|
||||
refresh_expires_at: datetime,
|
||||
) -> bool:
|
||||
self._ensure_available()
|
||||
refresh_key = self._refresh_key(refresh_hash)
|
||||
session_key = self._session_key(session_id)
|
||||
try:
|
||||
with self.client.pipeline() as pipe:
|
||||
pipe.watch(refresh_key, session_key)
|
||||
record = self._deserialize(pipe.get(session_key))
|
||||
if not record or record.refresh_hash != refresh_hash or record.refresh_expires_at <= _now():
|
||||
pipe.reset()
|
||||
return False
|
||||
next_record = SessionRecord(
|
||||
session_id=session_id,
|
||||
user_id=record.user_id,
|
||||
access_jti=access_jti,
|
||||
refresh_hash=refresh_hash_next,
|
||||
access_expires_at=access_expires_at,
|
||||
refresh_expires_at=refresh_expires_at,
|
||||
)
|
||||
pipe.multi()
|
||||
pipe.delete(refresh_key)
|
||||
pipe.set(session_key, self._serialize(next_record), ex=self._ttl(refresh_expires_at))
|
||||
pipe.set(self._refresh_key(refresh_hash_next), session_id, ex=self._ttl(refresh_expires_at))
|
||||
pipe.execute()
|
||||
return True
|
||||
except Exception as exc:
|
||||
if exc.__class__.__name__ == "WatchError":
|
||||
return False
|
||||
raise RedisUnavailableError("Redis 会话轮换失败") from exc
|
||||
|
||||
def revoke(self, session_id: str, refresh_hash: str | None = None) -> None:
|
||||
self._ensure_available()
|
||||
record = self._deserialize(self.client.get(self._session_key(session_id)))
|
||||
refresh_key = self._refresh_key(refresh_hash or record.refresh_hash) if record or refresh_hash else None
|
||||
keys = [self._session_key(session_id)]
|
||||
if refresh_key:
|
||||
keys.append(refresh_key)
|
||||
try:
|
||||
self.client.delete(*keys)
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 会话注销失败") from exc
|
||||
|
||||
def get_permission_context(self, user_id: str) -> dict | None:
|
||||
self._ensure_available()
|
||||
try:
|
||||
value = self.client.get(self._permission_key(user_id))
|
||||
if not value:
|
||||
return None
|
||||
payload = json.loads(value)
|
||||
return payload if isinstance(payload, dict) else None
|
||||
except (TypeError, json.JSONDecodeError) as exc:
|
||||
raise RedisUnavailableError("Redis 权限缓存数据无效") from exc
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 权限缓存读取失败") from exc
|
||||
|
||||
def set_permission_context(self, user_id: str, context: dict, ttl_seconds: int) -> None:
|
||||
self._ensure_available()
|
||||
try:
|
||||
self.client.set(self._permission_key(user_id), json.dumps(context), ex=max(1, ttl_seconds))
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 权限缓存写入失败") from exc
|
||||
|
||||
def invalidate_permission_cache(self, user_ids: list[str] | None = None) -> None:
|
||||
self._ensure_available()
|
||||
try:
|
||||
keys = [self._permission_key(user_id) for user_id in user_ids] if user_ids is not None else list(self.client.scan_iter(match=f"{self.prefix}:permission:*"))
|
||||
if keys:
|
||||
self.client.delete(*keys)
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 权限缓存失效失败") from exc
|
||||
|
||||
def allow_login_attempt(self, identity: str, limit: int, window_seconds: int) -> bool:
|
||||
self._ensure_available()
|
||||
key = self._login_limit_key(identity)
|
||||
try:
|
||||
attempts = int(self.client.incr(key))
|
||||
if attempts == 1:
|
||||
self.client.expire(key, max(1, window_seconds))
|
||||
return attempts <= limit
|
||||
except Exception as exc:
|
||||
raise RedisUnavailableError("Redis 登录限流不可用") from exc
|
||||
|
||||
|
||||
def get_admin_session_store() -> AdminSessionStore:
|
||||
return RedisAdminSessionStore()
|
||||
Reference in new issue
Block a user