Atlas-Install/source/mkw/Patch/PatchFile.py

from io import BytesIO
from pathlib import Path
from typing import Generator, IO, TYPE_CHECKING

from source.mkw.Patch import PathOutsidePatch, InvalidPatchMode, InvalidSourceMode
from source.mkw.Patch.PatchOperation import AbstractPatchOperation
from source.mkw.Patch.PatchObject import PatchObject
from source.progress import Progress
from source.wt.szs import SZSPath

if TYPE_CHECKING:
    from source.mkw.ExtractedGame import ExtractedGame


class PatchFile(PatchObject):
    """
    Represent a file from a patch
    """

    def get_source_path(self, game_subpath: Path) -> Path:
        """
        Return the path of the file that the patch is applied on.
        If the configuration mode is set to "game", then return the path of the file inside the Game
        Else, return the path of the file inside the Patch
        """
        match self.configuration["source"]:
            case "patch": return self.full_path
            case "game": return game_subpath
            case _: raise InvalidSourceMode(self, self.configuration["course"])

    def get_patched_file(self, game_subpath: Path) -> (str, BytesIO):
        """
        Return the name and the content of the patched file
        :param game_subpath: path to the game subfile
        :return: the name and the content of the patched file
        """
        patch_source: Path = self.get_source_path(game_subpath)
        patch_name: str = game_subpath.name
        patch_content: BytesIO = BytesIO(open(patch_source, "rb").read())
        # the file is converted into a BytesIO because if the source is "game",
        # the file is overwritten and if the content is untouched, the file will only be lost

        for operation_name, operation in self.configuration.get("operation", {}).items():
            # process every operation and get the new patch_path (if the name is changed)
            # and the new content of the patch
            patch_name, patch_content = AbstractPatchOperation.get(operation_name)(**operation).patch(
                self.patch, patch_name, patch_content
            )

        return patch_name, patch_content

    @staticmethod
    def write_patch(destination: Path, patch_content: IO) -> None:
        """
        Write a patch content to the destination. Automatically create the directory and seek to the start.
        :param destination: file where the content will be written
        :param patch_content: content of the file to write
        """
        destination.parent.mkdir(parents=True, exist_ok=True)
        with open(destination, "wb") as file:
            patch_content.seek(0)
            file.write(patch_content.read())

    @staticmethod
    def check_szs(extracted_game: "ExtractedGame", game_subpath: Path) -> None:
        """
        Check if game path is inside a szs archive. If yes, extract it.
        :param extracted_game: the extracted game object
        :param game_subpath: path to the game file that is being patched
        """
        for szs_subpath in filter(lambda path: path.suffix == ".d",
                                  game_subpath.relative_to(extracted_game.path).parents):

            szs_path = extracted_game.path / szs_subpath

            # if the archive is not already extracted  and the szs file in the game exists, extract it
            if not szs_path.exists() and szs_path.with_suffix(".szs").exists():
                SZSPath(szs_path.with_suffix(".szs")).extract_all(szs_path)

    def install(self, extracted_game: "ExtractedGame", game_subpath: Path) -> Generator[Progress, None, None]:
        """
        patch a subfile of the game with the PatchFile
        """
        yield Progress(description=f"Patching {game_subpath}")

        # check if the file should be patched
        if not self.is_enabled(extracted_game): return

        # check if it is patching a szs archive. If yes, extract it.
        self.check_szs(extracted_game, game_subpath)

        # apply operation on the file

        match self.configuration["mode"]:
            # if the mode is copy, replace the subfile in the game by the PatchFile
            case "copy":
                patch_name, patch_content = self.get_patched_file(game_subpath)
                self.write_patch(game_subpath.parent / patch_name, patch_content)
                patch_content.close()

            # if the mode is replace, only write if the file existed before
            case "replace":
                patch_name, patch_content = self.get_patched_file(game_subpath)
                if (game_subpath.parent / patch_name).exists():
                    self.write_patch(game_subpath.parent / patch_name, patch_content)
                    patch_content.close()

            # if the mode is match, replace all the subfiles that match match_regex by the PatchFile
            case "match":
                patch_content: BytesIO = BytesIO()

                # if the source is the patch, then directly calculate the patch content
                # so that it is not recalculated for every file with no reason
                if self.configuration["source"] == "patch":
                    _, patch_content = self.get_patched_file(game_subpath)

                for game_subfile in game_subpath.parent.glob(self.configuration["match_regex"]):
                    # disallow patching files outside of the game
                    if not game_subfile.relative_to(extracted_game.path):
                        raise PathOutsidePatch(game_subfile, extracted_game.path)

                    yield Progress(description=f"Patching {game_subfile}")

                    # if the source is the game, then recalculate the content for every game subfile
                    if self.configuration["source"] == "game":
                        _, patch_content = self.get_patched_file(game_subfile)

                    # patch the game with the patch content
                    self.write_patch(game_subfile, patch_content)

                patch_content.close()

            # ignore if mode is "ignore", useful if the file is used as a resource for an operation
            case "ignore": pass
            # else raise an error
            case _: raise InvalidPatchMode(self, self.configuration["mode"])