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}")