SaraView/Backend/Users.py

from __future__ import annotations

import hashlib
import hmac
import secrets
import tomllib

from pathlib import Path
from typing import Any

import tomli_w

from Settings import GetDataDirectory

from Groups import AddMember, AdminGroupName, GetGroupsForUser

Pbkdf2Iterations = 260_000

UserColourOptions = ["Red", "Orange", "Yellow", "Green", "Blue", "Purple"]
DefaultUserColour = "Purple"

AllowedAvatarExtensions = {".png", ".jpg", ".jpeg", ".gif", ".webp"}


def GetUsersPath():
    return GetDataDirectory() / "Users.toml"


def GetAvatarDirectory() -> Path:
    return GetDataDirectory() / "Avatars"


def HashPassword(Password: str) -> str:
    Salt = secrets.token_hex(16)
    DerivedKey = hashlib.pbkdf2_hmac(
        "sha256", Password.encode("utf-8"), bytes.fromhex(Salt), Pbkdf2Iterations
    )
    return f"pbkdf2_sha256${Pbkdf2Iterations}${Salt}${DerivedKey.hex()}"


def VerifyPassword(Password: str, StoredHash: str) -> bool:
    try:
        Algorithm, IterationsText, Salt, HashHex = StoredHash.split("$")
    except ValueError:
        return False

    if Algorithm != "pbkdf2_sha256":
        return False

    DerivedKey = hashlib.pbkdf2_hmac(
        "sha256", Password.encode("utf-8"), bytes.fromhex(Salt), int(IterationsText)
    )
    return hmac.compare_digest(DerivedKey.hex(), HashHex)


def LoadUsers() -> dict[str, dict[str, Any]]:
    UsersPath = GetUsersPath()

    if not UsersPath.is_file():
        return {}

    with UsersPath.open("rb") as UsersFile:
        Loaded = tomllib.load(UsersFile)

    return Loaded.get("Users", {})


def SaveUsers(Users: dict[str, dict[str, Any]]) -> None:
    GetDataDirectory().mkdir(parents=True, exist_ok=True)

    with GetUsersPath().open("wb") as UsersFile:
        tomli_w.dump({"Users": Users}, UsersFile)


def AnyUsersExist() -> bool:
    return len(LoadUsers()) > 0


def GetUser(Username: str) -> dict[str, Any] | None:
    return LoadUsers().get(Username)


def CreateUser(Username: str, Password: str, IsAdmin: bool, Colour: str = DefaultUserColour) -> None:
    if not Username or not Username.strip():
        raise ValueError("Username cannot be empty.")

    if not Password:
        raise ValueError("Password cannot be empty.")

    if Colour not in UserColourOptions:
        Colour = DefaultUserColour

    Users = LoadUsers()

    if Username in Users:
        raise ValueError(f"User '{Username}' already exists.")

    Users[Username] = {
        "PasswordHash": HashPassword(Password),
        "Colour": Colour,
    }
    SaveUsers(Users)

    if IsAdmin:
        AddMember(AdminGroupName, Username)


def UpdateUserColour(Username: str, Colour: str) -> None:
    if Colour not in UserColourOptions:
        raise ValueError(f"Invalid colour: {Colour}")

    Users = LoadUsers()

    if Username not in Users:
        raise ValueError("User not found.")

    Users[Username]["Colour"] = Colour
    SaveUsers(Users)


def RenameUser(OldUsername: str, NewUsername: str) -> None:
    NewUsername = NewUsername.strip()

    if not NewUsername:
        raise ValueError("Username cannot be empty.")

    if NewUsername == OldUsername:
        return

    Users = LoadUsers()

    if OldUsername not in Users:
        raise ValueError("User not found.")

    if NewUsername in Users:
        raise ValueError(f"User '{NewUsername}' already exists.")

    Users[NewUsername] = Users.pop(OldUsername)
    SaveUsers(Users)


def ChangePassword(Username: str, CurrentPassword: str, NewPassword: str) -> None:
    User = GetUser(Username)

    if User is None:
        raise ValueError("User not found.")

    if not VerifyPassword(CurrentPassword, User["PasswordHash"]):
        raise ValueError("Current password is incorrect.")

    if not NewPassword:
        raise ValueError("New password cannot be empty.")

    Users = LoadUsers()
    Users[Username]["PasswordHash"] = HashPassword(NewPassword)
    SaveUsers(Users)


def VerifyLogin(Username: str, Password: str) -> dict[str, Any] | None:
    User = GetUser(Username)

    if User is None:
        return None

    if not VerifyPassword(Password, User["PasswordHash"]):
        return None

    return {
        "Username": Username,
        "IsAdmin": AdminGroupName in GetGroupsForUser(Username),
        "Colour": User.get("Colour", DefaultUserColour),
    }


def ListUsers() -> list[dict[str, Any]]:
    Users = LoadUsers()
    return [
        {
            "Username": Name,
            "IsAdmin": AdminGroupName in GetGroupsForUser(Name),
            "Colour": Data.get("Colour", DefaultUserColour),
        }
        for Name, Data in Users.items()
    ]


def SaveAvatar(Username: str, FileBytes: bytes, Extension: str) -> None:
    Extension = Extension.lower()

    if Extension not in AllowedAvatarExtensions:
        raise ValueError(f"Unsupported image type: {Extension}")

    Directory = GetAvatarDirectory()
    Directory.mkdir(parents=True, exist_ok=True)

    for ExistingPath in Directory.glob(f"{Username}.*"):
        ExistingPath.unlink()

    AvatarPath = Directory / f"{Username}{Extension}"
    AvatarPath.write_bytes(FileBytes)


def GetAvatarPath(Username: str) -> Path | None:
    Directory = GetAvatarDirectory()

    if not Directory.is_dir():
        return None

    Matches = sorted(Directory.glob(f"{Username}.*"))
    return Matches[0] if Matches else None


def RenameUserAvatar(OldUsername: str, NewUsername: str) -> None:
    Directory = GetAvatarDirectory()

    if not Directory.is_dir():
        return

    for ExistingPath in Directory.glob(f"{OldUsername}.*"):
        ExistingPath.rename(Directory / f"{NewUsername}{ExistingPath.suffix}")