From 7133089d31ee468e3bc1b1357f0d3ea66c1af7df Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Thu, 30 Nov 2017 20:03:19 -0800
Subject: [PATCH] updated select_backup_destination()

---
 Scripts/functions/backup.py      | 36 +++++++++++++++++---------------
 Scripts/functions/winpe_menus.py |  2 --
 Scripts/winpe_root_menu.py       |  3 ++-
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py
index 62e45d1..4ab89a4 100644
--- a/Scripts/functions/backup.py
+++ b/Scripts/functions/backup.py
@@ -105,35 +105,37 @@ def prep_disk_for_backup(destination, disk, ticket_number):
             COLORS['YELLOW'], COLORS['CLEAR'])
     disk['Backup Warnings'] = warnings
 
-def select_backup_destination():
+def select_backup_destination(auto_select=True):
     # Build menu
-    dests = []
-    for server in BACKUP_SERVERS:
-         if server['Mounted']:
-            dests.append(server)
+    destinations = [s for s in BACKUP_SERVERS if s['Mounted']]
     actions = [
         {'Name': 'Main Menu', 'Letter': 'M'},
         ]
 
     # Size check
-    for dest in dests:
+    for dest in destinations:
         if 'IP' in dest:
-            dest['Usage'] = shutil.disk_usage('\\\\{IP}\\{Share}'.format(**dest))
+            dest['Usage'] = shutil.disk_usage(r'\\{IP}\{Share}'.format(**dest))
         else:
-            dest['Usage'] = shutil.disk_usage('{Letter}:\\'.format(**dest))
+            dest['Usage'] = shutil.disk_usage('{}:\\'.format(dest['Letter']))
         dest['Free Space'] = human_readable_size(dest['Usage'].free)
         dest['Display Name'] = '{Name} ({Free Space} available)'.format(**dest)
 
-    # Show menu or bail
-    if len(dests) > 0:
-        selection = menu_select('Where are we backing up to?', dests, actions)
-        if selection == 'M':
-            return None
-        else:
-            return dests[int(selection)-1]
-    else:
+    # Bail
+    if not destinations:
         print_warning('No backup destinations found.')
-        return None
+        raise GenericAbort
+    
+    # Skip menu?
+    if len(destinations) == 1 and auto_select:
+        return destinations[0]
+    
+    selection = menu_select(
+        'Where are we backing up to?', destinations, actions)
+    if selection == 'M':
+        raise GenericAbort
+    else:
+        return destinations[int(selection)-1]
 
 def verify_wim_backup(bin=None, par=None):
     # Bail early
diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index ddbcc31..2f40532 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/Scripts/functions/winpe_menus.py
@@ -71,8 +71,6 @@ def menu_backup():
 
     # Select destination
     destination = select_backup_destination()
-    if not destination:
-        raise GenericAbort
 
     # Select disk to backup
     disk = select_disk('For which drive are we creating backups?')
diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py
index 9ff1508..ad19b34 100644
--- a/Scripts/winpe_root_menu.py
+++ b/Scripts/winpe_root_menu.py
@@ -15,7 +15,8 @@ if __name__ == '__main__':
     try:
         menu_root()
     except GenericAbort:
-        pause('Press Enter to return to main menu... ')
+        # pause('Press Enter to return to main menu... ')
+        pass
     except SystemExit:
         pass
     except: