diff --git a/GeckoLoader.py b/GeckoLoader.py index f233c61..ff3b40f 100644 --- a/GeckoLoader.py +++ b/GeckoLoader.py @@ -5,12 +5,12 @@ import os import random import shutil import sys +import atexit from distutils.version import LooseVersion from dolreader import DolFile from kernel import CodeHandler, KernelLoader from tools import CommandLineParser, color_text -from fileutils import GC_File from versioncheck import Updater try: @@ -34,13 +34,15 @@ except ImportError: TRED = '' TREDLIT = '' -__version__ = 'v6.1.0' +__version__ = 'v6.1.5' def resource_path(relative_path: str): """ Get absolute path to resource, works for dev and for PyInstaller """ base_path = os.path.dirname(os.path.realpath(sys.argv[0])) return os.path.join(base_path, relative_path) +TMPDIR = resource_path(''.join(random.choice('1234567890-_abcdefghijklomnpqrstuvwxyz') for i in range(6)) + '-GeckoLoader') + def sort_file_args(fileA, fileB): if os.path.splitext(fileA)[1].lower() == '.dol': dolFile = fileA @@ -52,80 +54,90 @@ def sort_file_args(fileA, fileB): parser.error(color_text('No dol file was passed\n', defaultColor=TREDLIT)) return dolFile, gctFile -if __name__ == "__main__": - parser = CommandLineParser(prog='GeckoLoader ' + __version__, - description='Process files and allocations for GeckoLoader', - allow_abbrev=False) +@atexit.register +def clean_tmp_resources(): + if os.path.isdir(TMPDIR): + shutil.rmtree(TMPDIR) - parser.add_argument('dolfile', help='DOL file') - parser.add_argument('codelist', help='Folder or Gecko GCT|TXT file') - parser.add_argument('-a', '--alloc', - help='Define the size of the code allocation in hex, only applies when using the ARENA space', - metavar ='SIZE') - parser.add_argument('-i', '--init', - help='Define where GeckoLoader is initialized in hex', - metavar='ADDRESS') - parser.add_argument('-m', '--movecodes', - help='''["AUTO", "LEGACY", "ARENA"] Choose if GeckoLoader moves the codes to OSArenaHi, - or the legacy space. Default is "AUTO", - which auto decides where to insert the codes''', - default='AUTO', - choices=['AUTO', 'LEGACY', 'ARENA'], - metavar='TYPE') - parser.add_argument('-tc', '--txtcodes', - help='''["ACTIVE", "ALL"] What codes get parsed when a txt file is used. - "ALL" makes all codes get parsed, - "ACTIVE" makes only activated codes get parsed. - "ACTIVE" is the default''', - default='ACTIVE', - metavar='TYPE') - parser.add_argument('--handler', - help='''["MINI", "FULL"] Which codehandler gets used. "MINI" uses a smaller codehandler - which only supports (0x, 2x, Cx, and E0 types) and supports up to - 600 lines of gecko codes when using the legacy codespace. - "FULL" is the standard codehandler, supporting up to 350 lines of code - in the legacy codespace. "FULL" is the default''', - default='FULL', - choices=['MINI', 'FULL'], - metavar='TYPE') - parser.add_argument('--hooktype', +class GeckoLoaderCli(CommandLineParser): + + def __init__(self, name, version=None, description=''): + super().__init__(prog=(name+' '+version), description=description, allow_abbrev=False) + self.__version__ = version + self.__doc__ = description + + self.add_argument('dolfile', help='DOL file') + self.add_argument('codelist', help='Folder or Gecko GCT|TXT file') + self.add_argument('-a', '--alloc', + help='Define the size of the code allocation in hex, only applies when using the ARENA space', + metavar ='SIZE') + self.add_argument('-i', '--init', + help='Define where GeckoLoader is initialized in hex', + metavar='ADDRESS') + self.add_argument('-m', '--movecodes', + help='''["AUTO", "LEGACY", "ARENA"] Choose if GeckoLoader moves the codes to OSArenaHi, + or the legacy space. Default is "AUTO", + which auto decides where to insert the codes''', + default='AUTO', + choices=['AUTO', 'LEGACY', 'ARENA'], + metavar='TYPE') + self.add_argument('-tc', '--txtcodes', + help='''["ACTIVE", "ALL"] What codes get parsed when a txt file is used. + "ALL" makes all codes get parsed, + "ACTIVE" makes only activated codes get parsed. + "ACTIVE" is the default''', + default='ACTIVE', + metavar='TYPE') + self.add_argument('--handler', + help='''["MINI", "FULL"] Which codehandler gets used. "MINI" uses a smaller codehandler + which only supports (0x, 2x, Cx, and E0 types) and supports up to + 600 lines of gecko codes when using the legacy codespace. + "FULL" is the standard codehandler, supporting up to 350 lines of code + in the legacy codespace. "FULL" is the default''', + default='FULL', + choices=['MINI', 'FULL'], + metavar='TYPE') + self.add_argument('--hooktype', help='''["VI", "GX", "PAD"] The type of hook used for the RAM search. "VI" or "GX" are recommended, although "PAD" can work just as well. "VI" is the default''', default='VI', choices=['VI', 'GX', 'PAD'], metavar='HOOK') - parser.add_argument('--hookaddress', + self.add_argument('--hookaddress', help='Choose where the codehandler hooks to in hex, overrides auto hooks', metavar='ADDRESS') - parser.add_argument('-o', '--optimize', + self.add_argument('-o', '--optimize', help='''Optimizes the codelist by directly patching qualifying ram writes into the dol file, and removing them from the codelist''', action='store_true') - parser.add_argument('-p', '--protect', + self.add_argument('-p', '--protect', help='''Targets and nullifies the standard codehandler provided by loaders and Dolphin Emulator, only applies when the ARENA is used''', action='store_true') - parser.add_argument('--dest', + self.add_argument('--dest', help='Target path to put the modified DOL, can be a folder or file', metavar='PATH') - parser.add_argument('--check-update', + self.add_argument('--check-update', help='''Checks to see if a new update exists on the GitHub Repository releases page, this option overrides all other commands.''', action='store_true') - parser.add_argument('--encrypt', + self.add_argument('--encrypt', help='Encrypts the codelist on compile time, helping to slow the snoopers', action='store_true') - parser.add_argument('-q', '--quiet', + self.add_argument('-q', '--quiet', help='Print nothing to the console', action='store_true') - parser.add_argument('-v', '--verbose', + self.add_argument('-v', '--verbose', help='Print extra info to the console', default=0, action='count') + + def __str__(self): + return self.__doc__ - if len(sys.argv) == 1: - version = __version__.rjust(9, ' ') + def print_splash(self): helpMessage = 'Try option -h for more info on this program'.center(64, ' ') + version = self.__version__.rjust(9, ' ') logo = [' ', ' ╔═══════════════════════════════════════════════════════════╗ ', @@ -151,10 +163,13 @@ if __name__ == "__main__": ' ', f'{helpMessage}', ' '] + for line in logo: print(color_text(line, [('║', TREDLIT), ('╔╚╝╗═', TRED)], TGREENLIT)) + sys.exit(0) - elif '--check-update' in sys.argv: + + def check_updates(self): repoChecker = Updater('JoshuaMKW', 'GeckoLoader') tag, status = repoChecker.get_newest_version() @@ -162,21 +177,29 @@ if __name__ == "__main__": print('') if status is False: - parser.error(color_text(tag + '\n', defaultColor=TREDLIT), print_usage=False) + self.error(color_text(tag + '\n', defaultColor=TREDLIT), print_usage=False) - if LooseVersion(tag) > LooseVersion(__version__): + if LooseVersion(tag) > LooseVersion(self.__version__): print(color_text(f' :: A new update is live at {repoChecker.gitReleases.format(repoChecker.owner, repoChecker.repo)}', defaultColor=TYELLOWLIT)) - print(color_text(f' :: Current version is "{__version__}", Most recent version is "{tag}"', defaultColor=TYELLOWLIT)) - elif LooseVersion(tag) < LooseVersion(__version__): + print(color_text(f' :: Current version is "{self.__version__}", Most recent version is "{tag}"', defaultColor=TYELLOWLIT)) + elif LooseVersion(tag) < LooseVersion(self.__version__): print(color_text(' :: No update available', defaultColor=TGREENLIT)) - print(color_text(f' :: Current version is "{__version__}(dev)", Most recent version is "{tag}(release)"', defaultColor=TGREENLIT)) + print(color_text(f' :: Current version is "{self.__version__}(dev)", Most recent version is "{tag}(release)"', defaultColor=TGREENLIT)) else: print(color_text(' :: No update available', defaultColor=TGREENLIT)) - print(color_text(f' :: Current version is "{__version__}(release)", Most recent version is "{tag}(release)"', defaultColor=TGREENLIT)) + print(color_text(f' :: Current version is "{self.__version__}(release)", Most recent version is "{tag}(release)"', defaultColor=TGREENLIT)) print('') sys.exit(0) +if __name__ == "__main__": + parser = GeckoLoaderCli('GeckoLoader ' + __version__, __version__, description='Dol editing tool for allocating extended codespace') + + if len(sys.argv) == 1: + parser.print_splash() + elif '--check-update' in sys.argv: + parser.check_updates() + args = parser.parse_args() if args.alloc: @@ -198,11 +221,8 @@ if __name__ == "__main__": else: _codehook = None - if args.handler: - if args.handler == 'MINI': - codeHandlerFile = 'codehandler-mini.bin' - else: - codeHandlerFile = 'codehandler.bin' + if args.handler == 'MINI': + codeHandlerFile = 'codehandler-mini.bin' else: codeHandlerFile = 'codehandler.bin' @@ -213,14 +233,16 @@ if __name__ == "__main__": if not os.path.exists(args.codelist): parser.error(color_text(f'File/folder "{args.codelist}" does not exist\n', defaultColor=TREDLIT)) - tmpdir = ''.join(random.choice('1234567890-_abcdefghijklomnpqrstuvwxyz') for i in range(6)) + '-GeckoLoader' + with open(os.path.normpath(args.dolfile), 'rb') as dol: + dolFile = DolFile(dol) with open(resource_path(os.path.join('bin', os.path.normpath(codeHandlerFile))), 'rb') as handler: codeHandler = CodeHandler(handler) codeHandler.allocation = _allocation codeHandler.hookAddress = _codehook codeHandler.hookType = args.hooktype - codeHandler.includeAll = (args.txtcodes.lower() == 'all') + codeHandler.includeAll = args.txtcodes.lower() == 'all' + codeHandler.optimizeList = args.optimize with open(resource_path(os.path.join('bin', 'geckoloader.bin')), 'rb') as kernelfile: geckoKernel = KernelLoader(kernelfile) @@ -228,16 +250,11 @@ if __name__ == "__main__": if args.init is not None: geckoKernel.initAddress = int(args.init, 16) - geckoKernel.codeLocation = args.movecodes + geckoKernel.patchJob = args.movecodes geckoKernel.verbosity = args.verbose geckoKernel.quiet = args.quiet geckoKernel.encrypt = args.encrypt - - with GC_File(os.path.normpath(args.dolfile), 'rb') as dol: - dolFile = DolFile(dol) - - codeHandler.optimizeList = args.optimize - geckoKernel.protect = args.protect + geckoKernel.protect = args.protect if args.dest: if os.path.splitext(args.dest)[1] == "": @@ -245,17 +262,16 @@ if __name__ == "__main__": else: dest = os.path.normpath(os.path.join(os.getcwd(), args.dest.lstrip('\\').lstrip('/'))) else: - dest = os.path.normpath(os.path.join(os.getcwd(), "BUILD", os.path.basename(args.dolfile))) + dest = os.path.normpath(os.path.join(os.getcwd(), "geckoloader-build", os.path.basename(args.dolfile))) if not os.path.exists(dest) and os.path.dirname(dest) not in ('', '/'): os.makedirs(os.path.dirname(dest), exist_ok=True) - if not os.path.exists(os.path.abspath(tmpdir)): - os.mkdir(tmpdir) - - geckoKernel.build(parser, args.codelist, dolFile, codeHandler, tmpdir, dest) - - shutil.rmtree(tmpdir) + if not os.path.isdir(TMPDIR): + os.mkdir(TMPDIR) + + geckoKernel.build(parser, args.codelist, dolFile, codeHandler, TMPDIR, dest) + sys.exit(0) except FileNotFoundError as e: diff --git a/README.md b/README.md index 6e72188..29616a6 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,6 @@ 2. In command prompt, input `GeckoLoader -h` for help on syntax and options 3. Run the command `GeckoLoader ` filling in the variables as needed -Your new patched `dol` file will be in the folder ./BUILD/ +Your new patched `dol` file will be in the folder `./geckoloader-build` by default *NOTE: can be an Ocarina formatted txt file, a gct, or a folder containing the previous mentioned files. diff --git a/dolreader.py b/dolreader.py index 658a1a5..80aab71 100644 --- a/dolreader.py +++ b/dolreader.py @@ -3,9 +3,13 @@ from io import BytesIO import tools from fileutils import * +class UnmappedAddressError(Exception): pass +class SectionCountFullError(Exception): pass +class AddressOutOfRangeError(Exception): pass + class DolFile(object): - def __init__(self, f: GC_File=None): + def __init__(self, f=None): self.fileOffsetLoc = 0 self.fileAddressLoc = 0x48 self.fileSizeLoc = 0x90 @@ -47,12 +51,15 @@ class DolFile(object): f.seek(self.fileEntryLoc) self.entryPoint = read_uint32(f) - self._currLogicAddr = self.textSections[0][1] + self._currLogicAddr = self.get_first_section()[1] self.seek(self._currLogicAddr) f.seek(0) + + def __str__(self): + return "Nintendo DOL format executable for the Wii and Gamecube" # Internal function for - def resolve_address(self, gcAddr, raiseError=True) -> (None, tuple): + def resolve_address(self, gcAddr, raiseError=True) -> tuple: '''Returns the data of the section that houses the given address If raiseError is True, a RuntimeError is raised when the address is unmapped, otherwise it returns None''' @@ -65,7 +72,7 @@ class DolFile(object): return offset, address, size, data if raiseError: - raise RuntimeError("Unmapped address: 0x{:X}".format(gcAddr)) + raise UnmappedAddressError(f"Unmapped address: 0x{gcAddr:X}") return None @@ -82,23 +89,54 @@ class DolFile(object): gcAddr = address + size return gcAddr + @property + def sections(self) -> tuple: + """ Generator that yields each section's data """ + for i in self.textSections: + yield i + for i in self.dataSections: + yield i + + return + def get_final_section(self) -> tuple: largestOffset = 0 indexToTarget = 0 - targetType = 0 + targetType = "Text" for i, sectionData in enumerate(self.textSections): if sectionData[0] > largestOffset: largestOffset = sectionData[0] indexToTarget = i - targetType = 0 + targetType = "Text" for i, sectionData in enumerate(self.dataSections): if sectionData[0] > largestOffset: largestOffset = sectionData[0] indexToTarget = i - targetType = 1 + targetType = "Data" - if targetType == 0: + if targetType == "Text": + return self.textSections[indexToTarget] + else: + return self.dataSections[indexToTarget] + + def get_first_section(self) -> tuple: + smallestOffset = 0xFFFFFFFF + indexToTarget = 0 + targetType = "Text" + + for i, sectionData in enumerate(self.textSections): + if sectionData[0] < smallestOffset: + smallestOffset = sectionData[0] + indexToTarget = i + targetType = "Text" + for i, sectionData in enumerate(self.dataSections): + if sectionData[0] < smallestOffset: + smallestOffset = sectionData[0] + indexToTarget = i + targetType = "Data" + + if targetType == "Text": return self.textSections[indexToTarget] else: return self.dataSections[indexToTarget] @@ -134,14 +172,14 @@ class DolFile(object): self._currLogicAddr += where else: - raise RuntimeError("Unsupported whence type '{}'".format(whence)) + raise RuntimeError(f"Unsupported whence type '{whence}'") def tell(self) -> int: return self._currLogicAddr - def save(self, f: GC_File): + def save(self, f): f.seek(0) - f.write(b"\x00" * 0x100) + f.write(b"\x00" * self.get_full_size()) for i in range(self.maxTextSections + self.maxDataSections): if i < self.maxTextSections: @@ -156,48 +194,37 @@ class DolFile(object): continue f.seek(self.fileOffsetLoc + (i * 4)) - f.write_uint32(offset) #offset in file + write_uint32(f, offset) #offset in file f.seek(self.fileAddressLoc + (i * 4)) - f.write_uint32(address) #game address + write_uint32(f, address) #game address f.seek(self.fileSizeLoc + (i * 4)) - f.write_uint32(size) #size in file - - if offset > f.get_size(): - f.seek(0, 2) - f.write(b"\x00" * (offset - f.get_size())) + write_uint32(f, size) #size in file f.seek(offset) f.write(data.getbuffer()) - f.align_file(32) f.seek(self.fileBssInfoLoc) - f.write_uint32(self.bssAddress) - f.write_uint32(self.bssSize) + write_uint32(f, self.bssAddress) + write_uint32(f, self.bssSize) f.seek(self.fileEntryLoc) - f.write_uint32(self.entryPoint) - f.align_file(256) + write_uint32(f, self.entryPoint) + align_byte_size(f, 256) def get_full_size(self) -> int: - fullSize = 0x100 - for section in self.textSections: - fullSize += section[2] - for section in self.dataSections: - fullSize += section[2] - return fullSize + offset, _, size, _ = self.get_final_section() + return (0x100 + offset + size + 255) & -256 def get_section_size(self, sectionsList: list, index: int) -> int: return sectionsList[index][2] def append_text_sections(self, sectionsList: list) -> bool: - """ Follows the list format: [tuple(Data, GameAddress or None), tuple(Data... - - Returns True if the operation can be performed, otherwise it returns False """ + """ Follows the list format: [tuple(Data, GameAddress or None), tuple(Data... """ '''Write offset/address/size to each section in DOL file header''' for i, dataSet in enumerate(sectionsList): if len(self.textSections) >= self.maxTextSections: - return False + raise SectionCountFullError(f"Exceeded max text section limit of {self.maxTextSections}") fOffset, _, fSize, _ = self.get_final_section() _, pAddress, pSize, _ = self.textSections[len(self.textSections) - 1] @@ -217,21 +244,17 @@ class DolFile(object): address = self.seek_nearest_unmapped(pAddress + pSize, size) if address < 0x80000000 or address >= 0x81200000: - raise ValueError("Address '{:08X}' of text section {} is beyond scope (0x80000000 <-> 0x81200000)".format(address, i)) + raise AddressOutOfRangeError(f"Address '{address:08X}' of text section {i} is beyond scope (0x80000000 <-> 0x81200000)") self.textSections.append((offset, address, size, data)) - return True - def append_data_sections(self, sectionsList: list) -> bool: - """ Follows the list format: [tuple(Data, GameAddress or None), tuple(Data... - - Returns True if the operation can be performed, otherwise it returns False """ + """ Follows the list format: [tuple(Data, GameAddress or None), tuple(Data... """ '''Write offset/address/size to each section in DOL file header''' for i, dataSet in enumerate(sectionsList): if len(self.dataSections) >= self.maxDataSections: - return False + raise SectionCountFullError(f"Exceeded max data section limit of {self.maxDataSections}") fOffset, _, fSize, _ = self.get_final_section() _, pAddress, pSize, _ = self.dataSections[len(self.dataSections) - 1] @@ -251,15 +274,13 @@ class DolFile(object): address = self.seek_nearest_unmapped(pAddress + pSize, size) if address < 0x80000000 or address >= 0x81200000: - raise ValueError("Address '{:08X}' of data section {} is beyond scope (0x80000000 <-> 0x81200000)".format(address, i)) + raise AddressOutOfRangeError(f"Address '{address:08X}' of data section {i} is beyond scope (0x80000000 <-> 0x81200000)") self.dataSections.append((offset, address, size, data)) - return True - def insert_branch(self, to: int, _from: int, lk=0): self.seek(_from) - f.write_uint32(self, (to - _from) & 0x3FFFFFD | 0x48000000 | lk) + write_uint32(self, (to - _from) & 0x3FFFFFD | 0x48000000 | lk) def extract_branch_addr(self, bAddr: int) -> tuple: """ Returns the branch offset of the given instruction, @@ -267,7 +288,7 @@ class DolFile(object): self.seek(bAddr) - ppc = f.read_uint32(self) + ppc = read_uint32(self) conditional = False if (ppc >> 24) & 0xFF < 0x48: @@ -305,10 +326,26 @@ class DolFile(object): return string + def print_info(self): + print("|---DOL INFO---|".center(20, " ")) + for i, (offset, addr, size, _) in enumerate(self.textSections): + header = f"| Text section {i} |" + print("-"*len(header) + "\n" + header + "\n" + "-"*len(header) + f"\n File offset:\t0x{offset:X}\n Virtual addr:\t0x{addr:X}\n Size:\t\t0x{size:X}\n") + + for i, (offset, addr, size, _) in enumerate(self.dataSections): + header = f"| Data section {i} |" + print("-"*len(header) + "\n" + header + "\n" + "-"*len(header) + f"\n File offset:\t0x{offset:X}\n Virtual addr:\t0x{addr:X}\n Size:\t\t0x{size:X}\n") + + header = "| BSS section |" + print("-"*len(header) + "\n" + header + "\n" + "-"*len(header) + f"\n Virtual addr:\t0x{self.bssAddress:X}\n Size:\t\t0x{self.bssSize:X}\n End:\t\t0x{self.bssAddress+self.bssSize:X}\n") + + header = "| Miscellaneous Info |" + print("-"*len(header) + "\n" + header + "\n" + "-"*len(header) + f"\n Text sections:\t{len(self.textSections)}\n Data sections:\t{len(self.dataSections)}\n File length:\t0x{self.get_full_size():X} bytes\n") + if __name__ == "__main__": # Example usage (Reading global string "mario" from Super Mario Sunshine (NTSC-U)) - with GC_File("Start.dol", "rb") as f: + with open("Start.dol", "rb") as f: dol = DolFile(f) name = dol.read_string(addr=0x804165A0) diff --git a/fileutils.py b/fileutils.py index 48da8a3..396b1c4 100644 --- a/fileutils.py +++ b/fileutils.py @@ -1,5 +1,4 @@ -from io import FileIO -from tools import get_alignment +from tools import get_alignment, align_byte_size import struct def read_sbyte(f): @@ -55,35 +54,4 @@ def read_bool(f, vSize=1): def write_bool(f, val, vSize=1): if val is True: f.write(b'\x00'*(vSize-1) + b'\x01') - else: f.write(b'\x00' * vSize) - -class GC_File(FileIO): - - def __init__(self, *args, **kwargs): - self._args = args - self._kwargs = kwargs - - def __enter__(self): - self._filestream = open(*self._args, **self._kwargs) - return self._filestream - - def __exit__(self, *args): - self._filestream.close() - - def size(self, ofs: int = 0): - _pos = self.tell() - self.seek(0, 2) - _size = self.tell() - self.seek(_pos, 1) - return _size + ofs - - def size_alignment(self, alignment: int): - """ Return file alignment, 0 = aligned, non zero = misaligned """ - return get_alignment(self.size(), alignment) - - def align_file_size(self, alignment: int, fillchar='00'): - """ Align a file to be the specified size """ - self.write(bytes.fromhex(fillchar * self.size_alignment(alignment))) - - - \ No newline at end of file + else: f.write(b'\x00' * vSize) \ No newline at end of file diff --git a/kernel.py b/kernel.py index 14697e4..e2e1916 100644 --- a/kernel.py +++ b/kernel.py @@ -24,18 +24,21 @@ def timer(func): end = time.perf_counter() print(tools.color_text(f'\n :: Completed in {(end - start):0.4f} seconds!\n', defaultColor=tools.TGREENLIT)) return value + return wrapper + +class InvalidGeckoCodeError(Exception): pass class GCT(object): - def __init__(self, f: GC_File): + def __init__(self, f: open): self.codeList = BytesIO(f.read()) - self.rawLineCount = f.size() >> 3 + self.rawLineCount = tools.stream_size(self.codeList) >> 3 self.lineCount = self.rawLineCount - 2 - self.size = f.size() + self.size = tools.stream_size(self.codeList) f.seek(0) @staticmethod - def determine_codelength(codetype, info): + def determine_codelength(codetype, info) -> int: if codetype.startswith(b'\x06'): bytelength = int.from_bytes(info, byteorder='big', signed=False) padding = get_alignment(bytelength, 8) @@ -122,9 +125,9 @@ class GCT(object): dolFile.seek(address) value = int.from_bytes(info, byteorder='big', signed=False) - data = self.codeList.read(2).hex() - size = int(data[:-3], 16) - counter = int(data[1:], 16) + data = read_uint16(self.codeList) + size = data & 0x3000 + counter = data & 0xFFF address_increment = read_uint16(self.codeList) value_increment = read_uint32(self.codeList) @@ -186,13 +189,13 @@ class CodeHandler(object): self._rawData = BytesIO(f.read()) '''Get codelist pointer''' - f.seek(0xFA) - codelistUpper = f.read(2).hex() - f.seek(0xFE) - codelistLower = f.read(2).hex() + self._rawData.seek(0xFA) + codelistUpper = self._rawData.read(2).hex() + self._rawData.seek(0xFE) + codelistLower = self._rawData.read(2).hex() self.codeListPointer = int(codelistUpper[2:] + codelistLower[2:], 16) - self.handlerLength = len(self._rawData.getbuffer()) + self.handlerLength = tools.stream_size(self._rawData) self.initAddress = 0x80001800 self.startAddress = 0x800018A8 @@ -217,7 +220,7 @@ class CodeHandler(object): f.seek(0) - def gecko_parser(self, geckoText): + def gecko_parser(self, geckoText) -> str: with open(r'{}'.format(geckoText), 'rb') as gecko: result = chardet.detect(gecko.read()) encodeType = result['encoding'] @@ -252,7 +255,7 @@ class CodeHandler(object): return geckoCodes @staticmethod - def encrypt_key(key: int): + def encrypt_key(key: int) -> int: b1 = key & 0xFF b2 = (key >> 8) & 0xFF b3 = (key >> 16) & 0xFF @@ -277,7 +280,7 @@ class CodeHandler(object): except: break - def find_variable_data(self, variable): + def find_variable_data(self, variable) -> int: self._rawData.seek(0) if self._rawData.read(4) == variable: @@ -314,11 +317,11 @@ class KernelLoader(object): def __init__(self, f): self._rawData = BytesIO(f.read()) - self.initDataList = None - self.gpModDataList = None - self.gpDiscDataList = None - self.gpKeyAddrList = None - self.codeLocation = None + self._initDataList = None + self._gpModDataList = None + self._gpDiscDataList = None + self._gpKeyAddrList = None + self.patchJob = None self.initAddress = None self.protect = False self.verbosity = 0 @@ -327,18 +330,17 @@ class KernelLoader(object): def set_variables(self, entryPoint: list, baseOffset: int=0): self._rawData.seek(0) - if self.gpModDataList is None: - return - sample = self._rawData.read(2) + if self._gpModDataList is None: + return - while sample: + while sample := self._rawData.read(2): if sample == b'GH': self._rawData.seek(-2, 1) - write_uint16(self._rawData, self.gpModDataList[0]) + write_uint16(self._rawData, self._gpModDataList[0]) elif sample == b'GL': self._rawData.seek(-2, 1) - write_uint16(self._rawData, baseOffset + self.gpModDataList[1]) + write_uint16(self._rawData, baseOffset + self._gpModDataList[1]) elif sample == b'IH': self._rawData.seek(-2, 1) write_uint16(self._rawData, entryPoint[0]) @@ -347,29 +349,27 @@ class KernelLoader(object): write_uint16(self._rawData, entryPoint[1]) elif sample == b'KH': self._rawData.seek(-2, 1) - write_uint16(self._rawData, self.gpKeyAddrList[0]) + write_uint16(self._rawData, self._gpKeyAddrList[0]) elif sample == b'KL': self._rawData.seek(-2, 1) - write_uint16(self._rawData, baseOffset + self.gpKeyAddrList[1]) - - sample = self._rawData.read(2) + write_uint16(self._rawData, baseOffset + self._gpKeyAddrList[1]) def complete_data(self, codeHandler: CodeHandler, initpoint: list): - upperAddr, lowerAddr = ((self.initAddress >> 16) & 0xFFFF, self.initAddress & 0xFFFF) - key = random.randrange(0x100000000) + _upperAddr, _lowerAddr = ((self.initAddress >> 16) & 0xFFFF, self.initAddress & 0xFFFF) + _key = random.randrange(0x100000000) self._rawData.seek(0) sample = self._rawData.read(4) - while sample: + while sample := self._rawData.read(4): if sample == b'HEAP': #Found keyword "HEAP". Goes with the resize of the heap self._rawData.seek(-4, 1) gpModInfoOffset = self._rawData.tell() - if lowerAddr + gpModInfoOffset > 0x7FFF: #Absolute addressing - gpModUpperAddr = upperAddr + 1 + if _lowerAddr + gpModInfoOffset > 0x7FFF: #Absolute addressing + gpModUpperAddr = _upperAddr + 1 else: - gpModUpperAddr = upperAddr + gpModUpperAddr = _upperAddr if codeHandler.allocation == None: codeHandler.allocation = (codeHandler.handlerLength + codeHandler.geckoCodes.size + 7) & -8 @@ -400,22 +400,20 @@ class KernelLoader(object): self._rawData.seek(-4, 1) gpKeyOffset = self._rawData.tell() - if lowerAddr + gpKeyOffset > 0x7FFF: #Absolute addressing - gpKeyUpperAddr = upperAddr + 1 + if _lowerAddr + gpKeyOffset > 0x7FFF: #Absolute addressing + gpKeyUpperAddr = _upperAddr + 1 else: - gpKeyUpperAddr = upperAddr + gpKeyUpperAddr = _upperAddr - write_uint32(self._rawData, CodeHandler.encrypt_key(key)) - - sample = self._rawData.read(4) + write_uint32(self._rawData, CodeHandler.encrypt_key(_key)) - self.gpModDataList = (gpModUpperAddr, gpModInfoOffset) - self.gpKeyAddrList = (gpKeyUpperAddr, gpKeyOffset) + self._gpModDataList = (gpModUpperAddr, gpModInfoOffset) + self._gpKeyAddrList = (gpKeyUpperAddr, gpKeyOffset) - self.set_variables(initpoint, lowerAddr) + self.set_variables(initpoint, _lowerAddr) if self.encrypt: - codeHandler.encrypt_data(key) + codeHandler.encrypt_data(_key) def patch_arena(self, codeHandler: CodeHandler, dolFile: DolFile): @@ -425,26 +423,21 @@ class KernelLoader(object): self._rawData.write(codeHandler._rawData.getvalue() + codeHandler.geckoCodes.codeList.getvalue()) self._rawData.seek(0) - kernelData = self._rawData.getvalue() + _kernelData = self._rawData.getvalue() - status = dolFile.append_text_sections([(kernelData, self.initAddress)]) - - if status is True: - dolFile.entryPoint = self.initAddress - - return status + dolFile.append_text_sections([(_kernelData, self.initAddress)]) + dolFile.entryPoint = self.initAddress def patch_legacy(self, codeHandler: CodeHandler, dolFile: DolFile): codeHandler._rawData.seek(0) codeHandler.geckoCodes.codeList.seek(0) - handlerData = codeHandler._rawData.getvalue() + codeHandler.geckoCodes.codeList.getvalue() + _handlerData = codeHandler._rawData.getvalue() + codeHandler.geckoCodes.codeList.getvalue() - status = dolFile.append_text_sections([(handlerData, codeHandler.initAddress)]) - return status + dolFile.append_text_sections([(_handlerData, codeHandler.initAddress)]) def protect_game(self, codeHandler: CodeHandler): - oldpos = codeHandler.geckoCodes.codeList.tell() + _oldpos = codeHandler.geckoCodes.codeList.tell() protectdata = [b'\xC0\x00\x00\x00\x00\x00\x00\x17', b'\x7C\x08\x02\xA6\x94\x21\xFF\x70', @@ -477,11 +470,11 @@ class KernelLoader(object): codeHandler.geckoCodes.codeList.write(b'\xF0\x00\x00\x00\x00\x00\x00\x00') codeHandler.geckoCodes.codeList.seek(0, 2) codeHandler.geckoCodes.size = codeHandler.geckoCodes.codeList.tell() - codeHandler.geckoCodes.codeList.seek(oldpos) + codeHandler.geckoCodes.codeList.seek(_oldpos) @timer def build(self, parser: tools.CommandLineParser, gctFile, dolFile: DolFile, codeHandler: CodeHandler, tmpdir, dump): - with GC_File(dump, 'wb+') as final: + with open(dump, 'wb+') as final: if dolFile.get_full_size() < 0x100: parser.error(tools.color_text('DOL header is corrupted. Please provide a clean file\n', defaultColor=tools.TREDLIT), exit=False) @@ -507,7 +500,7 @@ class KernelLoader(object): else: with open(os.path.join(tmpdir, 'gct.bin'), 'wb+') as temp: - temp.write(bytes.fromhex('00D0C0DE'*2)) + temp.write(b'\x00\xD0\xC0\xDE'*2) for file in os.listdir(gctFile): if os.path.isfile(os.path.join(gctFile, file)): @@ -521,7 +514,7 @@ class KernelLoader(object): else: print(tools.color_text(f' :: HINT: {file} is not a .txt or .gct file', defaultColor=tools.TYELLOWLIT)) - temp.write(bytes.fromhex('F000000000000000')) + temp.write(b'\xF0\x00\x00\x00\x00\x00\x00\x00') temp.seek(0) codeHandler.geckoCodes = GCT(temp) @@ -529,14 +522,14 @@ class KernelLoader(object): parser.error(tools.color_text('No valid gecko code file found\n', defaultColor=tools.TREDLIT), exit=False) return - if self.protect and self.codeLocation == "ARENA": + if self.protect and self.patchJob == "ARENA": self.protect_game(codeHandler) - if self.codeLocation == 'AUTO': + if self.patchJob == 'AUTO': if codeHandler.initAddress + codeHandler.handlerLength + codeHandler.geckoCodes.size > 0x80002FFF: - self.codeLocation = 'ARENA' + self.patchJob = 'ARENA' else: - self.codeLocation = 'LEGACY' + self.patchJob = 'LEGACY' '''Get entrypoint (or BSS midpoint) for insert''' @@ -558,29 +551,24 @@ class KernelLoader(object): print(tools.color_text('\n :: All codes have been successfully pre patched', defaultColor=tools.TGREENLIT)) return - if self.codeLocation == 'LEGACY': + if self.patchJob == 'LEGACY': codeHandler.allocation = 0x80003000 - (codeHandler.initAddress + codeHandler.handlerLength) codeHandler.set_variables(dolFile) hooked = determine_codehook(dolFile, codeHandler, True) - status = self.patch_legacy(codeHandler, dolFile) + self.patch_legacy(codeHandler, dolFile) legacy = True else: hooked = determine_codehook(dolFile, codeHandler, False) - status = self.patch_arena(codeHandler, dolFile) + self.patch_arena(codeHandler, dolFile) legacy = False if not hooked: parser.error(tools.color_text('Failed to find a hook address. Try using option --codehook to use your own address\n', defaultColor=tools.TREDLIT)) - - if status is False: - parser.error(tools.color_text('Not enough text sections to patch the DOL file! Potentially due to previous mods?\n', defaultColor=tools.TREDLIT), exit=False) - return - - dolFile.save(final) - - if codeHandler.allocation < codeHandler.geckoCodes.size: - print(tools.color_text('\n :: WARNING: Allocated codespace was smaller than the given codelist. The game will crash if run', defaultColor=tools.TYELLOW)) + elif codeHandler.allocation < codeHandler.geckoCodes.size: + parser.error(tools.color_text('\n :: Error: Allocated codespace was smaller than the given codelist.\n', defaultColor=tools.TYELLOW)) + dolFile.save(final) + if self.quiet: return @@ -634,7 +622,7 @@ def determine_codehook(dolFile: DolFile, codeHandler: CodeHandler, hook=False): def assert_code_hook(dolFile: DolFile, codeHandler: CodeHandler): for _, address, size, _, in dolFile.textSections: - dolFile.seek(address, 0) + dolFile.seek(address) sample = dolFile.read(size) if codeHandler.hookType == 'VI': @@ -647,8 +635,7 @@ def assert_code_hook(dolFile: DolFile, codeHandler: CodeHandler): raise NotImplementedError(tools.color_text(f'Unsupported hook type specified ({codeHandler.hookType})', defaultColor=tools.TREDLIT)) if result >= 0: - dolFile.seek(address, 0) - dolFile.seek(result, 1) + dolFile.seek(address + result) else: if codeHandler.hookType == 'VI': result = sample.find(codeHandler.wiiVIHook) @@ -660,14 +647,12 @@ def assert_code_hook(dolFile: DolFile, codeHandler: CodeHandler): raise NotImplementedError(tools.color_text(f'Unsupported hook type specified ({codeHandler.hookType})', defaultColor=tools.TREDLIT)) if result >= 0: - dolFile.seek(address, 0) - dolFile.seek(result, 1) + dolFile.seek(address + result) else: continue - sample = read_uint32(dolFile) - while sample != 0x4E800020: - sample = read_uint32(dolFile) + while (sample := read_uint32(dolFile)) != 0x4E800020: + pass dolFile.seek(-4, 1) codeHandler.hookAddress = dolFile.tell() diff --git a/tools.py b/tools.py index ad2817f..ba4f3b1 100644 --- a/tools.py +++ b/tools.py @@ -1,6 +1,7 @@ import struct import sys import os +from io import IOBase from argparse import ArgumentParser try: @@ -16,27 +17,50 @@ try: TREDLIT = Style.BRIGHT + Fore.RED except ImportError: - TRESET = '' - TGREEN = '' - TGREENLIT = '' - TYELLOW = '' - TYELLOWLIT = '' - TRED = '' - TREDLIT = '' + TRESET = "" + TGREEN = "" + TGREENLIT = "" + TYELLOW = "" + TYELLOWLIT = "" + TRED = "" + TREDLIT = "" -def get_alignment(number, align: int): +def get_alignment(number: int, align: int) -> int: if number % align != 0: return align - (number % align) else: return 0 -def color_text(text: str, textToColor: list=[('', None)], defaultColor: str=None): +def stream_size(obj, ofs: int = 0) -> int: + if hasattr(obj, "getbuffer"): + return len(obj.getbuffer()) + ofs + elif hasattr(obj, "tell") and hasattr(obj, "seek"): + _pos = obj.tell() + obj.seek(0, 2) + _size = obj.tell() + obj.seek(_pos, 1) + return _size + ofs + else: + raise NotImplementedError(f"Getting the stream size of class {type(obj)} is unsupported") + +def align_byte_size(obj, alignment: int, fillchar="00"): + if isinstance(obj, bytes): + obj += bytes.fromhex(fillchar * get_alignment(len(obj), alignment)) + elif isinstance(obj, bytearray): + obj.append(bytearray.fromhex(fillchar * get_alignment(len(obj), alignment))) + elif issubclass(type(obj), IOBase): + _size = stream_size(obj) + obj.write(bytes.fromhex(fillchar * get_alignment(_size, alignment))) + else: + raise NotImplementedError(f"Aligning the size of class {type(obj)} is unsupported") + +def color_text(text: str, textToColor: list=[("", None)], defaultColor: str=None) -> str: currentColor = None - formattedText = '' + formattedText = "" format = False for itemPair in textToColor: - if itemPair[0] != '' and itemPair[1] is not None: + if itemPair[0] != "" and itemPair[1] is not None: format = True break @@ -46,7 +70,7 @@ def color_text(text: str, textToColor: list=[('', None)], defaultColor: str=None for char in text: handled = False for itemPair in textToColor: - if (char in itemPair[0] or r'\*' in itemPair[0]) and itemPair[1] is not None: + if (char in itemPair[0] or r"\*" in itemPair[0]) and itemPair[1] is not None: if currentColor != itemPair[1]: formattedText += TRESET formattedText += itemPair[1] @@ -77,11 +101,11 @@ class CommandLineParser(ArgumentParser): if prefix is None: if exit: - self.exit(2, f'{self.prog}: error: {message}\n') + self.exit(2, f"{self.prog}: error: {message}\n") else: - self._print_message(f'{self.prog}: error: {message}\n') + self._print_message(f"{self.prog}: error: {message}\n") else: if exit: - self.exit(2, f'{prefix} {message}\n') + self.exit(2, f"{prefix} {message}\n") else: - self._print_message(f'{prefix} {message}\n') \ No newline at end of file + self._print_message(f"{prefix} {message}\n") \ No newline at end of file diff --git a/versioncheck.py b/versioncheck.py index d7b25e6..2a001af 100644 --- a/versioncheck.py +++ b/versioncheck.py @@ -1,7 +1,7 @@ from urllib import request from bs4 import BeautifulSoup -class Updater: +class Updater(object): def __init__(self, owner: str, repository: str): self.owner = owner @@ -14,7 +14,7 @@ class Updater: html = response.read() return html - def get_newest_version(self): + def get_newest_version(self) -> str: '''Returns newest release version''' try: response = self.request_release_data()