Volume letter updates

* Added reassign_letter()
  * Attempts to reassign a volume to better ensure predictable letters
  * i.e. Local Windows source volume letter == 'I'
* Adjusted code to avoid the "hidden" assign_volume_letters() call in select_disk()

Scripts/functions/disk.py

@@ -4,6 +4,7 @@ from functions.common import *
 import partition_uids
 
 def assign_volume_letters():
+    remove_volume_letters()
     try:
         # Run script
         with open(DISKPART_SCRIPT, 'w') as script:
@@ -19,9 +20,6 @@ def get_attached_disk_info():
     disks = []
     print_info('Getting drive info...')
 
-    # Assign all the letters
-    assign_volume_letters()
-
     # Get disks
     disks = get_disks()
 
@@ -286,8 +284,24 @@ def prep_disk_for_formatting(disk=None):
                 q='"' if par['Name'] != '' else '',
                 **par)
 
-def remove_volume_letters(keep=''):
-    if keep is None:
+def reassign_volume_letter(letter, new_letter='I'):
+    if not letter:
+        # Ignore
+        return None
+    try:
+        # Run script
+        with open(DISKPART_SCRIPT, 'w') as script:
+            script.write('select volume {}\n'.format(letter))
+            script.write('remove noerr\n')
+            script.write('assign letter={}\n'.format(new_letter))
+        run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
+    except subprocess.CalledProcessError:
+        pass
+    else:
+        return new_letter
+
+def remove_volume_letters(keep=None):
+    if not keep:
         keep = ''
     try:
         # Run script

Scripts/functions/winpe_menus.py

@@ -75,6 +75,7 @@ def menu_backup():
     destination = select_backup_destination()
 
     # Select disk to backup
+    assign_volume_letters()
     disk = select_disk('For which drive are we creating backups?')
     if not disk:
         raise GenericAbort
@@ -200,16 +201,15 @@ def menu_setup():
 
     # Select the version of Windows to apply
     windows_version = select_windows_version()
-
-    # Select drive to use as the OS drive
-    dest_disk = select_disk('To which drive are we installing Windows?')
-    prep_disk_for_formatting(dest_disk)
     
     # Find Windows image
-    ## NOTE: Needs to happen AFTER select_disk() is called as there's a hidden assign_volume_letters().
-    ##       This changes the current letters thus preventing installing from a local source.
     windows_image = find_windows_image(bin, windows_version)
 
+    # Select drive to use as the OS drive
+    assign_volume_letters()
+    dest_disk = select_disk('To which drive are we installing Windows?')
+    prep_disk_for_formatting(dest_disk)
+
     # Display details for setup task
     os.system('cls')
     print('Setup Windows - Details:\n')
@@ -235,6 +235,9 @@ def menu_setup():
 
     # Release currently used volume letters (ensures that the drives will get S, T, & W as needed below)
     remove_volume_letters(keep=windows_image['Source'])
+    new_letter = reassign_volume_letter(letter=windows_image['Source'])
+    if new_letter:
+        windows_image['Source'] = new_letter
 
     # Format and partition drive
     print('\n    Formatting Drive...     \t\t', end='', flush=True)