343 lines
13 KiB
Python
343 lines
13 KiB
Python
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()
|