diff --git a/source/mkw/Game.py b/source/mkw/Game.py index ed617e4..5820c51 100644 --- a/source/mkw/Game.py +++ b/source/mkw/Game.py @@ -1,73 +1,29 @@ -import enum from pathlib import Path -from source.wt.wit import WITPath - - -class Extension(enum.Enum): - """ - Enum for game extension - """ - FST = ".dol" - WBFS = ".wbfs" - ISO = ".iso" - - @classmethod - def _missing_(cls, value: str) -> "Extension | None": - """ - if not found, search for the same value with lower case - :param value: value to search for - :return: None if nothing found, otherwise the found value - """ - value = value.lower() - for member in filter(lambda m: m.value == value, cls): return member - return None - - -class Region(enum.Enum): - """ - Enum for game region - """ - PAL = "PAL" - USA = "USA" - EUR = "EUR" - KOR = "KOR" +from source.wt.wit import WITPath, Region class Game: def __init__(self, path: Path | str): - self.path = Path(path) if isinstance(path, str) else path - - @property - def extension(self) -> Extension: - """ - Returns the extension of the game - :return: the extension of the game - """ - return Extension(self.path.suffix) - - @property - def id(self) -> str: - """ - Return the id of the game (RMCP01, RMCK01, ...) - :return: the id of the game - """ - return WITPath(self.path).analyze()["id6"] - - @property - def region(self) -> Region: - """ - Return the region of the game (PAL, USA, EUR, ...) - :return: the region of the game - """ - return Region(WITPath(self.path).analyze()["dol_region"]) + self.wit_path = WITPath(path) def is_mkw(self) -> bool: """ - Return True if the game is Mario Kart Wii, else otherwise + Return True if the game is Mario Kart Wii, False otherwise :return: is the game a MKW game """ - return WITPath(self.path).analyze()["dol_is_mkw"] == 1 + return self.wit_path.analyze()["dol_is_mkw"] == 1 def is_vanilla(self) -> bool: - ... + """ + Return True if the game is vanilla, False if the game is modded + :return: if the game is not modded + """ + return not any(self.wit_path[f"./files/rel/lecode-{region.value}.bin"].exists() for region in Region) + + def extract(self, dest: Path | str) -> Path: + """ + Extract the game to the destination directory. If the game is a FST, just copy to the destination + :param dest: destination directory + """ + return self.wit_path.extract_all(dest) diff --git a/source/wt/szs.py b/source/wt/szs.py index 673fe9b..4ebf66b 100644 --- a/source/wt/szs.py +++ b/source/wt/szs.py @@ -5,14 +5,18 @@ tools_path = tools_szs_dir / ("wszst.exe" if system == "win64" else "wszst") class SZSPath: - __slots__ = ("path",) + __slots__ = ("path", "_analyze") def __init__(self, path: Path | str): self.path: Path = path if isinstance(path, Path) else Path(path) + self._analyze = None - def __repr__(self): + def __repr__(self) -> str: return f"" + def __eq__(self, other: "SZSPath") -> bool: + return self.path == other.path + @better_error(tools_path) def _run(self, *args) -> bytes: """ @@ -62,10 +66,11 @@ class SZSPath: def analyze(self) -> dict: """ - Return the analyze of the szs + Return the analyze of the file :return: dictionnary of key and value of the analyze """ - return self._run_dict("ANALYZE", self.path) + if self._analyze is None: self._analyze = self._run_dict("ANALYZE", self.path) + return self._analyze def list_raw(self) -> list[str]: """ @@ -75,7 +80,11 @@ class SZSPath: # cycle though all of the output line of the command, check if the line are empty, and if not, # add it to the list. Finally, remove the first line because this is a description of the command - return [subfile.strip() for subfile in self._run("list", self.path).decode().splitlines() if subfile][1:] + return [ + subfile.strip() + for subfile in self._run("list", self.path).decode().splitlines() + if subfile.startswith("./") + ] def list(self) -> list["SZSSubPath"]: """ @@ -105,9 +114,12 @@ class SZSSubPath: self.szs_path = szs_path self.subfile = subfile - def __repr__(self): + def __repr__(self) -> str: return f"" + def __eq__(self, other: "SZSSubPath") -> bool: + return self.subfile == other.subfile and self.szs_path == other.szs_path + def extract(self, dest: Path | str) -> Path: """ Extract the subfile to a destination @@ -124,11 +136,30 @@ class SZSSubPath: return dest - def is_dir(self): + def exists(self): + """ + Return if the subfile exist in the szs + :return: True if the subfile exist, else otherwise + """ + return self in self.szs_path.list() + + def is_dir(self) -> bool: + """ + Return if the subfile is a directory + :return: True if the subfile is a directory, else otherwise + """ return self.subfile.endswith("/") - def is_file(self): + def is_file(self) -> bool: + """ + Return if the subfile is a file + :return: True if the subfile is a file, else otherwise + """ return not self.is_dir() - def basename(self): + def basename(self) -> str: + """ + Return the basename of the subfile + :return: the basename of the subfile + """ return self.subfile.rsplit("/", 1)[-1] diff --git a/source/wt/wit.py b/source/wt/wit.py index feb527f..d3f4c80 100644 --- a/source/wt/wit.py +++ b/source/wt/wit.py @@ -1,14 +1,54 @@ +import enum +import shutil + from source.wt import * from source.wt import _run, _run_dict tools_path = tools_wit_dir / ("wit.exe" if system == "win64" else "wit") -class WITPath: - __slots__ = ("path",) +class Extension(enum.Enum): + """ + Enum for game extension + """ + FST = ".dol" + WBFS = ".wbfs" + ISO = ".iso" - def __init__(self, path: Path): - self.path = path + @classmethod + def _missing_(cls, value: str) -> "Extension | None": + """ + if not found, search for the same value with lower case + :param value: value to search for + :return: None if nothing found, otherwise the found value + """ + value = value.lower() + for member in filter(lambda m: m.value == value, cls): return member + return None + + +class Region(enum.Enum): + """ + Enum for game region + """ + PAL = "PAL" + USA = "USA" + EUR = "EUR" + KOR = "KOR" + + +class WITPath: + __slots__ = ("path", "_analyze") + + def __init__(self, path: Path | str): + self.path: Path = path if isinstance(path, Path) else Path(path) + self._analyze = None + + def __repr__(self) -> str: + return f"" + + def __eq__(self, other: "WITPath") -> bool: + return self.path == other.path @better_error(tools_path) def _run(self, *args) -> bytes: @@ -28,9 +68,172 @@ class WITPath: """ return _run_dict(tools_path, *args) + def _get_fst_root(self) -> Path: + """ + If the game is a FST, return the root of the FST + :return: root of the FST + """ + # main.dol is located in ./sys/main.dol, so return parent of parent + if self.extension == Extension.FST: return self.path.parent.parent + def analyze(self) -> dict: """ Return the analyze of the file :return: dictionnary of key and value of the analyze """ - return self._run_dict("ANALYZE", self.path) + if self._analyze is None: self._analyze = self._run_dict("ANALYZE", self.path) + return self._analyze + + def list_raw(self) -> list[str]: + """ + Return the list of subfiles + :return: the list of subfiles + """ + if self.extension == Extension.FST: + return [ + str(file.relative_to(self._get_fst_root())) + for file in self._get_fst_root().rglob("*") + ] + + return [ + subfile.strip() for subfile + in self._run("files", self.path).decode().splitlines() + if subfile.startswith("./") + ] + + def list(self) -> list["WITSubPath"]: + """ + Return the list of subfiles + :return: the list of subfiles + """ + return [self.get_subfile(subfile) for subfile in self.list_raw()] + + def get_subfile(self, subfile: str) -> "WITSubPath": + """ + Return the subfile of the game + :return: the subfile + """ + return WITSubPath(self, subfile) + + def __getitem__(self, item): + return self.get_subfile(item) + + def __iter__(self): + return iter(self.list()) + + def extract_all(self, dest: Path | str) -> Path: + """ + Extract all the subfiles to the destination directory + :param dest: destination directory + :return: the extracted file path + """ + return self["./"].extract(dest, flat=False) + + @property + def extension(self) -> Extension: + """ + Returns the extension of the game + :return: the extension of the game + """ + return Extension(self.path.suffix) + + @property + def id(self) -> str: + """ + Return the id of the game (RMCP01, RMCK01, ...) + :return: the id of the game + """ + return self.analyze()["id6"] + + @property + def region(self) -> Region: + """ + Return the region of the game (PAL, USA, EUR, ...) + :return: the region of the game + """ + return Region(self.analyze()["dol_region"]) + + +class WITSubPath: + __slots__ = ("wit_path", "subfile") + + def __init__(self, wit_path: WITPath, subfile: str): + self.wit_path = wit_path + self.subfile = subfile.removeprefix("./").replace("\\", "/") + + def __repr__(self): + if self.wit_path.extension == Extension.FST: return f"" + return f"" + + def __eq__(self, other: "WITSubPath") -> bool: + return self.subfile == other.subfile and self.wit_path == other.wit_path + + def _get_fst_path(self) -> Path: + """ + Return the path of the subfile in the FST + :return: the path of the subfile in the FST + """ + return self.wit_path._get_fst_root() / self.subfile + + def extract(self, dest: Path | str, flat: bool = True) -> Path: + """ + Extract the subfile to the destination directory + :param flat: all files will be extracted directly in the directory, instead of creating subdirectory + :param dest: destination directory + :return: the extracted file path + """ + dest: Path = dest if isinstance(dest, Path) else Path(dest) + + if self.wit_path.extension == Extension.FST: + # if flat is used, extract the file / dir into the destination directory, without subdirectory + if flat: + os.makedirs(dest, exist_ok=True) + # if we are extracting a directory, we need to extract every file recursively + if self.is_dir(): + for file in (self._get_fst_path()).rglob("*"): + if file.is_file(): shutil.copy(file, dest / file.name) + # else we just copy the file + else: + shutil.copy(self._get_fst_path(), dest) + # if flat is not used, copy the structure of the directory, or just copy the file + else: + func = shutil.copytree if self.is_dir() else shutil.copy + func(self._get_fst_path(), dest / self.subfile) + + return dest / self.basename() + + else: + args = [] + if flat: args.append("--flat") + self.wit_path._run("EXTRACT", self.wit_path.path, f"--files=+{self.subfile}", "-d", dest, *args) + return dest / self.basename() + + def is_dir(self) -> bool: + """ + Return if the subfile is a directory + :return: True if the subfile is a directory, else otherwise + """ + if self.wit_path.extension == Extension.FST: + return self._get_fst_path().is_dir() + return self.subfile.endswith("/") + + def is_file(self) -> bool: + """ + Return if the subfile is a file + :return: True if the subfile is a file, else otherwise + """ + return not self.is_dir() + + def exists(self): + """ + Return if the subfile exist in the game + :return: True if the subfile exist, else otherwise + """ + return self in self.wit_path.list() + + def basename(self) -> str: + """ + Return the basename of the subfile + :return: the basename of the subfile + """ + return self.subfile.rsplit("/", 1)[-1]