add unsplit function to script

splitNSP.py

@@ -4,6 +4,7 @@
 # Purpose: Splits Nintendo Switch NSP files into parts for installation on FAT32
 
 import os
+import fnmatch
 import argparse
 import shutil
 from datetime import datetime
@@ -25,7 +26,7 @@ def splitQuick(filepath):
         return
 
     print('Splitting NSP into {0} parts...\n'.format(splitNum + 1))
-    
+
     # Create directory, delete if already exists
     dir = filepath[:-4] + '_split.nsp'
     if os.path.exists(dir):
@@ -68,12 +69,12 @@ def splitQuick(filepath):
             nspFile.seek(splitSize * -1, os.SEEK_END)
             nspFile.truncate()
             print('Part {:02} complete'.format(splitNum - (i + 1)))
-    
+
     # Print assurance statement for user
     print('Starting part 00\nPart 00 complete')
 
     print('\nNSP successfully split!\n')
-    
+
 def splitCopy(filepath, output_dir=""):
     fileSize = os.path.getsize(filepath)
     info = shutil.disk_usage(os.path.dirname(os.path.abspath(filepath)))
@@ -85,7 +86,7 @@ def splitCopy(filepath, output_dir=""):
     if splitNum == 0:
         print('This NSP is under 4GiB and does not need to be split.\n')
         return
-    
+
     print('Splitting NSP into {0} parts...\n'.format(splitNum + 1))
 
     # Create directory, delete if already exists
@@ -107,7 +108,7 @@ def splitCopy(filepath, output_dir=""):
             partSize = 0
             print('Starting part {:02}'.format(i))
             outFile = os.path.join(dir, '{:02}'.format(i))
-            with open(outFile, 'wb') as splitFile: 
+            with open(outFile, 'wb') as splitFile:
                 if remainingSize > splitSize:
                     while partSize < splitSize:
                         splitFile.write(nspFile.read(chunkSize))
@@ -120,30 +121,62 @@ def splitCopy(filepath, output_dir=""):
             print('Part {:02} complete'.format(i))
     print('\nNSP successfully split!\n')
 
+def undoSplit(splitDir):
+    if not os.path.exists(splitDir):
+        print('{0}: No such directory exists'.format(splitDir))
+        return
+    totalFileSize = sum(os.path.getsize(os.path.join(splitDir,f)) for f in os.listdir(splitDir) if os.path.isfile(os.path.join(splitDir,f)))
+
+    info = shutil.disk_usage(os.path.dirname(os.path.abspath(splitDir)))
+    if info.free < totalFileSize:
+        print('Not enough disk space. Needs {0}GiB of free space\n'.format(totalFileSize))
+        return
+    filename = splitDir.rstrip("_split.nsp")+".nsp"
+    if os.path.exists(filename) and os.path.getsize(filename) == totalFileSize:
+        print('NSP seems to be already recovered. No need to unsplit')
+        return
+    # Re-combine split files
+    with open(filename, 'w+b') as nspFile:
+        for f in sorted(fnmatch.filter(os.listdir(splitDir), "[0,9]*")):
+            partSize = 0
+            with open(os.path.join(splitDir, f), 'r+b') as splitFile:
+                fileSize = os.path.getsize(os.path.join(splitDir, f))
+                while partSize < fileSize:
+                    nspFile.write(splitFile.read(chunkSize))
+                    partSize += chunkSize
+
+    # Print assurance statement for user
+    print('Combined splitted NSP from directory {0} into {1}...\n'.format(splitDir, filename))
+
+    print('\nNSP successfully unsplit!\n')
+
 def main():
     print('\n========== NSP Splitter ==========\n')
 
     # Arg parser for program options
     parser = argparse.ArgumentParser(description='Split NSP files into FAT32 compatible sizes')
-    parser.add_argument('filepath', help='Path to NSP file')
+    parser.add_argument('filepath', help='Path to NSP file (or NSP split directory in case of `-u\')')
     group = parser.add_mutually_exclusive_group()
     group.add_argument('-q', '--quick', action='store_true', help='Splits file in-place without creating a copy. Only requires 4GiB free space to run')
     group.add_argument('-o', '--output-dir', type=str, default="",
                         help="Set alternative output dir")
+    group.add_argument('-u', '--undo', action='store_true', help='Undo previous split')
 
     # Check passed arguments
     args = parser.parse_args()
 
     filepath = args.filepath
 
     # Check if required files exist
-    if os.path.isfile(filepath) == False:
+    if not args.undo and os.path.isfile(filepath) == False:
         print('NSP cannot be found\n')
         return 1
-    
+
     # Split NSP file
     if args.quick:
         splitQuick(filepath)
+    elif args.undo:
+        undoSplit(filepath)
     else:
         splitCopy(filepath, args.output_dir)