import shutil from io import BytesIO from pathlib import Path from typing import Generator, IO from source.mkw.ModConfig import ModConfig from source.mkw.Patch.Patch import Patch from source.wt import szs, lec, wit from source.wt.wstrt import StrPath Game: any class PathOutsideMod(Exception): def __init__(self, forbidden_path: Path, allowed_range: Path): super().__init__(f"Error : path {forbidden_path} outside of allowed range {allowed_range}") class ExtractedGame: """ Class that represents an extracted game """ def __init__(self, path: "Path | str", original_game: "Game" = None): self.path = Path(path) self.original_game = original_game self._special_file: dict[str, IO] = {} def extract_autoadd(self, destination_path: "Path | str") -> Generator[dict, None, None]: """ Extract all the autoadd files from the game to destination_path :param destination_path: directory where the autoadd files will be extracted """ yield {"description": "Extracting autoadd files...", "determinate": False} szs.autoadd(self.path / "files/Race/Course/", destination_path) def extract_original_tracks(self, destination_path: "Path | str") -> Generator[dict, None, None]: """ Move all the original tracks to the destination path :param destination_path: destination of the track """ destination_path = Path(destination_path) destination_path.mkdir(parents=True, exist_ok=True) yield {"description": "Extracting original tracks...", "determinate": False} for track_file in (self.path / "files/Race/Course/").glob("*.szs"): yield {"description": f"Extracting original tracks ({track_file.name})...", "determinate": False} if not (destination_path / track_file.name).exists(): track_file.rename(destination_path / track_file.name) else: track_file.unlink() def install_mystuff(self, mystuff_path: "Path | str") -> Generator[dict, None, None]: """ Install mystuff directory. If any files of the game have the same name as a file at the root of the MyStuff Patch, then it is copied. :mystuff_path: path to the MyStuff directory :return: """ yield {"description": f"Installing MyStuff '{mystuff_path}'...", "determinate": False} mystuff_path = Path(mystuff_path) mystuff_rootfiles: dict[str, Path] = {} for mystuff_subpath in mystuff_path.glob("*"): if mystuff_subpath.is_file(): mystuff_rootfiles[mystuff_subpath.name] = mystuff_subpath else: shutil.copytree(mystuff_subpath, self.path / f"files/{mystuff_subpath.name}", dirs_exist_ok=True) for game_file in filter(lambda file: file.is_file(), (self.path / "files/").rglob("*")): if (mystuff_file := mystuff_rootfiles.get(game_file.name)) is None: continue shutil.copy(mystuff_file, game_file) def install_multiple_mystuff(self, mystuff_paths: list["Path | str"]) -> Generator[dict, None, None]: """ Install multiple mystuff patch :param mystuff_paths: paths to all the mystuff patch """ yield {"description": "Installing all the mystuff patchs"} for mystuff_path in mystuff_paths: yield from self.install_mystuff(mystuff_path) def prepare_special_file(self, mod_config: ModConfig) -> Generator[dict, None, None]: """ Prepare special files for the patch :return: the special files dict """ yield {"description": "Preparing ct_icon special file...", "determinate": False} ct_icons = BytesIO() mod_config.get_full_cticon().save(ct_icons, format="PNG") ct_icons.seek(0) self._special_file["ct_icons"] = ct_icons def prepare_dol(self) -> Generator[dict, None, None]: """ Prepare main.dol and StaticR.rel files (clean them and add lecode) """ yield {"description": "Preparing main.dol...", "determinate": False} StrPath(self.path / "sys/main.dol").patch(clean_dol=True, add_lecode=True) def recreate_all_szs(self) -> Generator[dict, None, None]: """ Repack all the .d directory into .szs files. """ yield {"description": f"Repacking all szs", "determinate": False} for extracted_szs in filter(lambda path: path.is_dir(), self.path.rglob("*.d")): # for every directory that end with a .d in the extracted game, recreate the szs yield {"description": f"Repacking {extracted_szs} to szs", "determinate": False} szs.create(extracted_szs, extracted_szs.with_suffix(".szs"), overwrite=True) shutil.rmtree(str(extracted_szs.resolve())) def patch_lecode(self, mod_config: ModConfig, cache_directory: Path | str, cttracks_directory: Path | str, ogtracks_directory: Path | str) -> Generator[dict, None, None]: """ install lecode on the mod :param cttracks_directory: directory to the customs tracks :param ogtracks_directory: directory to the originals tracks :param cache_directory: Path to the cache :param mod_config: mod configuration """ yield {"description": "Patching LECODE.bin"} cache_directory = Path(cache_directory) cttracks_directory = Path(cttracks_directory) ogtracks_directory = Path(ogtracks_directory) # write ctfile data to a file in the cache ct_file = (cache_directory / "ctfile.lpar") ct_file.write_text(mod_config.get_ctfile(template="-")) lpar_dir: Path = mod_config.path.parent / "_LPAR/" lpar: Path = lpar_dir / mod_config.multiple_safe_eval(mod_config.lpar_template) if not lpar.is_relative_to(lpar_dir): raise PathOutsideMod(lpar, lpar_dir) for lecode_file in (self.path / "files/rel/").glob("lecode-*.bin"): lec.patch( lecode_file=lecode_file, ct_file=ct_file, lpar=lpar, game_tracks_directory=self.path / "files/Race/Course/", copy_tracks_directories=[ogtracks_directory, cttracks_directory] ) def install_all_patch(self, mod_config: ModConfig) -> Generator[dict, None, None]: """ Install all patchs of the mod_config into the game :param mod_config: the mod to install """ yield {"description": "Installing all Patch...", "determinate": False} # prepare special files data yield from self.prepare_special_file(mod_config) # for all directory that are in the root of the mod, and don't start with an underscore, # for all the subdirectory named "_PATCH", apply the patch for part_directory in mod_config.get_mod_directory().glob("[!_]*"): for patch_directory in part_directory.glob("_PATCH/"): yield from Patch(patch_directory, mod_config, self._special_file).install(self) def convert_to(self, output_type: wit.Extension) -> Generator[dict, None, None]: """ Convert the extracted game to another format :param output_type: path to the destination of the game :output_type: format of the destination game """ yield {"description": f"Converting game to {output_type}", "determinate": False} if output_type == wit.Extension.FST: return destination_file = self.path.with_suffix(self.path.suffix + output_type.value) dest_name: str = destination_file.name i: int = 0 while destination_file.exists(): i += 1 destination_file = destination_file.with_name(dest_name + f" ({i})") wit.copy( source_directory=self.path, destination_file=destination_file, ) yield {"description": "Deleting the extracted game...", "determinate": False} shutil.rmtree(self.path)