data_studio/security.py

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())