import hashlib import shutil from io import BytesIO from pathlib import Path from typing import Generator, IO, TYPE_CHECKING from source import file_block_size from source.mkw import PathOutsideAllowedRange from source.mkw.ModConfig import ModConfig from source.mkw.Patch.Patch import Patch from source.mkw.collection.Extension import Extension from source.progress import Progress from source.utils import comp_dict_changes from source.wt import szs, lec, wit from source.wt.wstrt import StrPath from source.translation import translate as _ if TYPE_CHECKING: from source.mkw.Game import Game RIIVOLUTION_FOLDER_NAME: str = "riivolution" 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[Progress, 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 Progress(description=_("TEXT_EXTRACTING_AUTOADD"), determinate=False) szs.autoadd(self.path / "files/Race/Course/", destination_path) def extract_original_tracks(self, destination_path: "Path | str") -> Generator[Progress, 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) original_tracks: list[Path] = list((self.path / "files/Race/Course/").glob("*.szs")) yield Progress( determinate=True, max_step=len(original_tracks), set_step=0, ) for track_file in original_tracks: yield Progress(description=_("TEXT_EXTRACTING_ORIGINAL_TRACKS") % track_file.name, step=1) 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[Progress, 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 Progress(description=_("TEXT_INSTALLING_MYSTUFF") % mystuff_path) 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[Progress, None, None]: """ Install multiple mystuff patch :param mystuff_paths: paths to all the mystuff patch """ yield Progress(determinate=True, max_step=len(mystuff_paths), set_step=0) for mystuff_path in mystuff_paths: yield Progress(step=1) yield from self.install_mystuff(mystuff_path) def prepare_special_file(self, mod_config: ModConfig) -> Generator[Progress, None, None]: """ Prepare special files for the patch :return: the special files dict """ yield Progress(description=_("TEXT_PREPARING_SPECIAL_FILE") % "ct_icons", 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[Progress, None, None]: """ Prepare main.dol and StaticR.rel files (clean them and add lecode) """ yield Progress(description=_("TEXT_PREPARING_MAIN_DOL"), determinate=False) StrPath(self.path / "sys/main.dol").patch(clean_dol=True, add_lecode=True) def recreate_all_szs(self) -> Generator[Progress, None, None]: """ Repack all the .d directory into .szs files. """ all_extracted_szs: list[Path] = list(filter(lambda path: path.is_dir(), self.path.rglob("*.d"))) yield Progress( determinate=True, max_step=len(all_extracted_szs), set_step=0, ) for extracted_szs in all_extracted_szs: # for every directory that end with a .d in the extracted game, recreate the szs yield Progress( description=_("TEXT_REPACKING_ARCHIVE") % extracted_szs.relative_to(self.path), step=1 ) 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[Progress, 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 Progress(description=_("TEXT_PATCHING_LECODE")) 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 PathOutsideAllowedRange(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, patch_directory_name: str) -> Generator[Progress, None, None]: """ for all directory that are in the root of the mod, and don't start with an underscore, for all the subdirectory named by the patch_directory_name, apply the patch :param mod_config: the mod to install """ patch_directories: list[Path] = [] for part_directory in mod_config.get_mod_directory().glob("[!_]*"): for patch_directory in part_directory.glob(patch_directory_name): patch_directories.append(patch_directory) yield Progress(determinate=True, max_step=len(patch_directories)+1, set_step=0) for patch_directory in patch_directories: yield Progress(step=1) yield from Patch(patch_directory, mod_config, self._special_file).install(self) def install_all_prepatch(self, mod_config: ModConfig) -> Generator[Progress, None, None]: """ Install all patchs of the mod_config into the game. Used before the lecode patch is applied :param mod_config: the mod to install """ yield from self._install_all_patch(mod_config, "_PREPATCH/") def install_all_patch(self, mod_config: ModConfig) -> Generator[Progress, None, None]: """ Install all patchs of the mod_config into the game. Used after the lecode patch is applied :param mod_config: the mod to install """ yield from self._install_all_patch(mod_config, "_PATCH/") def convert_to(self, output_type: Extension) -> Generator[Progress, None, wit.WITPath | 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 Progress(description=_("TEXT_CONVERT_GAME_TO") % output_type.name, determinate=False) if output_type == Extension.FST: return destination_file = self.path.with_suffix(self.path.suffix + output_type.value) dest_stem: str = destination_file.stem i: int = 0 while destination_file.exists(): i += 1 destination_file = destination_file.with_stem(dest_stem + f" ({i})") converted_game: wit.WITPath = wit.copy( source_directory=self.path, destination_file=destination_file, ) yield Progress(description=_("TEXT_DELETING_EXTRACTED_GAME"), determinate=False) shutil.rmtree(self.path) return converted_game def convert_to_riivolution(self, mod_config: ModConfig, old_hash_map: dict[str, str] ) -> Generator[Progress, None, None]: """ Convert the extracted game into a riivolution patch :param mod_config: the mod configuration :param old_hash_map: hash map of all the games files created before all the modifications """ new_hash_map = yield from self.get_hash_map() game_files: list[Path] = list(filter(lambda file: file.is_file(), self.path.rglob("*"))) yield Progress( description=_("TEXT_CONVERTING_TO_RIIVOLUTION"), determinate=True, max_step=len(game_files), set_step=0, ) # get the files difference between the original game and the patched game diff_hash_map: dict[str, Path] = comp_dict_changes(old_hash_map, new_hash_map) for file in game_files: yield Progress(step=1) # if the file have not being patched, delete it if str(file.relative_to(self.path)) not in diff_hash_map: file.unlink() # get riivolution configuration content riivolution_config_content = "\n".join(( '', ' ', '', ' ', f'
', ' ', ' ', ' ', '
', '
', '', ' ', f' ', f' ', ' ', '', ' ', f' ', f' ', ' ', '', ' ', ' ', ' ', '
', )) # get riivolution configuration path riivolution_config_path: Path = self.path.parent / f"{RIIVOLUTION_FOLDER_NAME}/{str(mod_config)}.xml" riivolution_config_path.parent.mkdir(parents=True, exist_ok=True) riivolution_config_path.write_text(riivolution_config_content, encoding="utf-8") def get_hash_map(self) -> Generator[Progress, None, dict[str, str]]: """ Return a dictionary associating all the game subfiles to a hash :return: a dictionary associating all the game subfiles to a hash """ md5_map: dict[str, str] = {} game_files: list[Path] = list(filter(lambda fp: fp.is_file(), self.path.rglob("*"))) yield Progress( determinate=True, max_step=len(game_files), set_step=0 ) for fp in game_files: hasher = hashlib.md5() rel_path: str = str(fp.relative_to(self.path)) yield Progress( description=_("TEXT_CALCULATING_HASH") % rel_path, step=1 ) with open(fp, "rb") as file: while block := file.read(file_block_size): hasher.update(block) md5_map[rel_path] = hasher.hexdigest() return md5_map