data_studio/auth.py

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)