LC_ModManager/ModManager.py

import glob
import hashlib
import json
import logging
import os
import shutil
import sys
import tempfile
from typing import Tuple, List
from zipfile import ZipFile

from PyQt6 import QtWidgets
from PyQt6.QtWidgets import QApplication

from files.settings import LCSettings
from window.main_window import MainWindow


class ModManager:
    VERSION = "0.2"

    def __init__(self, log_level: int = logging.INFO):
        self.__logger = logging.getLogger("ModManager")
        sh = logging.StreamHandler(sys.stdout)
        sh.setFormatter(logging.Formatter(
            fmt="%(asctime)s.%(msecs)03d [%(name)-12.12s] [%(levelname)-5.5s]  %(message)s"))
        self.__logger.addHandler(sh)
        self.__logger.setLevel(log_level)

        self.__app = QApplication([])
        self.__window = MainWindow(self, self.__logger, version=ModManager.VERSION)
        self.__settings = LCSettings()

        self.available_mods = dict()
        self.installed_mods = dict()

        if self.__settings.get_game_folder() is None:
            dialog = QtWidgets.QMessageBox()
            dialog.setWindowTitle("No Game folder")
            dialog.setInformativeText(f"The mod manager requires the path to your Lethal Company game folder!\n"
                                      "Please choose it in the following window.")
            dialog.setIcon(QtWidgets.QMessageBox.Icon.Information)
            dialog.exec()
            self.__window.on_action_set_game_folder()
        if self.__settings.get_game_folder() is None:
            dialog = QtWidgets.QMessageBox()
            dialog.setWindowTitle("No Game folder")
            dialog.setInformativeText(f"Can't do anything without a valid game folder path.\nExiting!")
            dialog.setIcon(QtWidgets.QMessageBox.Icon.Information)
            dialog.exec()
            self.__app.exit(0)
            sys.exit(0)

        self.create_manager_folder()
        self.index_stored_mods()
        self.index_installed_mods()

    def run(self):
        self.__logger.debug(f"Starting LCMod Manager {ModManager.VERSION} ...")
        self.__window.show()

        status = self.__app.exec()
        sys.exit(status)

    @staticmethod
    def get_mod_key(mod_name: str, version: str) -> Tuple[str, str]:
        """
        Creates the mod_key used for referencing this mod in the dictionaries
        :param mod_name:
        :param version:
        :return:
        """
        return mod_name, version

    @staticmethod
    def __adjust_folder_paths(path: str, includes_bepinx=True, includes_folders=True) -> str:
        """
        Adjusts a given path containing common issues to one compatible with BepInEx

        :param path:
        :param includes_bepinx: bool toggle if the mod archive included paths with "BepInEx", if not it will be added
        :param includes_folders: bool toggle if the mod archive included paths with folders, if not within plugins will be assumed
        :return:
        """
        if not includes_bepinx:
            if path.startswith("config") or path.startswith("plugins"):
                path = os.path.join("BepInEx", path)
            else:
                path = os.path.join("BepInEx", "plugins", path)

        if path.endswith(".dll") and "BepInEx" not in path:
            path = os.path.join("BepInEx", "plugins", path)
        # if path.startswith("config" + os.path.sep) or path.startswith("plugins" + os.path.sep):
        #    path = os.path.join("BepInEx", path)
        if path.lower().startswith("bepinex") and not path.startswith("BepInEx"):
            path = "BepInEx" + path[len("BepinEx"):]

        return path

    def create_manager_folder(self):
        game_folder = self.__settings.get_game_folder()
        if game_folder is not None:
            manager_folder = os.path.join(game_folder, LCSettings.FOLDER_NAME)
            if not os.path.isdir(manager_folder):
                self.__logger.debug(f"Creating folder: {manager_folder}")
                os.mkdir(manager_folder)

    def is_folder_valid_game_folder(self, dir_path: str) -> bool:
        req_files = ["Lethal Company.exe", "doorstop_config.ini"]
        req_folders = ["Lethal Company_Data", "MonoBleedingEdge", "BepInEx"]
        for entry in os.listdir(dir_path):
            fullpath_entry = os.path.join(dir_path, entry)
            if os.path.isfile(fullpath_entry):
                if entry in req_files:
                    req_files.remove(entry)
            elif os.path.isdir(fullpath_entry):
                if entry in req_folders:
                    req_folders.remove(entry)

        if len(req_folders) == len(req_files) == 0:
            return True
        else:
            if len(req_files) == 1 and req_files[0] == "doorstop_config.ini" or \
                    len(req_folders) == 1 and req_folders[0] == "BepInEx":
                self.__logger.error(f"new game path rejected! Following files and folders were "
                                    f"expected but not present: {req_files + req_folders}\nBepInEx not installed?")
                return False
            self.__logger.error(f"new game path rejected! Following files and folders were "
                                f"expected but not present: {req_files + req_folders}")
            return False

    def set_game_folder(self, dir_path: str) -> bool:
        """
        Sets the game_path to dir_path in the settings
        :param dir_path:
        :return: indicates if the new game folder path was accepted
        """
        if len(dir_path) == 0:
            self.__logger.debug("new game path selection got probably cancelled.")
            return True

        is_valid = self.is_folder_valid_game_folder(dir_path)
        if is_valid:
            s = self.__settings.get_settings()
            s["game_path"] = dir_path
            self.__settings.apply_changes()
            self.create_manager_folder()
            self.index_stored_mods()
            self.index_installed_mods()
        return is_valid

    def is_mod_installed(self, mod_name: str, mod_version: str) -> bool:
        """
        Checks if a given mod has been installed to the game already
        :param mod_name:
        :param mod_version:
        :return:
        """
        r = self.available_mods[(mod_name, mod_version)]
        mod_files = r["mod_files"]
        for i in range(len(mod_files)):
            file = mod_files[i]
            orig_file = r["orig_mod_files"][i]
            if not os.path.exists(os.path.join(self.__settings.get_game_folder(), file)):
                self.__logger.debug(f"Checking if mod \"{mod_name} - {mod_version}\" is installed ... FALSE")
                return False
            else:
                if os.path.isfile(os.path.join(self.__settings.get_game_folder(), file)):
                    hash_installed = self.get_file_hash(
                        open(os.path.join(self.__settings.get_game_folder(), file), 'rb'))
                    modzip = os.path.join(self.__settings.get_mod_folder(), self.available_mods[(mod_name, mod_version)]["path"])
                    hash_in_storage = self.get_file_hash(ZipFile(modzip).open(orig_file, 'r'))
                    if hash_installed != hash_in_storage:
                        self.__logger.debug(f"Checking if mod \"{mod_name} - {mod_version}\" is installed ... FALSE")
                        return False
        self.__logger.debug(f"Checking if mod \"{mod_name} - {mod_version}\" is installed ... TRUE")
        return True

    def is_valid_mod_file(self, file_path: str):
        """
        Checks if the given file is a valid mod zip file
        :param file_path:
        :return:
        """
        if not os.path.isfile(file_path):
            return False
        zip = ZipFile(file_path)
        contents = zip.namelist()

        if "manifest.json" in contents:
            return True
        return False

    def get_file_hash(self, file_buffer) -> str:
        md5_obj = hashlib.md5()
        while True:
            buffer = file_buffer.read(8096)
            if not buffer:
                break
            md5_obj.update(buffer)
        return md5_obj.hexdigest()

    def get_mod_hash(self, file_path: str) -> str:
        md5_obj = hashlib.md5()
        with open(file_path, 'rb') as file:
            while True:
                buffer = file.read(8096)
                if not buffer:
                    break
                md5_obj.update(buffer)
        return md5_obj.hexdigest()

    def get_mod_info(self, file_path: str) -> tuple[str, str, list[str], list[str]]:
        """
        Returns the name and version string of a given mod file
        :param file_path:
        :return:
        """
        self.__logger.debug(f"Trying to get mod info of file \"{file_path}\" ...")
        if not self.is_valid_mod_file(file_path):
            self.__logger.error(f"Tried to get mod info of an invalid file: {file_path}")
            raise ValueError(f"Tried to get mod info of an invalid file: {file_path}")

        zip = ZipFile(file_path)
        f = zip.open('manifest.json')
        contents = zip.namelist()

        contains_bepinex = False
        contains_folders = False
        for i in range(len(contents)):
            if "/" in contents[i]:
                contains_folders = True
            if contents[i].lower().startswith("bepinex"):
                contains_bepinex = True

        orig_files = []
        files = []
        for file in contents:
            if "icon.png" in file or "manifest.json" in file or file.endswith(".md"):
                continue
            orig_files.append(file)
            file = self.__adjust_folder_paths(file, includes_bepinx=contains_bepinex, includes_folders=contains_folders)

            files.append(file)

        manifest = json.load(f)
        return str(manifest["name"]), str(manifest["version_number"]), files, orig_files

    def index_installed_mods(self):
        """
        Checks all installed mods against the known mod files.

        For unknown mods a placeholder is created
        :return:
        """
        self.__logger.debug("Indexing all installed mods ...")
        self.installed_mods = dict()
        if self.__settings.get_game_folder() is None:
            return

        files = []
        to_search = [self.__settings.get_plugin_folder()]
        while len(to_search) > 0:
            curr = to_search.pop()
            if os.path.isdir(curr):
                r = os.listdir(curr)
                for i in r:
                    to_search.append(os.path.join(curr, i))
            elif os.path.isfile(curr) and curr.endswith('.dll'):
                files.append(curr)

        for mod_name, mod_version in self.available_mods.keys():
            if self.is_mod_installed(mod_name, mod_version):
                self.installed_mods[(mod_name, mod_version)] = self.available_mods[(mod_name, mod_version)]

        unknown_mod = dict()
        unresolved_files = files.copy()
        for file in files:
            for mod in self.installed_mods.keys():
                file_rel = file[len(self.__settings.get_game_folder()) + 1:]
                if file_rel.replace(os.path.sep, "/") in self.installed_mods[mod]["mod_files"]:
                    unresolved_files.remove(file)
            if file in unresolved_files:
                unknown_mod[os.path.basename(file)] = {"mod_files": [file]}
        for key in unknown_mod.keys():
            self.installed_mods[key] = unknown_mod[key]

        self.__window.set_installed_mods(self.installed_mods)
        self.__window.set_available_mods(self.available_mods)

    def index_stored_mods(self):
        """
        Goes through all mods in ModManager.FOLDER_NAME and tries to add them to the available mod list.

        Ignores mods that are already on the list
        :return:
        """
        self.__logger.debug("Indexing all stored mods ...")
        self.available_mods = dict()
        if self.__settings.get_game_folder() is None:
            return
        for file in os.listdir(os.path.join(self.__settings.get_game_folder(), LCSettings.FOLDER_NAME)):
            full_path = os.path.join(self.__settings.get_mod_folder(), str(file))
            if not self.is_valid_mod_file(full_path):
                self.__logger.warning(f"File {file} is not a valid mod file but inside the mod storage folder!")
            mod_name, mod_version, mod_files, orig_mod_files = self.get_mod_info(full_path)
            if mod_name is None:
                self.__logger.warning(f"Mod \"{full_path}\" did not have the expected path. Ignoring it ...")
                continue

            self.available_mods[ModManager.get_mod_key(mod_name, mod_version)] = {"path": file, "version": mod_version,
                                                                                  "mod_files": mod_files,
                                                                                  "orig_mod_files": orig_mod_files}
        self.__window.set_available_mods(self.available_mods)

    def add_mod_file(self, file_path: str):
        """
        Adds the mod to the mod list
        :param file_path:
        :return:
        """
        self.__logger.debug(f"Trying to add mod file \"{file_path}\" to manager ...")
        if self.__settings.get_settings()["game_path"] is None:
            self.__logger.error("Can't add a mod without a valid game path!")

        if not self.is_valid_mod_file(file_path):
            self.__logger.warning(f"File {file_path} was not a mod file!")
            return

        dst = os.path.join(self.__settings.get_mod_folder(),
                           os.path.basename(file_path))
        if os.path.isfile(dst):
            hash1 = self.get_mod_hash(file_path)
            hash2 = self.get_mod_hash(dst)

            if hash1 != hash2:
                self.__logger.info("Given file is different than the one stored. Overwriting it!")
            else:
                self.__logger.info("Given file is the same as the one stored. Ignoring it!")
                return
        self.create_manager_folder()
        self.__logger.info(f"File {file_path} added as a mod file ...")
        shutil.copy(file_path,
                    os.path.join(self.__settings.get_mod_folder(), os.path.basename(file_path)))
        self.index_stored_mods()
        self.index_installed_mods()

    def install_mod(self, mod_name: str, mod_version: str):
        """
        Installs the given mod by extracting all files into the game directory.

        This assumes that the .zip is structured like the game folder as all mods should be
        :param mod_name:
        :param mod_version:
        :return:
        """
        self.__logger.debug(f"Trying to install mod \"{mod_name} - {mod_version}\"")
        if (mod_name, mod_version) not in self.available_mods.keys():
            self.__logger.critical(f"Tried to install a mod that doesn't exist: {mod_name} - {mod_version}")
            return
        self.__logger.info(f"Installing mod \"{mod_name} - {mod_version}\" ...")
        mod_zip = self.available_mods[(mod_name, mod_version)]["path"]

        with ZipFile(os.path.join(self.__settings.get_mod_folder(), mod_zip), 'r') as zip_ref:
            contents = zip_ref.namelist()

            contains_bepinex = False
            contains_folders = False
            for i in range(len(contents)):
                if "/" in contents[i]:
                    contains_folders = True
                if contents[i].lower().startswith("bepinex"):
                    contains_bepinex = True

            with tempfile.TemporaryDirectory() as tmp_dir:
                for file in contents:
                    content = file
                    if "icon.png" in file or "manifest.json" in file or file.endswith(".md"):
                        continue
                    file = self.__adjust_folder_paths(file, includes_bepinx=contains_bepinex,
                                                      includes_folders=contains_folders)

                    zip_ref.extract(content, tmp_dir)
                    if content.endswith(os.path.sep):
                        self.__logger.debug(f"Skipped moving of {os.path.join(tmp_dir, content)}")
                        continue
                    parent_dir = os.path.join(self.__settings.get_game_folder(), file).split(os.path.basename(file))[0]\
                        .replace("/", os.path.sep)
                    if not os.path.exists(parent_dir):
                        os.mkdir(parent_dir)
                    if not os.path.exists(os.path.join(self.__settings.get_game_folder().replace("/", os.path.sep),
                                                       file.replace("/", os.path.sep))):
                        self.__logger.debug(
                            f"Extracting file \"{os.path.join(tmp_dir, content.replace('/', os.path.sep))}\" to " +
                            f"\"{os.path.join(self.__settings.get_game_folder().replace('/', os.path.sep), file.replace('/', os.path.sep))}\"")
                        shutil.move(os.path.join(tmp_dir, content.replace("/", os.path.sep)),
                                    os.path.join(self.__settings.get_game_folder().replace("/", os.path.sep),
                                                 file.replace("/", os.path.sep)))
        self.index_installed_mods()

    def uninstall_mod(self, mod_name: str, mod_version: str):
        """
        Uninstalls the given mod by removing all files (not folders) that the manager is aware of

        Note: For untracked mods this will only be the .dll in the plugins folder
        :param mod_name:
        :param mod_version:
        :return:
        """
        self.__logger.info(f"Uninstalling mod \"{mod_name} - {mod_version}\" ...")
        for file in self.installed_mods[(mod_name, mod_version)]["mod_files"]:
            full_path = os.path.join(self.__settings.get_game_folder(), file)
            if os.path.isfile(full_path):
                self.__logger.debug(f"Deleting file \"{full_path}\" ...")
                os.remove(full_path)
        self.index_installed_mods()

    def nuke_manager_files(self):
        self.__logger.info("Deleting all manager related files ...")
        self.__logger.debug(f"Deleting folder \"{self.__settings.get_mod_folder()}\"")
        files = glob.glob(os.path.join(self.__settings.get_mod_folder(), '*.zip'))
        for f in files:
            self.__logger.debug(f"Deleting file \"{f}\"")
            os.remove(f)
        os.removedirs(self.__settings.get_mod_folder())
        self.__logger.debug(f"Deleting configuration file \"{self.__settings.file_path}\"")
        os.remove(self.__settings.file_path)
        self.__app.exit(0)