Optimized WiiFS for speed

Author: niemasd

README.md

@@ -58,7 +58,7 @@ with open(target_path, 'rb') as target_file:
 ```
 
 ### [`WiiFS`](https://niema.net/NiemaFS/#niemafs.WiiFS) — Nintendo Wii DVD
-Note that, due to the need to decrypt the filesystem, this is extremely slow and memory-intensive (each decrypted partition is loaded into memory). Optimizations are possible and may be explored in the future.
+Note that, due to the need to decrypt the filesystem, this is extremely memory-intensive (each partition is loaded into memory to process in parallel for speed).
 
 ```python
 from niemafs import WiiFS

niemafs/__init__.py

@@ -1,15 +1,26 @@
 #! /usr/bin/env python
-from niemafs.common import clean_string, FileSystem, open_file
+# standard imports
+from warnings import warn
+
+# NiemaFS imports
+from niemafs.common import clean_string, FileSystem, open_file, safename
 from niemafs.dir import DirFS
 from niemafs.gcm import GcmFS
 from niemafs.iso import IsoFS
-from niemafs.wii import WiiFS
 from niemafs.zip import ZipFS
+
+# build __all__
 __all__ = [
-    'clean_string', 'FileSystem', 'open_file', # common.py
-    'DirFS',                                   # dir.py
-    'GcmFS',                                   # gcm.py
-    'IsoFS',                                   # iso.py
-    'WiiFS',                                   # wii.py
-    'ZipFS',                                   # zip.py
+    'clean_string', 'FileSystem', 'open_file', 'safename', # common.py
+    'DirFS',                                               # dir.py
+    'GcmFS',                                               # gcm.py
+    'IsoFS',                                               # iso.py
+    'ZipFS',                                               # zip.py
 ]
+
+# WiiFS depends on PyCryptodome, so import it afterwards
+try:
+    from niemafs.wii import WiiFS
+    __all__.append('WiiFS')
+except:
+    warn("Unable to import WiiFS. Ensure that you have PyCryptodome installed")

niemafs/common.py

@@ -13,21 +13,33 @@
 # constants
 DEFAULT_BUFFER_SIZE = 8388608 # 8 MB
 DEFAULT_COMPRESS_LEVEL = 9
+SAFE_CHARS = set('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-')
 
 def clean_string(s):
-    '''Clean a string (binary or normal) by right-stripping 0x00 and spaces
+    '''Clean a string (binary or normal) by right-stripping 0x00 and spaces.
 
     Args:
-        `s` (`bytes`): The ISO 9660 string to clean
+        `s` (`bytes`): The ISO 9660 string to clean.
 
     Returns:
-        `str`: The cleaned string
+        `str`: The cleaned string.
     '''
     if isinstance(s, bytes):
         return s.rstrip(b'\x00').decode().rstrip()
     else:
         return s.rstrip()
 
+def safename(s):
+    '''Convert a string into a version that is safe for a filename
+    
+    Args:
+        `s` (`str`): The original string.
+
+    Returns:
+        `str`: A version of `s` that is safe for a filename.
+    '''
+    return ''.join(c if c in SAFE_CHARS else '_' for c in s)
+
 def open_file(path, mode='rb', buffering=DEFAULT_BUFFER_SIZE, compresslevel=DEFAULT_COMPRESS_LEVEL):
     '''Open a file for reading, writing, or appending. Automatically handles GZIP compression.
 

niemafs/wii.py

@@ -4,11 +4,13 @@
 '''
 
 # NiemaFS imports
-from niemafs.common import clean_string, FileSystem
+from niemafs.common import clean_string, FileSystem, safename
 from niemafs.gcm import GcmFS
 
 # imports
+from Crypto.Cipher import AES
 from datetime import datetime
+from multiprocessing import Pool
 from pathlib import Path
 from struct import unpack
 from warnings import warn
@@ -30,6 +32,12 @@
     4: 'Korea',
 }
 
+# helper function for multiprocessing of AES decryption
+def aes_helper(aes_tup):
+    decrypt_key, cluster_data_encrypted = aes_tup
+    aes = AES.new(decrypt_key, AES.MODE_CBC, cluster_data_encrypted[0x3D0 : 0x3E0])
+    return aes.decrypt(cluster_data_encrypted[0x0400:])
+
 class WiiFS(FileSystem):
     '''Class to represent a `Nintendo Wii DVD <https://wiibrew.org/wiki/Wii_disc#%22System_Area%22>`_.'''
     def __init__(self, file_obj, path=None):
@@ -272,9 +280,6 @@ def parse_ticket(data):
         return out
 
     def __iter__(self):
-        # import PyCryptodome within WiiFS to allow people to use the other NiemaFS classes without it
-        from Crypto.Cipher import AES
-
         # parse each partition table
         for volume_num, parsed_partition_table in enumerate(self.parse_partition_tables()):
             volume_path = Path('Volume %d' % volume_num)
@@ -317,19 +322,16 @@ def __iter__(self):
                 data_start_offset = partition['offset'] + partition_data_offset
                 data_end_offset = data_start_offset + partition_data_size
                 if WiiFS.parse_header(self.get_header())['disable_disc_encryption'] == 0:
-                    data_decrypted = b''
-                    for cluster_offset in range(data_start_offset, data_end_offset, 0x8000):
-                        cluster_data_encrypted = self.read_file(cluster_offset, 0x8000)
-                        aes = AES.new(decrypt_key, AES.MODE_CBC, cluster_data_encrypted[0x3D0 : 0x3E0])
-                        data_decrypted += aes.decrypt(cluster_data_encrypted[0x0400:])
+                    with Pool(processes=None) as pool:
+                        data_decrypted = b''.join(pool.map(aes_helper, [(decrypt_key, self.read_file(cluster_offset, 0x8000)) for cluster_offset in range(data_start_offset, data_end_offset, 0x8000)]))
                 else:
                     data_decrypted = self.read_file(data_start_offset, partition_data_size)
 
                 # yield partition header data
                 partition_header = WiiFS.parse_header(data_decrypted[0x0000 : 0x0420])
-                partition_path = volume_path / ('Partition %d (%s) (%s%s) (%s)' % (partition_num, partition_type.capitalize(), partition_header['game_code'], partition_header['maker_code'], partition_header['game_name']))
+                partition_path = volume_path / ('Partition %d (%s) (%s%s) (%s)' % (partition_num, partition_type.capitalize(), partition_header['game_code'], partition_header['maker_code'], safename(partition_header['game_name'])))
                 yield (partition_path, None, None)
-                partition_header_path = partition_path / '__PARTITION_HEADER_NIEMAFS'
+                partition_header_path = partition_path / '___PARTITION_HEADER_NIEMAFS'
                 yield (partition_header_path, None, None)
                 yield (partition_header_path / 'ticket.bin', None, ticket)
                 yield (partition_header_path / 'tmd.bin', None, tmd)

setup.py

@@ -14,7 +14,7 @@
 here = path.abspath(path.dirname(__file__))
 setup(
     name='niemafs',  # Required
-    version='0.0.9',  # Required
+    version='0.0.10',  # Required
     description="Niema's Python library for reading data from various file system standards", # Required
     long_description="Niema's Python library for reading data from various file system standards", # Optional
     long_description_content_type='text/plain',  # Optional (see note above)