Update backup/restore code

* Added support for local backups
* Added volume label detection for local backups
* Replace spaces in backup_prefix with underscores

.bin/Scripts/functions/backup.py

@@ -1,5 +1,7 @@
 # Wizard Kit: Functions - Backup
 
+import ctypes
+
 from functions.disk import *
 
 # Regex
@@ -30,7 +32,39 @@ def backup_partition(disk, par):
 
 def fix_path(path):
     """Replace invalid filename characters with underscores."""
-    return REGEX_BAD_PATH_NAMES.sub('_', path)
+    local_drive = path[1:2] == ':'
+    new_path = REGEX_BAD_PATH_NAMES.sub('_', path)
+    if local_drive:
+        new_path = '{}:{}'.format(new_path[0:1], new_path[2:])
+    return new_path
+
+def get_volume_display_name(mountpoint):
+    """Get display name from volume mountpoint and label, returns str."""
+    name = mountpoint
+    try:
+        kernel32 = ctypes.windll.kernel32
+        vol_name_buffer = ctypes.create_unicode_buffer(1024)
+        fs_name_buffer = ctypes.create_unicode_buffer(1024)
+        serial_number = None
+        max_component_length = None
+        file_system_flags = None
+        
+        vol_info = kernel32.GetVolumeInformationW(
+            ctypes.c_wchar_p(mountpoint),
+            vol_name_buffer,
+            ctypes.sizeof(vol_name_buffer),
+            serial_number,
+            max_component_length,
+            file_system_flags,
+            fs_name_buffer,
+            ctypes.sizeof(fs_name_buffer)
+            )
+        
+        name = '{} "{}"'.format(name, vol_name_buffer.value)
+    except:
+        pass
+    
+    return name
 
 def prep_disk_for_backup(destination, disk, backup_prefix):
     """Gather details about the disk and its partitions.
@@ -114,7 +148,20 @@ def select_backup_destination(auto_select=True):
     actions = [
         {'Name': 'Main Menu', 'Letter': 'M'},
         ]
-
+    
+    # Add local disks
+    for d in psutil.disk_partitions():
+        if re.search(r'^{}'.format(global_vars['Env']['SYSTEMDRIVE']), d.mountpoint, re.IGNORECASE):
+            # Skip current OS drive
+            pass
+        elif 'fixed' in d.opts:
+            # Skip DVD, etc
+            destinations.append({
+                'Name':     'Local Disk - {}'.format(
+                    get_volume_display_name(d.mountpoint)),
+                'Letter':   re.sub(r'^(\w):\\.*$', r'\1', d.mountpoint),
+                })
+    
     # Size check
     for dest in destinations:
         if 'IP' in dest:

.bin/Scripts/functions/winpe_menus.py

@@ -83,7 +83,14 @@ def menu_backup():
         backup_prefix = ticket_number
     else:
         backup_prefix = get_simple_string(prompt='Enter backup name prefix')
+        backup_prefix = backup_prefix.replace(' ', '_')
 
+    # Assign drive letters
+    try_and_print(
+        message = 'Assigning letters...',
+        function = assign_volume_letters,
+        other_results = other_results)
+    
     # Mount backup shares
     mount_backup_shares(read_write=True)
 
@@ -91,10 +98,6 @@ def menu_backup():
     destination = select_backup_destination(auto_select=False)
 
     # Scan disks
-    try_and_print(
-        message = 'Assigning letters...',
-        function = assign_volume_letters,
-        other_results = other_results)
     result = try_and_print(
         message = 'Getting disk info...',
         function = scan_disks,
@@ -111,7 +114,7 @@ def menu_backup():
         raise GenericAbort
     
     # "Prep" disk
-    prep_disk_for_backup(destination, disk, dest_prefix)
+    prep_disk_for_backup(destination, disk, backup_prefix)
 
     # Display details for backup task
     clear_screen()

.bin/Scripts/user_data_transfer.py

@@ -25,6 +25,7 @@ if __name__ == '__main__':
             backup_prefix = ticket_number
         else:
             backup_prefix = get_simple_string(prompt='Enter backup name prefix')
+            backup_prefix = backup_prefix.replace(' ', '_')
         
         # Set destination
         folder_path = r'{}\Transfer'.format(KIT_NAME_SHORT)