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

---
 Scripts/functions/backup.py      | 38 ++++++++++++--------------------
 Scripts/functions/common.py      |  7 ++++--
 Scripts/functions/winpe_menus.py | 16 +++++++++++---
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py
index e33004ba..f91c692f 100644
--- a/Scripts/functions/backup.py
+++ b/Scripts/functions/backup.py
@@ -2,31 +2,21 @@
 
 from functions.common import *
 
-def backup_partition(bin=None, disk=None, par=None):
-    # Bail early
-    if bin is None:
-        raise Exception('bin path not specified.')
-    if disk is None:
-        raise Exception('Disk not specified.')
-    if par is None:
-        raise Exception('Partition not specified.')
+def backup_partition(disk, partition):
+    if par['Image Exists'] or par['Number'] in disk['Bad Partitions']:
+        raise GenericAbort
     
-    print('    Partition {Number} Backup...\t\t'.format(**par), end='', flush=True)
-    if par['Number'] in disk['Bad Partitions']:
-        print_warning('Skipped.')
-    else:
-        cmd = '{bin}\\wimlib\\wimlib-imagex capture {Letter}:\\ "{Image Path}\\{Image File}" "{Image Name}" "{Image Name}" --compress=none'.format(bin=bin, **par)
-        if par['Image Exists']:
-            print_warning('Skipped.')
-        else:
-            try:
-                os.makedirs('{Image Path}'.format(**par), exist_ok=True)
-                run_program(cmd)
-                print_success('Complete.')
-            except subprocess.CalledProcessError as err:
-                print_error('Failed.')
-                par['Error'] = err.stderr.decode().splitlines()
-                raise BackupError
+    cmd = [
+        global_vars['Tools']['wimlib-imagex'],
+        'capture'
+        '{}:\\'.format(par['Letter']),
+        r'{}\{}'.format(par['Image Path'], par['Image File']),
+        par['Image Name'], # Image name
+        par['Image Name'], # Image description
+        ' --compress=none',
+        ]
+    os.makedirs(par['Image Path'], exist_ok=True)
+    run_program(cmd)
 
 def prep_disk_for_backup(dest=None, disk=None, ticket_id=None):
     disk['Backup Warnings'] = '\n'
diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py
index fcb2dd51..4cf147e7 100644
--- a/Scripts/functions/common.py
+++ b/Scripts/functions/common.py
@@ -37,6 +37,9 @@ class BIOSKeyNotFoundError(Exception):
 class BinNotFoundError(Exception):
     pass
 
+class GenericAbort(Exception):
+    pass
+
 class GenericError(Exception):
     pass
 
@@ -52,7 +55,7 @@ class NotInstalledError(Exception):
 class NoProfilesError(Exception):
     pass
 
-class PathNotFoundException(Exception):
+class PathNotFoundError(Exception):
     pass
 
 class UnsupportedOSError(Exception):
@@ -432,7 +435,7 @@ def try_and_print(message='Trying...',
         err = traceback.format_exc()
 
     # Return or raise?
-    if bool(err) and not catch_all:
+    if err and not catch_all:
         raise
     else:
         return {'CS': not bool(err), 'Error': err}
diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index b03802b4..2056b4fd 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/Scripts/functions/winpe_menus.py
@@ -53,6 +53,14 @@ PE_TOOLS = {
 def menu_backup():
     """Take backup images of partition(s) in the WIM format and save them to a backup share"""
     errors = False
+    other_results = {
+        'Error': {
+            'CalledProcessError':   'Unknown Error',
+        },
+        'Warning': {
+            'GenericAbort':         'Skipped',
+            'GenericRepair':        'Repaired',
+        }}
 
     # Set ticket ID
     os.system('cls')
@@ -89,10 +97,12 @@ def menu_backup():
     # Backup partition(s)
     print('\n\nStarting task.\n')
     for par in disk['Partitions']:
-        try:
-            backup_partition(bin, disk, par)
-        except BackupError:
+        message = 'Partition {} Backup...'.format(par['Number'])
+        result = try_and_print(message=message, function=backup_partition,
+            other_results=other_results, disk=disk, partition=par)
+        if not result['CS']:
             errors = True
+            par['Error'] = result['Error']
     
     # Verify backup(s)
     if disk['Valid Partitions'] > 1: