From fbedd79aa360662691885b9709dac46c8f092279 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Thu, 30 Nov 2017 22:10:45 -0800
Subject: [PATCH] updated scan_disks()

* Renamed from get_attached_disk_info()
---
 Scripts/functions/common.py      |  3 +-
 Scripts/functions/disk.py        | 52 ++++++++++++++------------------
 Scripts/functions/winpe_menus.py | 30 +++++++++++++++---
 3 files changed, 50 insertions(+), 35 deletions(-)

diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py
index 5204b45f..59826f0f 100644
--- a/Scripts/functions/common.py
+++ b/Scripts/functions/common.py
@@ -414,6 +414,7 @@ def try_and_print(message='Trying...',
         and the result string will be printed in the correct color.
     catch_all=False will result in unspecified exceptions being re-raised."""
     err = None
+    out = None
     w_exceptions = other_results.get('Warning', {}).keys()
     w_exceptions = tuple(get_exception(e) for e in w_exceptions)
     e_exceptions = other_results.get('Error', {}).keys()
@@ -449,7 +450,7 @@ def try_and_print(message='Trying...',
     if err and not catch_all:
         raise
     else:
-        return {'CS': not bool(err), 'Error': err}
+        return {'CS': not bool(err), 'Error': err, 'Out': out}
 
 def upload_data(path, file):
     """Add CLIENT_INFO_SERVER to authorized connections and upload file."""
diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py
index 4a265300..15343053 100644
--- a/Scripts/functions/disk.py
+++ b/Scripts/functions/disk.py
@@ -18,32 +18,6 @@ def assign_volume_letters():
     except subprocess.CalledProcessError:
         pass
 
-def get_attached_disk_info():
-    """Get details about the attached disks"""
-    disks = []
-    print_info('Getting drive info...')
-
-    # Get disks
-    disks = get_disks()
-
-    # Get disk details
-    for disk in disks:
-        # Get partition style
-        disk['Table'] = get_table_type(disk)
-
-        # Get disk name/model and physical details
-        disk.update(get_disk_details(disk))
-
-        # Get partition info for disk
-        disk['Partitions'] = get_partitions(disk)
-        
-        for par in disk['Partitions']:
-            # Get partition details
-            par.update(get_partition_details(disk, par))
-
-    # Done
-    return disks
-
 def get_boot_mode():
     boot_mode = 'Legacy'
     try:
@@ -318,10 +292,30 @@ def remove_volume_letters(keep=None):
     except subprocess.CalledProcessError:
         pass
 
-def select_disk(title='Which disk?'):
-    """Select a disk from the attached disks"""
-    disks = get_attached_disk_info()
+def scan_disks():
+    """Get details about the attached disks"""
+    disks = get_disks()
 
+    # Get disk details
+    for disk in disks:
+        # Get partition style
+        disk['Table'] = get_table_type(disk)
+
+        # Get disk name/model and physical details
+        disk.update(get_disk_details(disk))
+
+        # Get partition info for disk
+        disk['Partitions'] = get_partitions(disk)
+        
+        for par in disk['Partitions']:
+            # Get partition details
+            par.update(get_partition_details(disk, par))
+
+    # Done
+    return disks
+
+def select_disk(title='Which disk?', disks):
+    """Select a disk from the attached disks"""
     # Build menu
     disk_options = []
     for disk in disks:
diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index ceda3aa1..69b551fe 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/Scripts/functions/winpe_menus.py
@@ -74,13 +74,21 @@ def menu_backup():
     # Select destination
     destination = select_backup_destination()
 
+    # Scan disks
+    try_and_print(message='Assigning letters...', function=assign_volume_letters, other_results=other_results)
+    result = try_and_print(message='Getting drive info...', function=scan_disks, other_results=other_results)
+    if result['CS']:
+        disks = result['Out']
+    else:
+        print_error('ERROR: No disks found.')
+        raise GenericAbort
+    
     # Select disk to backup
-    assign_volume_letters()
-    disk = select_disk('For which drive are we creating backups?')
+    disk = select_disk('For which drive are we creating backups?', disks)
     if not disk:
         raise GenericAbort
     
-    # "Prep" disk?
+    # "Prep" disk
     prep_disk_for_backup(destination, disk, ticket_number)
 
     # Display details for backup task
@@ -205,9 +213,21 @@ def menu_setup():
     # Find Windows image
     windows_image = find_windows_image(bin, windows_version)
 
+    # Scan disks
+    try_and_print(message='Assigning letters...', function=assign_volume_letters, other_results=other_results)
+    result = try_and_print(message='Getting drive info...', function=scan_disks, other_results=other_results)
+    if result['CS']:
+        disks = result['Out']
+    else:
+        print_error('ERROR: No disks found.')
+        raise GenericAbort
+    
     # Select drive to use as the OS drive
-    assign_volume_letters()
-    dest_disk = select_disk('To which drive are we installing Windows?')
+    dest_disk = select_disk('To which drive are we installing Windows?', disks)
+    if not disk:
+        raise GenericAbort
+    
+    # "Prep" disk
     prep_disk_for_formatting(dest_disk)
 
     # Display details for setup task