WizardKit/Scripts/functions/backup.py
2017-11-30 17:17:00 -08:00

149 lines
6.4 KiB
Python

# Wizard Kit PE: Functions - Backup
from functions.common import *
def backup_partition(disk, partition):
if par['Image Exists'] or par['Number'] in disk['Bad Partitions']:
raise GenericAbort
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'
disk['Clobber Risk'] = []
width = len(str(len(disk['Partitions'])))
# Bail early
if dest is None:
raise Exception('Destination not provided.')
if disk is None:
raise Exception('Disk not provided.')
if ticket_id is None:
raise Exception('Ticket ID not provided.')
# Get partition totals
disk['Bad Partitions'] = [par['Number'] for par in disk['Partitions'] if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE)]
disk['Valid Partitions'] = len(disk['Partitions']) - len(disk['Bad Partitions'])
# Bail if no valid partitions are found (those that can be imaged)
if disk['Valid Partitions'] <= 0:
abort_to_main_menu(' No partitions can be imaged for the selected drive')
# Prep partitions
for par in disk['Partitions']:
if par['Number'] in disk['Bad Partitions']:
par['Display String'] = '{YELLOW} * Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS}){CLEAR}'.format(
width=width,
q='"' if par['Name'] != '' else '',
**par,
**COLORS)
else:
# Update info for WIM capturing
par['Image Name'] = str(par['Name'])
if par['Image Name'] == '':
par['Image Name'] = 'Unknown'
if 'IP' in dest:
par['Image Path'] = '\\\\{IP}\\{Share}\\{ticket}'.format(ticket=ticket_id, **dest)
else:
par['Image Path'] = '{Letter}:\\{ticket}'.format(ticket=ticket_id, **dest)
par['Image File'] = '{Number}_{Image Name}'.format(**par)
par['Image File'] = '{fixed_name}.wim'.format(fixed_name=re.sub(r'\W', '_', par['Image File']))
# Check for existing backups
par['Image Exists'] = False
if os.path.exists('{Image Path}\\{Image File}'.format(**par)):
par['Image Exists'] = True
disk['Clobber Risk'].append(par['Number'])
par['Display String'] = '{BLUE} + '.format(**COLORS)
else:
par['Display String'] = '{CLEAR} '.format(**COLORS)
# Append rest of Display String for valid/clobber partitions
par['Display String'] += 'Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}{CLEAR}'.format(
width=width,
q='"' if par['Name'] != '' else '',
**par,
**COLORS)
# Set description for bad partitions
if len(disk['Bad Partitions']) > 1:
disk['Backup Warnings'] += '{YELLOW} * Unable to backup these partitions{CLEAR}\n'.format(**COLORS)
elif len(disk['Bad Partitions']) == 1:
print_warning(' * Unable to backup this partition')
disk['Backup Warnings'] += '{YELLOW} * Unable to backup this partition{CLEAR}\n'.format(**COLORS)
# Set description for partitions that would be clobbered
if len(disk['Clobber Risk']) > 1:
disk['Backup Warnings'] += '{BLUE} + These partitions already have backup images on {Name}{CLEAR}\n'.format(**dest, **COLORS)
elif len(disk['Clobber Risk']) == 1:
disk['Backup Warnings'] += '{BLUE} + This partition already has a backup image on {Name}{CLEAR}\n'.format(**dest, **COLORS)
# Set warning for skipped partitions
if len(disk['Clobber Risk']) + len(disk['Bad Partitions']) > 1:
disk['Backup Warnings'] += '\n{YELLOW}If you continue the partitions marked above will NOT be backed up.{CLEAR}\n'.format(**COLORS)
if len(disk['Clobber Risk']) + len(disk['Bad Partitions']) == 1:
disk['Backup Warnings'] += '\n{YELLOW}If you continue the partition marked above will NOT be backed up.{CLEAR}\n'.format(**COLORS)
def select_backup_destination():
# Build menu
dests = []
for server in BACKUP_SERVERS:
if server['Mounted']:
dests.append(server)
actions = [
{'Name': 'Main Menu', 'Letter': 'M'},
]
# Size check
for dest in dests:
if 'IP' in dest:
dest['Usage'] = shutil.disk_usage('\\\\{IP}\\{Share}'.format(**dest))
else:
dest['Usage'] = shutil.disk_usage('{Letter}:\\'.format(**dest))
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:
print_warning('No backup destinations found.')
return None
def verify_wim_backup(bin=None, par=None):
# Bail early
if bin is None:
raise Exception('bin path not specified.')
if par is None:
raise Exception('Partition not specified.')
# Verify hiding all output for quicker verification
print(' Partition {Number} Image...\t\t'.format(**par), end='', flush=True)
cmd = '{bin}\\wimlib\\wimlib-imagex verify "{Image Path}\\{Image File}" --nocheck'.format(bin=bin, **par)
if not os.path.exists('{Image Path}\\{Image File}'.format(**par)):
print_error('Missing.')
else:
try:
run_program(cmd)
print_success('OK.')
except subprocess.CalledProcessError as err:
print_error('Damaged.')
par['Error'] = par.get('Error', []) + err.stderr.decode().splitlines()
raise BackupError
if __name__ == '__main__':
print("This file is not meant to be called directly.")