Finished formatting, added section type to tuple

dolreader.py

@@ -9,6 +9,10 @@ class AddressOutOfRangeError(Exception): pass
 
 class DolFile(object):
 
+    class SectionType():
+        Text = 0
+        Data = 1
+
     def __init__(self, f=None):
         self.fileOffsetLoc = 0
         self.fileAddressLoc = 0x48
@@ -40,9 +44,9 @@ class DolFile(object):
                 f.seek(offset)
                 data = BytesIO(f.read(size))
                 if i < self.maxTextSections:
-                    self.textSections.append((offset, address, size, data))
+                    self.textSections.append([offset, address, size, data, DolFile.SectionType.Text])
                 else:
-                    self.dataSections.append((offset, address, size, data))
+                    self.dataSections.append([offset, address, size, data, DolFile.SectionType.Data])
         
         f.seek(self.fileBssInfoLoc)
         self.bssAddress = read_uint32(f)
@@ -59,31 +63,27 @@ class DolFile(object):
         return "Nintendo DOL format executable for the Wii and Gamecube"
         
     # Internal function for 
-    def resolve_address(self, gcAddr, raiseError=True) -> tuple:
+    def resolve_address(self, gcAddr) -> tuple:
-        '''Returns the data of the section that houses the given address
+        """ Returns the data of the section that houses the given address\n
-           If raiseError is True, a RuntimeError is raised when the address is unmapped,
+            UnmappedAddressError is raised when the address is unmapped """
-           otherwise it returns None'''
 
-        for offset, address, size, data in self.textSections:
+        for offset, address, size, data, sectiontype in self.textSections:
             if address <= gcAddr < address+size:
-                return offset, address, size, data
+                return offset, address, size, data, sectiontype
-        for offset, address, size, data in self.dataSections:
+        for offset, address, size, data, sectiontype in self.dataSections:
             if address <= gcAddr < address+size:
-                return offset, address, size, data
+                return offset, address, size, data, sectiontype
         
-        if raiseError:
+        raise UnmappedAddressError(f"Unmapped address: 0x{gcAddr:X}")
-            raise UnmappedAddressError(f"Unmapped address: 0x{gcAddr:X}")
-
-        return None
 
     def seek_nearest_unmapped(self, gcAddr, buffer=0) -> int:
         '''Returns the nearest unmapped address (greater) if the given address is already taken by data'''
         
-        for _, address, size, _ in self.textSections:
+        for _, address, size, _, _ in self.textSections:
             if address > (gcAddr + buffer) or address+size < gcAddr:
                 continue
             gcAddr = address + size
-        for _, address, size, _ in self.dataSections:
+        for _, address, size, _, _ in self.dataSections:
             if address > (gcAddr + buffer) or address+size < gcAddr:
                 continue
             gcAddr = address + size
@@ -92,6 +92,7 @@ class DolFile(object):
     @property
     def sections(self) -> tuple:
         """ Generator that yields each section's data """
+
         for i in self.textSections:
             yield i
         for i in self.dataSections:
@@ -100,43 +101,47 @@ class DolFile(object):
         return
 
     def get_final_section(self) -> tuple:
+        """ Returns the last section in the dol file as sorted by internal offset """
+
         largestOffset = 0
         indexToTarget = 0
-        targetType = "Text"
+        targetType = DolFile.SectionType.Text
 
         for i, sectionData in enumerate(self.textSections):
             if sectionData[0] > largestOffset:
                 largestOffset = sectionData[0]
                 indexToTarget = i
-                targetType = "Text"
+                targetType = DolFile.SectionType.Text
         for i, sectionData in enumerate(self.dataSections):
             if sectionData[0] > largestOffset:
                 largestOffset = sectionData[0]
                 indexToTarget = i
-                targetType = "Data"
+                targetType = DolFile.SectionType.Data
         
-        if targetType == "Text":
+        if targetType == DolFile.SectionType.Text:
             return self.textSections[indexToTarget]
         else:
             return self.dataSections[indexToTarget]
 
     def get_first_section(self) -> tuple:
+        """ Returns the first section in the dol file as sorted by internal offset """
+
         smallestOffset = 0xFFFFFFFF
         indexToTarget = 0
-        targetType = "Text"
+        targetType = DolFile.SectionType.Text
 
         for i, sectionData in enumerate(self.textSections):
             if sectionData[0] < smallestOffset:
                 smallestOffset = sectionData[0]
                 indexToTarget = i
-                targetType = "Text"
+                targetType = DolFile.SectionType.Text
         for i, sectionData in enumerate(self.dataSections):
             if sectionData[0] < smallestOffset:
                 smallestOffset = sectionData[0]
                 indexToTarget = i
-                targetType = "Data"
+                targetType = DolFile.SectionType.Data
         
-        if targetType == "Text":
+        if targetType == DolFile.SectionType.Text:
             return self.textSections[indexToTarget]
         else:
             return self.dataSections[indexToTarget]
@@ -144,35 +149,35 @@ class DolFile(object):
     # Unsupported: Reading an entire dol file 
     # Assumption: A read should not go beyond the current section 
     def read(self, _size) -> bytes:
-        _, address, size, data = self.resolve_address(self._currLogicAddr)
+        _, address, size, data, _ = self.resolve_address(self._currLogicAddr)
         if self._currLogicAddr + _size > address + size:
-            raise RuntimeError("Read goes over current section")
+            raise UnmappedAddressError("Read goes over current section")
             
         self._currLogicAddr += _size  
         return data.read(_size)
         
     # Assumption: A write should not go beyond the current section 
     def write(self, _data):
-        _, address, size, data = self.resolve_address(self._currLogicAddr)
+        _, address, size, data, _ = self.resolve_address(self._currLogicAddr)
         if self._currLogicAddr + len(_data) > address + size:
-            raise RuntimeError("Write goes over current section")
+            raise UnmappedAddressError("Write goes over current section")
             
         data.write(_data)
         self._currLogicAddr += len(_data)
     
     def seek(self, where, whence=0):
         if whence == 0:
-            _, address, _, data = self.resolve_address(where)
+            _, address, _, data, _ = self.resolve_address(where)
             data.seek(where - address)
             
             self._currLogicAddr = where
         elif whence == 1:
-            _, address, _, data = self.resolve_address(self._currLogicAddr + where)
+            _, address, _, data, _ = self.resolve_address(self._currLogicAddr + where)
             data.seek((self._currLogicAddr + where) - address)
             
             self._currLogicAddr += where
         else:
-            raise RuntimeError(f"Unsupported whence type '{whence}'")
+            raise NotImplementedError(f"Unsupported whence type '{whence}'")
         
     def tell(self) -> int:
         return self._currLogicAddr
@@ -184,12 +189,12 @@ class DolFile(object):
         for i in range(self.maxTextSections + self.maxDataSections):
             if i < self.maxTextSections:
                 if i < len(self.textSections):
-                    offset, address, size, data = self.textSections[i]
+                    offset, address, size, data, _ = self.textSections[i]
                 else:
                     continue
             else:
                 if i - self.maxTextSections < len(self.dataSections):
-                    offset, address, size, data = self.dataSections[i - self.maxTextSections]
+                    offset, address, size, data, _ = self.dataSections[i - self.maxTextSections]
                 else:
                     continue
 
@@ -212,26 +217,36 @@ class DolFile(object):
         align_byte_size(f, 256)
 
     def get_full_size(self) -> int:
-        offset, _, size, _ = self.get_final_section()
+        offset, _, size, _, _ = self.get_final_section()
-        return (0x100 + offset + size + 255) & -256
+        return (offset + size + 255) & -256
+
+    def get_section_size(self, index: int, section: SectionType) -> int:
+        """ Return the current size of the specified section\n
+            section: DolFile.SectionType """
+
+        if section == DolFile.SectionType.Text:
+            return self.textSections[index][2]
+        else:
+            return self.dataSections[index][2]
 
-    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(<Bytes>Data, <Int>GameAddress or None), tuple(<Bytes>Data... """
 
-        '''Write offset/address/size to each section in DOL file header'''
         for i, dataSet in enumerate(sectionsList):
             if len(self.textSections) >= self.maxTextSections:
                 raise SectionCountFullError(f"Exceeded max text section limit of {self.maxTextSections}")
 
-            fOffset, _, fSize, _ = self.get_final_section()
+            fOffset, _, fSize, _, _ = self.get_final_section()
-            _, pAddress, pSize, _ = self.textSections[len(self.textSections) - 1]
+            _, pAddress, pSize, _, _ = self.textSections[len(self.textSections) - 1]
             data, address = dataSet
             
-            if not isinstance(data, BytesIO):
+            if not hasattr(data, "getbuffer"):
-                data = BytesIO(data)
+                if hasattr(data, "read"):
+                    data.seek(0)
+                    data = BytesIO(data.read())
+                else:
+                    data = BytesIO(data)
 
             offset = fOffset + fSize
 
@@ -246,22 +261,25 @@ class DolFile(object):
             if address < 0x80000000 or address >= 0x81200000:
                 raise AddressOutOfRangeError(f"Address '{address:08X}' of text section {i} is beyond scope (0x80000000 <-> 0x81200000)")
 
-            self.textSections.append((offset, address, size, data))
+            self.textSections.append((offset, address, size, data, DolFile.SectionType.Text))
 
     def append_data_sections(self, sectionsList: list) -> bool:
         """ Follows the list format: [tuple(<Bytes>Data, <Int>GameAddress or None), tuple(<Bytes>Data... """
 
-        '''Write offset/address/size to each section in DOL file header'''
         for i, dataSet in enumerate(sectionsList):
             if len(self.dataSections) >= self.maxDataSections:
                 raise SectionCountFullError(f"Exceeded max data section limit of {self.maxDataSections}")
 
-            fOffset, _, fSize, _ = self.get_final_section()
+            fOffset, _, fSize, _, _ = self.get_final_section()
-            _, pAddress, pSize, _ = self.dataSections[len(self.dataSections) - 1]
+            _, pAddress, pSize, _, _ = self.dataSections[len(self.dataSections) - 1]
             data, address = dataSet
 
-            if not isinstance(data, BytesIO):
+            if not hasattr(data, "getbuffer"):
-                data = BytesIO(data)
+                if hasattr(data, "read"):
+                    data.seek(0)
+                    data = BytesIO(data.read())
+                else:
+                    data = BytesIO(data)
 
             offset = fOffset + fSize
 
@@ -276,9 +294,14 @@ class DolFile(object):
             if address < 0x80000000 or address >= 0x81200000:
                 raise AddressOutOfRangeError(f"Address '{address:08X}' of data section {i} is beyond scope (0x80000000 <-> 0x81200000)")
 
-            self.dataSections.append((offset, address, size, data))
+            self.dataSections.append((offset, address, size, data, DolFile.SectionType.Data))
 
     def insert_branch(self, to: int, _from: int, lk=0):
+        """ Insert a branch instruction at _from\n
+            to:    address to branch to\n
+            _from: address to branch from\n
+            lk:    0 | 1, is branch linking? """
+
         self.seek(_from)
         write_uint32(self, (to - _from) & 0x3FFFFFD | 0x48000000 | lk)
 
@@ -328,6 +351,7 @@ class DolFile(object):
 
     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")

kernel.py

@@ -621,7 +621,7 @@ def determine_codehook(dolFile: DolFile, codeHandler: CodeHandler, hook=False):
 
 
 def assert_code_hook(dolFile: DolFile, codeHandler: CodeHandler):
-    for _, address, size, _, in dolFile.textSections:
+    for _, address, size, _, _ in dolFile.textSections:
         dolFile.seek(address)
         sample = dolFile.read(size)