UWINE/ulwgl_dl_util.py

from pathlib import Path
from os import environ
from requests import get
from tarfile import open as tar_open
from requests import Response
from requests import Timeout
from typing import Dict, List, Tuple, Any, Union
from hashlib import sha512
from shutil import rmtree


def get_ulwgl_proton(env: Dict[str, str]) -> Union[Dict[str, str], None]:
    """Attempt to find Proton and downloads the latest if PROTONPATH is not set.

    Only fetches the latest if not first found in the Steam compat
    Cache is only referred to last
    """
    files: List[Tuple[str, str]] = []

    try:
        files = _fetch_releases()
    except Timeout:
        print("Offline.\nContinuing ...")

    cache: Path = Path(Path().home().as_posix() + "/.cache/ULWGL")
    steam_compat: Path = Path(
        Path().home().as_posix() + "/.local/share/Steam/compatibilitytools.d"
    )

    cache.mkdir(exist_ok=True, parents=True)
    steam_compat.mkdir(exist_ok=True, parents=True)

    # Prioritize the Steam compat
    if _get_from_steamcompat(env, steam_compat, cache, files):
        return env

    # Use the latest Proton in the cache if it exists
    if _get_from_cache(env, steam_compat, cache, files, True):
        return env

    # Download the latest if Proton is not in Steam compat
    # If the digests mismatched, refer to the cache in the next block
    if _get_latest(env, steam_compat, cache, files):
        return env

    # Cache
    # Refer to an old version previously installed
    # Reached on digest mismatch, user interrupt or download failure/no internet
    if _get_from_cache(env, steam_compat, cache, files, False):
        return env

    # No internet and cache/compat tool is empty, just return and raise an exception from the caller
    return env


def _fetch_releases() -> List[Tuple[str, str]]:
    """Fetch the latest releases from the Github API."""
    resp: Response = get(
        "https://api.github.com/repos/GloriousEggroll/proton-ge-custom/releases",
        timeout=30,
    )
    # The file name and its URL as one element
    # Checksum will be the first element, GE-Proton second
    files: List[Tuple[str, str]] = []

    if not resp or not resp.status_code == 200:
        return files

    # Attempt to acquire the tarball and checksum from the JSON data
    releases: List[Dict[str, Any]] = resp.json()
    for release in releases:
        if "assets" in release:
            assets: List[Dict[str, Any]] = release["assets"]

            for asset in assets:
                if (
                    "name" in asset
                    and (
                        asset["name"].endswith("sum")
                        or (
                            asset["name"].endswith("tar.gz")
                            and asset["name"].startswith("GE-Proton")
                        )
                    )
                    and "browser_download_url" in asset
                ):
                    if asset["name"].endswith("sum"):
                        files.append((asset["name"], asset["browser_download_url"]))
                    else:
                        files.append((asset["name"], asset["browser_download_url"]))

                if len(files) == 2:
                    break
        break

    return files


def _fetch_proton(
    env: Dict[str, str], steam_compat: Path, cache: Path, files: List[Tuple[str, str]]
) -> Dict[str, str]:
    """Download the latest ULWGL-Proton and set it as PROTONPATH."""
    hash, hash_url = files[0]
    proton, proton_url = files[1]
    stored_digest: str = ""

    # TODO: Parallelize this
    print(f"Downloading {hash} ...")
    resp_hash: Response = get(hash_url, timeout=30)
    print(f"Downloading {proton} ...")
    resp: Response = get(proton_url, timeout=150)
    if (
        not resp_hash
        and resp_hash.status_code != 200
        and not resp
        and resp.status_code != 200
    ):
        err: str = "Failed.\nFalling back to cache directory ..."
        raise ValueError(err)

    print("Completed.")

    # Download the hash
    with Path(f"{cache.as_posix()}/{hash}").open(mode="wb") as file:
        file.write(resp_hash.content)
    stored_digest = Path(f"{cache.as_posix()}/{hash}").read_text().split(" ")[0]

    # If checksum fails, raise an error and fallback to the cache
    with Path(f"{cache.as_posix()}/{proton}").open(mode="wb") as file:
        file.write(resp.content)

        if sha512(resp.content).hexdigest() != stored_digest:
            err: str = "Digests mismatched.\nFalling back to the cache ..."
            raise ValueError(err)
        print(f"{proton}: SHA512 is OK")

    _extract_dir(Path(f"{cache.as_posix()}/{proton}"), steam_compat)
    environ["PROTONPATH"] = steam_compat.joinpath(proton[: proton.find(".")]).as_posix()
    env["PROTONPATH"] = environ["PROTONPATH"]

    return env


def _extract_dir(proton: Path, steam_compat: Path) -> None:
    """Extract from the cache and to another location."""
    with tar_open(proton.as_posix(), "r:gz") as tar:
        print(f"Extracting {proton} -> {steam_compat.as_posix()} ...")
        tar.extractall(path=steam_compat.as_posix())
        print("Completed.")


def _cleanup(tarball, proton, cache, steam_compat) -> None:
    """Remove files that may have been left in an incomplete state to avoid corruption.

    We want to do this when a download for a new release is interrupted
    """
    print("Keyboard Interrupt received.\nCleaning ...")

    if cache.joinpath(tarball).is_file():
        print(f"Purging {tarball} in {cache} ...")
        cache.joinpath(tarball).unlink()
    if steam_compat.joinpath(proton).is_dir():
        print(f"Purging {proton} in {steam_compat} ...")
        rmtree(steam_compat.joinpath(proton).as_posix())


def _get_from_steamcompat(
    env: Dict[str, str], steam_compat: Path, cache: Path, files: List[Tuple[str, str]]
) -> Dict[str, str]:
    """Refer to Steam compat folder for any existing Proton directories."""
    for proton in steam_compat.glob("GE-Proton*"):
        print(f"{proton.name} found in: {steam_compat.as_posix()}")
        environ["PROTONPATH"] = proton.as_posix()
        env["PROTONPATH"] = environ["PROTONPATH"]

        # Notify the user that they're not using the latest
        if len(files) == 2 and proton.name != files[1][0][: files[1][0].find(".")]:
            print(
                "GE-Proton is outdated and requires manual intervention.\nFor latest release, please download "
                + files[1][1]
            )

        return env

    return None


def _get_from_cache(
    env: Dict[str, str],
    steam_compat: Path,
    cache: Path,
    files: List[Tuple[str, str]],
    latest=True,
) -> Dict[str, str]:
    """Refer to ULWGL cache directory.

    Use the latest in the cache when present. Older Proton versions are only referred to when: digests mismatch, user interrupt, or download failure/no internet
    """
    if files and latest:
        tarball: str = files[1][0]  # GE-Proton*.tar.gz
        proton: str = tarball[: tarball.find(".")]  # GE-Proton\d+\-\d\d

        print(tarball + " found in: " + cache.as_posix())
        try:
            _extract_dir(
                Path(Path().home().as_posix() + "/.cache/ULWGL").joinpath(tarball),
                steam_compat,
            )

            # Set PROTONPATH to .local/share/Steam/compatibilitytools.d/GE-Proton*
            environ["PROTONPATH"] = steam_compat.joinpath(proton).as_posix()
            env["PROTONPATH"] = environ["PROTONPATH"]

            return env
        except KeyboardInterrupt:
            # Exit cleanly
            # Clean up only the extracted data
            if steam_compat.joinpath(proton).is_dir():
                print(f"Purging {proton} in {steam_compat} ...")
                rmtree(steam_compat.joinpath(proton).as_posix())

            raise

    # Refer to an old version previously installed
    # Reached on digest mismatch, user interrupt or download failure/no internet
    for tarball in cache.glob("GE-Proton*.tar.gz"):
        # Ignore the mismatched file
        if files and tarball == cache.joinpath(files[1][0]):
            continue

        print(f"{tarball.name} found in: {cache.as_posix()}")
        try:
            _extract_dir(tarball, steam_compat)
            environ["PROTONPATH"] = steam_compat.joinpath(
                tarball.name[: tarball.name.find(".")]
            ).as_posix()
            env["PROTONPATH"] = environ["PROTONPATH"]

            break
        except KeyboardInterrupt:
            proton: str = tarball.name[: tarball.name.find(".")]

            if steam_compat.joinpath(proton).is_dir():
                print(f"Purging {proton} in {steam_compat} ...")
                rmtree(steam_compat.joinpath(proton).as_posix())
            raise

    return env


def _get_latest(
    env: Dict[str, str], steam_compat: Path, cache: Path, files: List[Tuple[str, str]]
) -> Dict[str, str]:
    """Download the latest Proton for new installs -- empty cache and Steam compat.

    When the digests mismatched or when interrupted, refer to cache for an old version
    """
    if files:
        tarball: str = files[1][0]

        print("Fetching latest release ...")
        try:
            _fetch_proton(env, steam_compat, cache, files)
            env["PROTONPATH"] = environ["PROTONPATH"]
        except ValueError:
            # Digest mismatched or download failed
            # Refer to the cache for old version next
            return None
        except KeyboardInterrupt:
            # Exit cleanly
            # Clean up extracted data and cache to prevent corruption/errors
            # Refer to the cache for old version next
            _cleanup(tarball, tarball[: tarball.find(".")], cache, steam_compat)
            return None

    return env