Rare/rare/utils/compat/proton.py

import platform as pf
import os
import shlex
from dataclasses import dataclass
from enum import StrEnum
from hashlib import md5
from logging import getLogger
from typing import Optional, Union, List, Dict, Set

if pf.system() in {"Linux", "FreeBSD"}:
    # noinspection PyUnresolvedReferences
    import vdf  # pylint: disable=E0401

logger = getLogger("Proton")

steam_compat_client_install_paths = [os.path.expanduser("~/.local/share/Steam")]


def find_steam() -> str:
    # return the first valid path
    for path in steam_compat_client_install_paths:
        if os.path.isdir(path) and os.path.isfile(os.path.join(path, "steam.sh")):
            return path


def find_libraries(steam_path: str) -> Set[str]:
    vdf_path = os.path.join(steam_path, "steamapps", "libraryfolders.vdf")
    with open(vdf_path, "r") as f:
        libraryfolders = vdf.load(f)["libraryfolders"]
    # libraries = [os.path.join(folder["path"], "steamapps") for key, folder in libraryfolders.items()]
    libraries = {os.path.join(folder["path"], "steamapps") for key, folder in libraryfolders.items()}
    return libraries


# Notes:
# Anything older than 'Proton 5.13' doesn't have the 'require_tool_appid' attribute.
# Anything older than 'Proton 7.0' doesn't have the 'compatmanager_layer_name' attribute.
# In addition to that, the 'Steam Linux Runtime 1.0 (scout)' runtime lists the
# 'Steam Linux Runtime 2.0 (soldier)' runtime as a dependency and is probably what was
# being used for any version before 5.13.
#
# As a result the following implementation will list versions from 7.0 onwards which honestly
# is a good trade-off for the amount of complexity supporting everything would ensue.


class SteamVerb(StrEnum):
    RUN = "run"
    WAIT_FOR_EXIT_AND_RUN = "waitforexitandrun"
    RUN_IN_PREFIX = "runinprefix"
    DESTROY_PREFIX = "destroyprefix"
    GET_COMPAT_PATH = "getcompatpath"
    GET_NATIVE_PATH = "getnativepath"
    DEFAULT = WAIT_FOR_EXIT_AND_RUN


@dataclass
class SteamBase:
    steam_path: str
    tool_path: str
    toolmanifest: Dict

    def __eq__(self, other):
        return self.tool_path == other.tool_path

    def __hash__(self):
        return hash(self.tool_path)

    @property
    def required_tool(self) -> Optional[str]:
        return self.toolmanifest["manifest"].get("require_tool_appid", None)

    def command(self, verb: SteamVerb = SteamVerb.DEFAULT) -> List[str]:
        tool_path = os.path.normpath(self.tool_path)
        cmd = "".join([shlex.quote(tool_path), self.toolmanifest["manifest"]["commandline"]])
        # NOTE: "waitforexitandrun" seems to be the verb used in by steam to execute stuff
        # `run` is used when setting up the environment, so use that if we are setting up the prefix.
        cmd = cmd.replace("%verb%", str(verb))
        return shlex.split(cmd)

    def as_str(self, verb: SteamVerb = SteamVerb.DEFAULT):
        return " ".join(map(shlex.quote, self.command(verb)))

    @property
    def checksum(self) -> str:
        return md5(self.as_str().encode("utf-8")).hexdigest()


@dataclass
class SteamRuntime(SteamBase):
    steam_library: str
    appmanifest: Dict

    @property
    def name(self) -> str:
        return self.appmanifest["AppState"]["name"]

    @property
    def appid(self) -> str:
        return self.appmanifest["AppState"]["appid"]


@dataclass
class ProtonTool(SteamRuntime):
    runtime: SteamRuntime = None

    def __bool__(self) -> bool:
        if appid := self.required_tool:
            return self.runtime is not None and self.runtime.appid == appid
        return True

    def command(self, verb: SteamVerb = SteamVerb.DEFAULT) -> List[str]:
        cmd = self.runtime.command(verb)
        cmd.extend(super().command(verb))
        return cmd


@dataclass
class CompatibilityTool(SteamBase):
    compatibilitytool: Dict
    runtime: SteamRuntime = None

    def __bool__(self) -> bool:
        if appid := self.required_tool:
            return self.runtime is not None and self.runtime.appid == appid
        return True

    @property
    def name(self) -> str:
        name, data = list(self.compatibilitytool["compatibilitytools"]["compat_tools"].items())[0]
        return data["display_name"]

    def command(self, verb: SteamVerb = SteamVerb.DEFAULT) -> List[str]:
        cmd = self.runtime.command(verb) if self.runtime is not None else []
        cmd.extend(super().command(verb))
        return cmd


def find_appmanifests(library: str) -> List[dict]:
    appmanifests = []
    for entry in os.scandir(library):
        if entry.is_file() and entry.name.endswith(".acf"):
            with open(os.path.join(library, entry.name), "r") as f:
                appmanifest = vdf.load(f)
            appmanifests.append(appmanifest)
    return appmanifests


def find_protons(steam_path: str, library: str) -> List[ProtonTool]:
    protons = []
    appmanifests = find_appmanifests(library)
    common = os.path.join(library, "common")
    for appmanifest in appmanifests:
        folder = appmanifest["AppState"]["installdir"]
        tool_path = os.path.join(common, folder)
        if os.path.isfile(vdf_file := os.path.join(tool_path, "toolmanifest.vdf")):
            with open(vdf_file, "r") as f:
                toolmanifest = vdf.load(f)
                if toolmanifest["manifest"]["compatmanager_layer_name"] == "proton":
                    protons.append(
                        ProtonTool(
                            steam_path=steam_path,
                            steam_library=library,
                            appmanifest=appmanifest,
                            tool_path=tool_path,
                            toolmanifest=toolmanifest,
                        )
                    )
    return protons


def find_compatibility_tools(steam_path: str) -> List[CompatibilityTool]:
    compatibilitytools_paths = {
        "/usr/share/steam/compatibilitytools.d",
        os.path.expanduser(os.path.join(steam_path, "compatibilitytools.d")),
        os.path.expanduser("~/.steam/compatibilitytools.d"),
        os.path.expanduser("~/.steam/root/compatibilitytools.d"),
    }
    compatibilitytools_paths = {
        os.path.realpath(path) for path in compatibilitytools_paths if os.path.isdir(path)
    }
    tools = []
    for path in compatibilitytools_paths:
        for entry in os.scandir(path):
            if entry.is_dir():
                tool_path = os.path.join(path, entry.name)
                tool_vdf = os.path.join(tool_path, "compatibilitytool.vdf")
                manifest_vdf = os.path.join(tool_path, "toolmanifest.vdf")
                if os.path.isfile(tool_vdf) and os.path.isfile(manifest_vdf):
                    with open(tool_vdf, "r") as f:
                        compatibilitytool = vdf.load(f)
                    with open(manifest_vdf, "r") as f:
                        manifest = vdf.load(f)
                    tools.append(
                        CompatibilityTool(
                            steam_path=steam_path,
                            tool_path=tool_path,
                            toolmanifest=manifest,
                            compatibilitytool=compatibilitytool,
                        )
                    )
    return tools


def find_runtimes(steam_path: str, library: str) -> Dict[str, SteamRuntime]:
    runtimes = {}
    appmanifests = find_appmanifests(library)
    common = os.path.join(library, "common")
    for appmanifest in appmanifests:
        folder = appmanifest["AppState"]["installdir"]
        tool_path = os.path.join(common, folder)
        if os.path.isfile(vdf_file := os.path.join(tool_path, "toolmanifest.vdf")):
            with open(vdf_file, "r") as f:
                toolmanifest = vdf.load(f)
                if toolmanifest["manifest"]["compatmanager_layer_name"] == "container-runtime":
                    runtimes.update(
                        {
                            appmanifest["AppState"]["appid"]: SteamRuntime(
                                steam_path=steam_path,
                                steam_library=library,
                                appmanifest=appmanifest,
                                tool_path=tool_path,
                                toolmanifest=toolmanifest,
                            )
                        }
                    )
    return runtimes


def find_runtime(
    tool: Union[ProtonTool, CompatibilityTool], runtimes: Dict[str, SteamRuntime]
) -> Optional[SteamRuntime]:
    required_tool = tool.required_tool
    if required_tool is None:
        return None
    return runtimes.get(required_tool, None)


def get_steam_environment(
    tool: Optional[Union[ProtonTool, CompatibilityTool]] = None, compat_path: Optional[str] = None
) -> Dict:
    # If the tool is unset, return all affected env variable names
    # IMPORTANT: keep this in sync with the code below
    environ = {"STEAM_COMPAT_DATA_PATH": compat_path if compat_path else ""}
    if tool is None:
        environ["STEAM_COMPAT_DATA_PATH"] = ""
        environ["STEAM_COMPAT_CLIENT_INSTALL_PATH"] = ""
        environ["STEAM_COMPAT_LIBRARY_PATHS"] = ""
        environ["STEAM_COMPAT_MOUNTS"] = ""
        environ["STEAM_COMPAT_TOOL_PATHS"] = ""
        return environ

    environ["STEAM_COMPAT_CLIENT_INSTALL_PATH"] = tool.steam_path
    if isinstance(tool, ProtonTool):
        environ["STEAM_COMPAT_LIBRARY_PATHS"] = tool.steam_library
    if tool.runtime is not None:
        compat_mounts = [tool.tool_path, tool.runtime.tool_path]
        environ["STEAM_COMPAT_MOUNTS"] = ":".join(compat_mounts)
    tool_paths = [tool.tool_path]
    if tool.runtime is not None:
        tool_paths.append(tool.runtime.tool_path)
    environ["STEAM_COMPAT_TOOL_PATHS"] = ":".join(tool_paths)
    return environ


def find_tools() -> List[Union[ProtonTool, CompatibilityTool]]:
    steam_path = find_steam()
    logger.debug("Using Steam install in %s", steam_path)
    steam_libraries = find_libraries(steam_path)
    logger.debug("Searching for tools in libraries:")
    logger.debug("%s", steam_libraries)

    runtimes = {}
    for library in steam_libraries:
        runtimes.update(find_runtimes(steam_path, library))

    tools = []
    for library in steam_libraries:
        tools.extend(find_protons(steam_path, library))
    tools.extend(find_compatibility_tools(steam_path))

    for tool in tools:
        runtime = find_runtime(tool, runtimes)
        tool.runtime = runtime

    tools = list(filter(lambda t: bool(t), tools))

    return tools


if __name__ == "__main__":
    from pprint import pprint

    _tools = find_tools()
    pprint(_tools)

    for tool in _tools:
        print(get_steam_environment(tool))
        print(tool.name)
        print(tool.command(SteamVerb.RUN))
        print(" ".join(tool.command(SteamVerb.RUN_IN_PREFIX)))