Launcher/core_lib/game/account.py

"""
A launcher profile module. Contains many useful function to manage minecraft account
INFO: The account data generated by this module is not compatible with official launcher profile!

This module is vibed too!
"""

from __future__ import annotations

import base64
import binascii
import datetime
import json
import logging
import os
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path

from core_lib.common.common import bytes_to_data_url
from core_lib.common.exception import (
    AccountException,
    AccountKeyNotFoundException,
    AccountLoadException,
    AccountSaveException,
)

logger = logging.getLogger("Launcher.CoreLib")


def get_account_sample() -> dict:
    return {"accounts": {}, "activeAccountId": None}


get_accounts_sample = get_account_sample


@dataclass
class Account:
    account_id: str
    username: str
    access_token: str
    refresh_token: str
    access_token_expires_at: str | datetime.datetime
    refresh_token_expires_at: str | datetime.datetime
    minecraft_uuid: str = ""
    minecraft_name: str = ""
    avatar: str = ""
    type: str = "msa"
    user_properties: list = field(default_factory=list)


def _datetime_string(value: str | datetime.datetime, field_name: str) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, datetime.datetime):
        return value.isoformat().replace("+00:00", "Z")
    raise TypeError(f'"{field_name}" must be a datetime.datetime or str')


def create_account_data_from_object(account: Account) -> dict:
    """
    Convert Account object to dictionary
    :param account:
    :return:
        account_dict: dict
    """
    return {
        "accessToken": account.access_token,
        "accessTokenExpiresAt": _datetime_string(
            account.access_token_expires_at, "access_token_expires_at"
        ),
        "refreshToken": account.refresh_token,
        "refreshTokenExpiresAt": _datetime_string(
            account.refresh_token_expires_at, "refresh_token_expires_at"
        ),
        "avatar": account.avatar,
        "accountID": account.account_id,
        "minecraftProfile": {
            "uuid": account.minecraft_uuid,
            "name": account.minecraft_name,
        },
        "type": account.type,
        "userProperties": deepcopy(account.user_properties),
        "username": account.username,
    }

SUPPORT_DATA_URL_IMAGE_TYPE = [
    "image/apng",
    "image/avif",
    "image/bmp",
    "image/gif",
    "image/jpeg",
    "image/png",
    "image/svg+xml",
    "image/tiff",
    "image/vnd.microsoft.icon",  # .ico
    "image/webp",
]

def is_valid_account_avatar(avatar: str | bytes, support_type: list[str] | None=None) -> bool:
    if isinstance(avatar, bytes):
        try:
            avatar = bytes_to_data_url(avatar)
        except (TypeError, ValueError):
            return False

    if not isinstance(avatar, str):
        return False

    # For empty avatar (not set)
    if isinstance(avatar, str) and not avatar:
        return True

    if support_type is None:
        support_type = SUPPORT_DATA_URL_IMAGE_TYPE

    try:
        header, encoded_data = avatar.split(",", 1)
    except ValueError:
        # For keyword avatar (This need launcher to replace the keyword to mapped image)
        return True

    # Check header
    if not header.startswith("data:image/"):
        return False

    if not header.endswith(";base64"):
        return False

    image_type = header.removeprefix("data:").removesuffix(";base64")

    if image_type not in support_type:
        logger.info(f"Unsupported image type: {image_type}")
        return False

    # Try to decode the avatar
    try:
        base64.b64decode(encoded_data, validate=True)
        return True
    except (binascii.Error, ValueError) as exc:
        logger.info(f"Unable to decode account avatar: {exc}")

    return False

def _check_account(account: dict, account_id: str | None = None, fix_wrong: bool = False) -> bool:
    """
    Check if an account is valid.
    :param account:
    :param account_id:
    :param fix_wrong:
    :return:
        is_valid: bool
    """
    if not isinstance(account, dict):
        return False

    required_strings = (
        "accessToken",
        "accessTokenExpiresAt",
        "refreshToken",
        "refreshTokenExpiresAt",
        "accountID",
        "type",
        "username",
    )

    # Check all required key is exiting
    if any(not isinstance(account.get(key), str) or not account[key] for key in required_strings):
        return False

    # Check account id is match as provided
    if account_id is not None and account["accountID"] != account_id:
        return False

    # Verify profile data (Must contain uuid and name)
    minecraft_profile = account.get("minecraftProfile")
    if not isinstance(minecraft_profile, dict):
        return False
    if any(not isinstance(minecraft_profile.get(key), str) for key in ("uuid", "name")):
        return False

    # User properties (Left from official launcher, maybe this field will put some custom account settings in the feture
    properties = account.get("userProperties")
    if properties is None and fix_wrong:
        account["userProperties"] = []
    elif not isinstance(properties, list):
        return False

    # Account avatar
    avatar = account.get("avatar")
    if avatar is None and fix_wrong:
        account["avatar"] = ""
    elif not isinstance(avatar, str):
        return False

    if not is_valid_account_avatar(avatar):
        return False

    return True


def is_account_valid(account_data: dict, fix_wrong: bool = False) -> bool:
    return _check_account(account_data, fix_wrong=fix_wrong)


def is_accounts_valid(accounts: dict, fix_wrong: bool = False, ignore_invalid_account: bool = False) -> bool:
    """
    Check if accountData (contain multiple account and some settings) is valid.
    :param accounts:
    :param fix_wrong:
    :param ignore_invalid_account:
    :return:
        is_valid: bool
    """
    if not isinstance(accounts, dict) or not isinstance(accounts.get("accounts"), dict):
        return False

    account_map = accounts["accounts"]
    active_id = accounts.get("activeAccountId")

    if active_id is not None and not isinstance(active_id, str):
        return False

    invalid_ids = [
        account_id
        for account_id, account in account_map.items()
        if not isinstance(account_id, str)
        or not _check_account(account, account_id=account_id, fix_wrong=fix_wrong)
    ]

    if invalid_ids and not ignore_invalid_account:
        return False

    for account_id in invalid_ids:
        account_map.pop(account_id, None)

    if active_id is not None and active_id not in account_map:
        if fix_wrong or ignore_invalid_account:
            accounts["activeAccountId"] = next(iter(account_map), None)
        else:
            return False

    return True


def list_accounts(accounts: dict) -> dict:
    """
    Get account map from accountData.
    :param accounts:
    :return:
    """
    value = accounts.get("accounts", {}) if isinstance(accounts, dict) else {}
    return value if isinstance(value, dict) else {}


def is_account_exists(accounts: dict, account_id: str) -> bool:
    """
    Check if an account exists in accountData.
    :param accounts:
    :param account_id:
    :return:
    """
    return account_id in list_accounts(accounts)


def get_account(accounts: dict, account_id: str) -> dict:
    """
    Get an account from accountData.
    :param accounts:
    :param account_id:
    :return:
        account_dict: dict
    """
    try:
        return list_accounts(accounts)[account_id]
    except KeyError as exc:
        raise AccountKeyNotFoundException(f"Specified account id {account_id} not found.") from exc


def get_current_account_id(accounts: dict) -> str | None:
    """
    Get current account id from accountData.
    :param accounts:
    :return:
        account_dict: dict
    """
    return accounts.get("activeAccountId") if isinstance(accounts, dict) else None


get_active_account_id = get_current_account_id


def get_current_account(accounts: dict) -> dict | None:
    """
    Get current account from accountData.
    :param accounts:
    :return:
        account_dict: dict
    """
    account_id = get_current_account_id(accounts)

    return get_account(accounts, account_id) if account_id is not None else None


get_active_account = get_current_account


def set_current_account_id(accounts: dict, account_id: str | None) -> None:
    """
    Set current account id
    :param accounts:
    :param account_id:
    :return:
    """
    if account_id is not None and not is_account_exists(accounts, account_id):
        raise AccountKeyNotFoundException(f"Specified account id {account_id} not found.")

    accounts["activeAccountId"] = account_id


set_active_account_id = set_current_account_id


def add_account(accounts: dict, account: dict, *, overwrite: bool = False) -> tuple[dict, str]:
    """
    Add an account to accountData.
    :param accounts:
    :param account:
    :param overwrite:
    :return:
        new_account_data: dict
        new_account_id: str
    """
    if not is_account_valid(account):
        raise AccountException("Provided account data is not valid.")

    account_id = account["accountID"]
    account_map = list_accounts(accounts)

    if account_id in account_map and not overwrite:
        raise AccountException(f"Account id {account_id} already exists.")

    account_map[account_id] = account

    return accounts, account_id


def create_account(account: Account, accounts: dict | None = None) -> tuple[dict, str]:
    """
    Create new account
    :param account:
    :param accounts:
    :return:
        account_dict: dict
        new_account_id: str
    """
    if accounts is None:
        accounts = get_account_sample()

    account_data = create_account_data_from_object(account)
    _, account_id = add_account(accounts, account_data)

    return account_data, account_id


def update_account(accounts: dict, account_id: str, account_data: dict, fix_wrong: bool = False) -> tuple[dict, str]:
    """
    Update target account data value
    :param accounts:
    :param account_id:
    :param account_data:
    :param fix_wrong:
    :return:
        updated_account_data: dict
        updated_account_id: str
    """
    if not is_account_exists(accounts, account_id):
        raise AccountKeyNotFoundException(f"Specified account id {account_id} not found.")

    if not _check_account(account_data, account_id=account_id, fix_wrong=fix_wrong):
        raise AccountException("Provided account data is not valid.")

    accounts["accounts"][account_id] = account_data

    return accounts, account_id


def remove_account(accounts: dict, account_id: str) -> None:
    """
    Remove account from accountData.
    INFO: If the active account ID is same as the account that provided. Active account will be set to None.
    :param accounts:
    :param account_id:
    :return:
    """
    account_map = list_accounts(accounts)

    try:
        del account_map[account_id]
    except KeyError as exc:
        raise AccountKeyNotFoundException(f"Specified account id {account_id} not found.") from exc

    if get_current_account_id(accounts) == account_id:
        accounts["activeAccountId"] = next(iter(account_map), None)


def convert_account_data_to_object(accounts: dict, account_id: str) -> Account:
    """
    Convert account data to Account
    :param accounts:
    :param account_id:
    :return:
        account_object: Account
    """
    data = get_account(accounts, account_id)
    minecraft_profile = data.get("minecraftProfile", {})
    return Account(
        account_id=data.get("accountID", account_id),
        username=data.get("username", ""),
        access_token=data.get("accessToken", ""),
        refresh_token=data.get("refreshToken", ""),
        access_token_expires_at=data.get("accessTokenExpiresAt", ""),
        refresh_token_expires_at=data.get("refreshTokenExpiresAt", ""),
        minecraft_uuid=minecraft_profile.get("uuid", ""),
        minecraft_name=minecraft_profile.get("name", ""),
        avatar=data.get("avatar", ""),
        type=data.get("type", "msa"),
        user_properties=deepcopy(data.get("userProperties", [])),
    )


def read_account_file(account_filepath: str | Path) -> dict:
    """
    Read account file
    :param account_filepath:
    :return:
        accountData: dict
    """
    path = Path(account_filepath)
    if not path.exists():
        raise AccountLoadException(f"Account file does not exist: {path}")
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise AccountLoadException(f"Unable to read account file: {exc}") from exc
    if not is_accounts_valid(data):
        raise AccountLoadException("Account file contains invalid account data.")
    return data


read_accounts_file = read_account_file


def create_or_save_account_file(accounts: dict, account_filepath: str | Path,
                                account_file_bak_filepath: str | Path | None = None,
                                overwrite: bool = True,
                                create_backup: bool = True) -> None:
    """
    Create or save account file
    :param accounts:
    :param account_filepath:
    :param account_file_bak_filepath:
    :param overwrite:
    :param create_backup:
    :return:
    """
    if not is_accounts_valid(accounts):
        raise AccountSaveException("Refusing to save invalid account data.")

    path = Path(account_filepath)
    if path.exists() and not overwrite:
        raise AccountSaveException(f"Account file already exists: {path}")
    backup = (
        Path(account_file_bak_filepath)
        if account_file_bak_filepath is not None
        else path.with_suffix(path.suffix + ".bak")
    )
    temporary = path.with_suffix(path.suffix + ".tmp")

    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        temporary.write_text(
            json.dumps(accounts, indent=4, ensure_ascii=False), encoding="utf-8"
        )
        if create_backup and path.exists():
            if backup.exists():
                backup.unlink()
            os.replace(path, backup)
        os.replace(temporary, path)
    except OSError as exc:
        if create_backup and backup.exists() and not path.exists():
            os.replace(backup, path)
        try:
            temporary.unlink(missing_ok=True)
        except OSError:
            pass
        raise AccountSaveException(f"Unable to save account file: {exc}") from exc

def get_account_avatar(accounts: dict, account_id: str) -> str | None:
    """
    Get account avatar fron raw data
    :param accounts:
    :param account_id:
    :return:
        avatar_data_url: str | None
    """
    account = get_account(accounts, account_id)

    if "avatar" not in account:
        raise AccountKeyNotFoundException(f"Account {account_id} missing 'avatar' field")

    return account.get("avatar", None)

def save_account_avatar(accounts: dict, account_id: str, avatar_raw: bytes | str) -> None:
    """
    Save account avatar to account data
    :param accounts:
    :param account_id:
    :param avatar_raw: Set it to empty string as default
    :return:
    """
    if not is_account_exists(accounts, account_id):
        raise AccountKeyNotFoundException(f"Specified account id {account_id} not found.")

    if isinstance(avatar_raw, bytes):
        try:
            avatar_raw = bytes_to_data_url(avatar_raw)
        except (TypeError, ValueError) as exc:
            raise AccountSaveException("Invalid account avatar bytes.") from exc

    if not is_valid_account_avatar(avatar_raw):
        raise AccountSaveException(f"Invalid account avatar: {avatar_raw}")

    accounts["accounts"][account_id]["avatar"] = avatar_raw

    return None

create_or_save_accounts_file = create_or_save_account_file