Cleanup

src/retro_data_structures/disc/game_disc.py

@@ -102,15 +102,23 @@ def parse(cls, file_path: Path) -> GameDisc:
 
     def _get_file_entry(self, name: str) -> FileEntry:
         file_entry = self._file_tree
-        for segment in name.split("/"):
-            file_entry = file_entry[segment]
+
+        try:
+            for segment in name.split("/"):
+                file_entry = file_entry[segment]
+        except KeyError:
+            raise FileNotFoundError(f"{name} does not exist")
 
         if isinstance(file_entry, FileEntry):
             return file_entry
         else:
             raise OSError(f"{name} is a directory")
 
     def files(self) -> list[str]:
+        """
+        Lists all files in this disc. For Wii, it's only the data partition.
+        :return:
+        """
         result = []
 
         def recurse(parent: str, tree: FileTree) -> None:
@@ -133,15 +141,29 @@ def _open_data_at_offset(self, offset: int, size: int) -> disc_common.DiscFileRe
             return disc_common.DiscFileReader(self._file_path, offset, size)
 
     def open_binary(self, name: str) -> disc_common.DiscFileReader:
+        """
+        Returns an IOBase for reading a file with the given name.
+        :param name:
+        :return:
+        """
         entry = self._get_file_entry(name)
         return self._open_data_at_offset(entry.offset, entry.size)
 
     def read_binary(self, name: str) -> bytes:
+        """
+        Returns the entire contents of the file with given name.
+        :param name:
+        :return:
+        """
         entry = self._get_file_entry(name)
         with self._open_data_at_offset(entry.offset, entry.size) as file:
             return file.read(entry.size)
 
     def get_dol(self) -> bytes:
+        """
+        Gets the main dol for this disc. With Wii discs, returns the dol in the data partition.
+        :return:
+        """
         disc_header = self._raw.data_partition.disc_header
         with self._open_data_at_offset(disc_header.main_executable_offset, -1) as file:
             header = dol.DolHeader.parse_stream(file)

src/retro_data_structures/disc/wii_disc.py

@@ -222,13 +222,14 @@ def __init__(self, file: Path | typing.BinaryIO, size: int, dec_key: bytes, base
         self._initial_offset = offset
         self._offset = offset
         self._cur_block = -1
+        self._dec_buf = bytearray(0x8000 - 0x400)
 
     def _decrypt_block(self, block: int) -> None:
         self._cur_block = block
         self._file.seek(self._base_offset + self._cur_block * 0x8000)
-        enc_buf = self._file.read(0x8000)
+        enc_buf = memoryview(self._file.read(0x8000))
         aes = AES.new(key=self._dec_key, mode=AES.MODE_CBC, iv=enc_buf[0x3D0:0x3E0])
-        self._dec_buf = aes.decrypt(enc_buf[0x400:])
+        aes.decrypt(enc_buf[0x400:], self._dec_buf)
 
     def read(self, size: int = -1) -> bytes:
         block_quot = self._offset // 0x7C00
@@ -237,7 +238,7 @@ def read(self, size: int = -1) -> bytes:
         if size == -1:
             size = self._size - (self._offset - self._initial_offset)
 
-        ret = b""
+        ret = bytearray()
         rem = size
 
         while rem > 0:
@@ -248,7 +249,7 @@ def read(self, size: int = -1) -> bytes:
             if cache_size + block_rem > 0x7C00:
                 cache_size = 0x7C00 - block_rem
 
-            ret += self._dec_buf[block_rem : block_rem + cache_size]
+            ret += memoryview(self._dec_buf)[block_rem : block_rem + cache_size]
             rem -= cache_size
             block_rem = 0
             block_quot += 1
@@ -309,7 +310,7 @@ def __init__(self, source: typing.BinaryIO, part_info: construct.Container):
 
     def begin_read_stream(
         self, file_io: Path | typing.BinaryIO, offset: int, file_size: int
-    ) -> EncryptedDiscFileReader:
+    ) -> disc_common.DiscFileReader:
         return EncryptedDiscFileReader(file_io, file_size, self._dec_key, self._data_offset, offset)