From 70e4797b7ae3d690f3a79c755790b266de80e696 Mon Sep 17 00:00:00 2001
From: 2Shirt <1923621+2Shirt@users.noreply.github.com>
Date: Mon, 21 May 2018 15:43:05 -0600
Subject: [PATCH] Refactored windows_image code

* Better names (e.g. 'Letter' instead of 'Source')
* Added 'Local' bool to simplify network-vs-local installs
* Only reassign the source letter if necessary
* Bugfix: Update 'Path' when reassigning a volume letter
* Fixes #33 (again)
---
 .bin/Scripts/functions/windows_setup.py |  6 ++++--
 .bin/Scripts/functions/winpe_menus.py   | 20 ++++++++++++++++----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/.bin/Scripts/functions/windows_setup.py b/.bin/Scripts/functions/windows_setup.py
index bf21069a..cd7f8444 100644
--- a/.bin/Scripts/functions/windows_setup.py
+++ b/.bin/Scripts/functions/windows_setup.py
@@ -50,7 +50,8 @@ def find_windows_image(windows_version):
             path = '{}images\{}.{}'.format(d.mountpoint, imagefile, ext)
             if os.path.isfile(path) and wim_contains_image(path, imagename):
                 image['Path'] = path
-                image['Source'] = d.mountpoint[:1]
+                image['Letter'] = d.mountpoint[:1].upper()
+                image['Local'] = True
                 if ext == 'swm':
                     image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4])
                 break
@@ -68,7 +69,8 @@ def find_windows_image(windows_version):
                     ext)
                 if os.path.isfile(path) and wim_contains_image(path, imagename):
                     image['Path'] = path
-                    image['Source'] = None
+                    image['Letter'] = None
+                    image['Local'] = False
                     if ext == 'swm':
                         image['Glob'] = '--ref="{}*.swm"'.format(
                             image['Path'][:-4])
diff --git a/.bin/Scripts/functions/winpe_menus.py b/.bin/Scripts/functions/winpe_menus.py
index 6d8d10bc..aed9839d 100644
--- a/.bin/Scripts/functions/winpe_menus.py
+++ b/.bin/Scripts/functions/winpe_menus.py
@@ -337,10 +337,22 @@ def menu_setup():
         raise GenericAbort
 
     # Remove volume letters so S, T, & W can be used 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
+    try_and_print(
+        message = 'Removing volume letters...',
+        function = remove_volume_letters,
+        other_results = other_results,
+        keep=windows_image['Letter'])
+   
+    # Assign new letter for local source if necessary
+    if windows_image['Local'] and windows_image['Letter'] in ['S', 'T', 'W']:
+        new_letter = try_and_print(
+            message = 'Reassigning source volume letter...',
+            function = reassign_volume_letter,
+            other_results = other_results,
+            letter=windows_image['Letter'])
+        windows_image['Path'] = '{}{}'.format(
+            new_letter, windows_image['Path'][1:])
+        windows_image['Letter'] = new_letter
 
     # Format and partition disk
     result = try_and_print(