from __future__ import annotations import hashlib import secrets from datetime import timedelta from typing import Any from pymongo.database import Database from data_studio.auth import authenticate_user, get_active_user, is_admin_username, is_single_user_owner, serialize_user from data_studio.utils import utc_now SESSION_HOURS = 12 ROLE_PERMISSIONS = { "reviewer": {"review"}, "manager": { "review", "import", "export", "manage_sources", "governed_import", "manage_rubrics", "manage_benchmarks", "manage_jobs", "audit", }, "admin": { "review", "import", "export", "manage_sources", "manage_users", "governed_import", "manage_rubrics", "manage_benchmarks", "manage_jobs", "audit", "approve_sources", "promote", "revoke", }, } ALL_PERMISSIONS = sorted({permission for permissions in ROLE_PERMISSIONS.values() for permission in permissions}) def _token_hash(token: str) -> str: return hashlib.sha256(token.encode("utf-8")).hexdigest() def permissions_for_role(role: str) -> list[str]: return sorted(ROLE_PERMISSIONS.get(role, set())) def permissions_for_user(user: dict[str, Any]) -> list[str]: if is_single_user_owner(user["username"]) or is_admin_username(user["username"]): return ALL_PERMISSIONS return permissions_for_role(user["role"]) def create_session(db: Database, username: str, password: str) -> tuple[str, dict[str, Any]] | None: user = authenticate_user(db, username, password) if user is None: return None token = secrets.token_urlsafe(32) now = utc_now() db["sessions"].insert_one( { "token_hash": _token_hash(token), "username": user["username"], "created_at": now, "expires_at": now + timedelta(hours=SESSION_HOURS), } ) user["permissions"] = permissions_for_user(user) return token, user def get_session_user(db: Database, token: str) -> dict[str, Any] | None: now = utc_now() session = db["sessions"].find_one({"token_hash": _token_hash(token), "expires_at": {"$gt": now}}) if session is None: return None user = serialize_user(get_active_user(db, session["username"])) if user is None: return None user["permissions"] = permissions_for_user(user) return user def delete_session(db: Database, token: str) -> None: db["sessions"].delete_one({"token_hash": _token_hash(token)}) def has_permission(user: dict[str, Any], permission: str) -> bool: if is_single_user_owner(user["username"]) or is_admin_username(user["username"]): return True return permission in ROLE_PERMISSIONS.get(user.get("role", "reviewer"), set())