custominstall: fix incorrect return value causing TypeError in ci-gui

Prevents the issue of one having a corrupt header or something causing an issue in the middle of writing.

ci-gui.py

@@ -37,7 +37,7 @@ if is_windows:
     except ModuleNotFoundError:
         pass
 
-file_parent = dirname(__file__)
+file_parent = dirname(abspath(__file__))
 
 # automatically load boot9 if it's in the current directory
 b9_paths.insert(0, join(file_parent, 'boot9.bin'))
@@ -191,6 +191,18 @@ class CustomInstallGUI(ttk.Frame):
         add_cias = ttk.Button(listbox_buttons, text='Add CIAs', command=add_cias_callback)
         add_cias.grid(row=0, column=0)
 
+        def add_cias_callback():
+            d = fd.askdirectory(parent=parent, title='Select folder containing title contents in CDN format',
+                                initialdir=file_parent)
+            if d:
+                if isfile(join(d, 'tmd')):
+                    self.add_cia(d)
+                else:
+                    self.show_error('tmd file not found in the CDN directory:\n' + d)
+
+        add_cias = ttk.Button(listbox_buttons, text='Add CDN title folder', command=add_cias_callback)
+        add_cias.grid(row=0, column=1)
+
         def add_dirs_callback():
             d = fd.askdirectory(parent=parent, title='Select folder containing CIA files', initialdir=file_parent)
             if d:
@@ -199,7 +211,7 @@ class CustomInstallGUI(ttk.Frame):
                         self.add_cia(f.path)
 
         add_dirs = ttk.Button(listbox_buttons, text='Add folder', command=add_dirs_callback)
-        add_dirs.grid(row=0, column=1)
+        add_dirs.grid(row=0, column=2)
 
         def remove_selected_callback():
             indexes = self.cia_listbox.curselection()
@@ -209,7 +221,7 @@ class CustomInstallGUI(ttk.Frame):
                 n += 1
 
         remove_selected = ttk.Button(listbox_buttons, text='Remove selected', command=remove_selected_callback)
-        remove_selected.grid(row=0, column=2)
+        remove_selected.grid(row=0, column=3)
 
         # ---------------------------------------------------------------- #
         # create listbox
@@ -355,7 +367,6 @@ class CustomInstallGUI(ttk.Frame):
         installer = CustomInstall(boot9=boot9,
                                   seeddb=seeddb,
                                   movable=movable_sed,
-                                  cias=cias,
                                   sd=sd_root,
                                   skip_contents=self.skip_contents_var.get() == 1,
                                   overwrite_saves=self.overwrite_saves_var.get() == 1)
@@ -393,12 +404,21 @@ class CustomInstallGUI(ttk.Frame):
         installer.event.on_error += ci_on_error
         installer.event.on_cia_start += ci_on_cia_start
 
+        try:
+            installer.prepare_titles(cias)
+        except Exception as e:
+            for line in format_exception(*exc_info()):
+                for line2 in line.split('\n')[:-1]:
+                    installer.log(line2)
+            self.show_error('An error occurred when trying to read the files.')
+            self.open_console()
+
         if taskbar:
             taskbar.SetProgressState(self.hwnd, tbl.TBPF_NORMAL)
 
         def install():
             try:
-                result, copied_3dsx = installer.start(continue_on_fail=False)
+                result, copied_3dsx = installer.start()
                 if result is True:
                     self.log('Done!')
                     if copied_3dsx:

custominstall.py

@@ -6,12 +6,14 @@
 
 from argparse import ArgumentParser
 from os import makedirs, scandir
-from os.path import dirname, join, isfile
+from os.path import dirname, join, isdir, isfile
 from random import randint
 from hashlib import sha256
 from locale import getpreferredencoding
+from pprint import pformat
 from shutil import copyfile
 import sys
+from sys import platform, executable
 from tempfile import TemporaryDirectory
 from traceback import format_exception
 from typing import BinaryIO, TYPE_CHECKING
@@ -19,22 +21,26 @@ import subprocess
 
 if TYPE_CHECKING:
     from os import PathLike
-    from typing import Union
+    from typing import List, Union
 
 from events import Events
 
 from pyctr.crypto import CryptoEngine, Keyslot, load_seeddb, get_seed
-from pyctr.type.cia import CIAReader, CIASection
+from pyctr.type.cdn import CDNReader
+from pyctr.type.cia import CIAReader, CIAError
 from pyctr.type.ncch import NCCHSection
 from pyctr.util import roundup
 
-is_windows = sys.platform == 'win32'
+is_windows = platform == 'win32'
+
+if platform == 'msys':
+    platform = 'win32'
 
 # used to run the save3ds_fuse binary next to the script
 frozen = getattr(sys, 'frozen', False)
 script_dir: str
 if frozen:
-    script_dir = dirname(sys.executable)
+    script_dir = dirname(executable)
 else:
     script_dir = dirname(__file__)
 
@@ -141,10 +147,7 @@ def save_cifinish(path: 'Union[PathLike, bytes, str]', data: dict):
 
 
 class CustomInstall:
-
-    cia: CIAReader
-
-    def __init__(self, boot9, seeddb, movable, cias, sd, cifinish_out=None,
+    def __init__(self, boot9, seeddb, movable, sd, cifinish_out=None,
                  overwrite_saves=False, skip_contents=False):
         self.event = Events()
         self.log_lines = []  # Stores all info messages for user to view
@@ -152,7 +155,7 @@ class CustomInstall:
         self.crypto = CryptoEngine(boot9=boot9)
         self.crypto.setup_sd_key_from_file(movable)
         self.seeddb = seeddb
-        self.cias = cias
+        self.readers: 'List[Union[CDNReader, CIAReader]]' = []
         self.sd = sd
         self.skip_contents = skip_contents
         self.overwrite_saves = overwrite_saves
@@ -170,12 +173,30 @@ class CustomInstall:
             total_read = size - left
             if fire_event:
                 self.event.update_percentage((total_read / size) * 100, total_read / 1048576, size / 1048576)
-    
-    def start(self, continue_on_fail=True):
+
+    def prepare_titles(self, paths: 'List[PathLike]'):
+        readers = []
+        for path in paths:
+            self.log(f'Reading {path}')
+            if isdir(path):
+                # try the default tmd file
+                reader = CDNReader(join(path, 'tmd'))
+            else:
+                try:
+                    reader = CIAReader(path)
+                except CIAError:
+                    # if there was an error with parsing the CIA header,
+                    # the file would be tried in CDNReader next (assuming it's a tmd)
+                    # any other error should be propagated to the caller
+                    reader = CDNReader(path)
+            readers.append(reader)
+        self.readers = readers
+
+    def start(self):
         if frozen:
             save3ds_fuse_path = join(script_dir, 'bin', 'save3ds_fuse')
         else:
-            save3ds_fuse_path = join(script_dir, 'bin', sys.platform, 'save3ds_fuse')
+            save3ds_fuse_path = join(script_dir, 'bin', platform, 'save3ds_fuse')
         if is_windows:
             save3ds_fuse_path += '.exe'
         if not isfile(save3ds_fuse_path):
@@ -204,22 +225,10 @@ class CustomInstall:
 
         # Now loop through all provided cia files
         
-        for idx, c in enumerate(self.cias):
-            self.log('Reading ' + c)
-
-            try:
-                cia = CIAReader(c)
-            except Exception as e:
-                self.event.on_error(sys.exc_info())
-                if continue_on_fail:
-                    continue
-                else:
-                    return None, False
+        for idx, cia in enumerate(self.readers):
 
             self.event.on_cia_start(idx)
 
-            self.cia = cia
-            
             tid_parts = (cia.tmd.title_id[0:8], cia.tmd.title_id[8:16])
 
             try:
@@ -282,10 +291,9 @@ class CustomInstall:
                 # write the tmd
                 enc_path = content_root_cmd + '/' + tmd_filename
                 self.log(f'Writing {enc_path}...')
-                with cia.open_raw_section(CIASection.TitleMetadata) as s:
-                    with open(join(content_root, tmd_filename), 'wb') as o:
-                        self.copy_with_progress(s, o, cia.sections[CIASection.TitleMetadata].size, enc_path,
-                                                fire_event=False)
+                with open(join(content_root, tmd_filename), 'wb') as o:
+                    with self.crypto.create_ctr_io(Keyslot.SD, o, self.crypto.sd_path_to_iv(enc_path)) as e:
+                        e.write(bytes(cia.tmd))
 
                 # write each content
                 for co in cia.content_info:
@@ -439,7 +447,10 @@ class CustomInstall:
                 if out.returncode:
                     for l in out.stdout.split('\n'):
                         self.log(l)
-                    return False
+                    self.log('Command line:')
+                    for l in pformat(out.args).split('\n'):
+                        self.log(l)
+                    return False, False
 
                 for title_id, entry in title_info_entries.items():
                     # write the title info entry to the temp directory
@@ -456,6 +467,9 @@ class CustomInstall:
                 if out.returncode:
                     for l in out.stdout.split('\n'):
                         self.log(l)
+                    self.log('Command line:')
+                    for l in pformat(out.args).split('\n'):
+                        self.log(l)
                     return False, False
 
             finalize_3dsx_orig_path = join(script_dir, 'custom-install-finalize.3dsx')
@@ -536,7 +550,6 @@ if __name__ == "__main__":
 
     installer = CustomInstall(boot9=args.boot9,
                               seeddb=args.seeddb,
-                              cias=args.cia,
                               movable=args.movable,
                               sd=args.sd,
                               overwrite_saves=args.overwrite_saves,
@@ -558,7 +571,9 @@ if __name__ == "__main__":
     installer.event.update_percentage += percent_handle
     installer.event.on_error += error
 
-    result, copied_3dsx = installer.start(continue_on_fail=False)
+    installer.prepare_titles(args.cia)
+
+    result, copied_3dsx = installer.start()
     if result is False:
         # save3ds_fuse failed
         installer.log('NOTE: Once save3ds_fuse is fixed, run the same command again with --skip-contents')

requirements.txt

@@ -1,3 +1,3 @@
 pycryptodomex==3.9.8
 events==0.3
-pyctr==0.4.1
+pyctr==0.4.3