from __future__ import annotations import base64 import hashlib import hmac import os import secrets from typing import Any from pymongo.database import Database from data_studio.utils import utc_now_iso PBKDF2_ITERATIONS = 600_000 PBKDF2_ALGORITHM = "sha256" def normalize_username(username: str) -> str: return username.strip().lower() def single_user_mode_enabled() -> bool: return os.getenv("SINGLE_USER_MODE", "false").strip().lower() in {"1", "true", "yes"} def single_user_username() -> str: return normalize_username(os.getenv("SINGLE_USER_USERNAME", "")) def admin_usernames() -> set[str]: raw = os.getenv("ADMIN_USERNAMES", "prashanth").strip() if not raw: return set() return {normalize_username(item) for item in raw.split(",") if item.strip()} def is_admin_username(username: str) -> bool: return normalize_username(username) in admin_usernames() def is_single_user_owner(username: str) -> bool: return single_user_mode_enabled() and normalize_username(username) == single_user_username() def _pbkdf2_hash(password: str, salt: bytes, iterations: int = PBKDF2_ITERATIONS) -> bytes: return hashlib.pbkdf2_hmac( PBKDF2_ALGORITHM, password.encode("utf-8"), salt, iterations, ) def hash_password(password: str) -> dict[str, Any]: salt = os.urandom(16) password_hash = _pbkdf2_hash(password, salt) return { "algorithm": f"pbkdf2_{PBKDF2_ALGORITHM}", "iterations": PBKDF2_ITERATIONS, "salt_b64": base64.b64encode(salt).decode("ascii"), "hash_b64": base64.b64encode(password_hash).decode("ascii"), } def generate_temporary_password(length: int = 16) -> str: alphabet = "ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz23456789" return "".join(secrets.choice(alphabet) for _ in range(length)) def verify_password(password: str, password_record: dict[str, Any]) -> bool: salt = base64.b64decode(password_record["salt_b64"]) expected = base64.b64decode(password_record["hash_b64"]) iterations = int(password_record.get("iterations", PBKDF2_ITERATIONS)) candidate = _pbkdf2_hash(password, salt, iterations) return hmac.compare_digest(candidate, expected) def serialize_user(user: dict[str, Any] | None) -> dict[str, Any] | None: if user is None: return None normalized_username = normalize_username(user["username"]) role = "owner" if is_single_user_owner(normalized_username) else "admin" if is_admin_username(normalized_username) else user.get("role", "reviewer") return { "username": normalized_username, "display_name": user.get("display_name") or normalized_username, "role": role, "is_active": bool(user.get("is_active", True)), "last_login_at": user.get("last_login_at"), } def count_users(db: Database) -> int: return db["users"].count_documents({}) def get_user(db: Database, username: str) -> dict[str, Any] | None: return db["users"].find_one({"username": normalize_username(username)}) def get_active_user(db: Database, username: str) -> dict[str, Any] | None: user = get_user(db, username) if user is None or not bool(user.get("is_active", True)): return None return user def create_or_update_user( db: Database, username: str, password: str, display_name: str | None = None, role: str = "reviewer", is_active: bool = True, update_existing: bool = False, ) -> dict[str, Any]: normalized_username = normalize_username(username) if single_user_mode_enabled() and normalized_username != single_user_username(): raise ValueError( f"Single-user mode only allows the configured owner account: {single_user_username()}" ) existing = get_user(db, normalized_username) if existing is not None and not update_existing: raise ValueError(f"User already exists: {normalized_username}") now = utc_now_iso() password_record = hash_password(password) payload = { "username": normalized_username, "display_name": (display_name or normalized_username).strip(), "role": "admin" if (is_single_user_owner(normalized_username) or is_admin_username(normalized_username)) else (role.strip() or "reviewer"), "is_active": bool(is_active), "password": password_record, "updated_at": now, } update_doc = { "$set": payload, "$setOnInsert": {"created_at": now, "last_login_at": None}, } db["users"].update_one({"username": normalized_username}, update_doc, upsert=True) return serialize_user(get_user(db, normalized_username)) def authenticate_user(db: Database, username: str, password: str) -> dict[str, Any] | None: normalized_username = normalize_username(username) if single_user_mode_enabled() and normalized_username != single_user_username(): return None user = get_active_user(db, normalized_username) if user is None: return None if not verify_password(password, user["password"]): return None now = utc_now_iso() db["users"].update_one({"username": user["username"]}, {"$set": {"last_login_at": now}}) user["last_login_at"] = now return serialize_user(user)