From 3114b72178cf821234ebe5becb8792d10c30001e Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 17:57:45 -0600 Subject: [PATCH 1/7] Add ticket toggle var and get_simple_string func * get_simple_string() will be used for backup prefixes --- .bin/Scripts/functions/common.py | 11 +++++++++++ .bin/Scripts/settings/main.py | 1 + 2 files changed, 12 insertions(+) diff --git a/.bin/Scripts/functions/common.py b/.bin/Scripts/functions/common.py index 050205e7..ed1d1938 100644 --- a/.bin/Scripts/functions/common.py +++ b/.bin/Scripts/functions/common.py @@ -196,6 +196,8 @@ def extract_item(item, filter='', silent=False): def get_ticket_number(): """Get TicketNumber from user, save in LogDir, and return as str.""" + if not ENABLED_TICKET_NUMBERS: + return None ticket_number = None while ticket_number is None: _input = input('Enter ticket number: ') @@ -208,6 +210,15 @@ def get_ticket_number(): f.write(ticket_number) return ticket_number +def get_simple_string(prompt='Enter string'): + """Get string from user (only alphanumeric/space chars) and return as str.""" + simple_string = None + while simple_string is None: + _input = input('{}: '.format(prompt)) + if re.match(r'^(\w|-| )+$', _input, re.ASCII): + simple_string = _input.strip() + return simple_string + def human_readable_size(size, decimals=0): """Convert size in bytes to a human-readable format and return a str.""" # Prep string formatting diff --git a/.bin/Scripts/settings/main.py b/.bin/Scripts/settings/main.py index 86b0fcee..c49da96e 100644 --- a/.bin/Scripts/settings/main.py +++ b/.bin/Scripts/settings/main.py @@ -2,6 +2,7 @@ # Features ENABLED_UPLOAD_DATA = False +ENABLED_TICKET_NUMBERS = False # STATIC VARIABLES (also used by BASH and BATCH files) ## NOTE: There are no spaces around the = for easier parsing in BASH and BATCH From bec947262b8bd9147317a002d78dea588b9f4786 Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 18:03:53 -0600 Subject: [PATCH 2/7] Updated WinPE sections --- .bin/Scripts/functions/backup.py | 6 +++--- .bin/Scripts/functions/winpe_menus.py | 14 ++++++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/.bin/Scripts/functions/backup.py b/.bin/Scripts/functions/backup.py index fe0935fb..a787a0ed 100644 --- a/.bin/Scripts/functions/backup.py +++ b/.bin/Scripts/functions/backup.py @@ -32,7 +32,7 @@ def fix_path(path): """Replace invalid filename characters with underscores.""" return REGEX_BAD_PATH_NAMES.sub('_', path) -def prep_disk_for_backup(destination, disk, ticket_number): +def prep_disk_for_backup(destination, disk, backup_prefix): """Gather details about the disk and its partitions. This includes partitions that can't be backed up, @@ -71,10 +71,10 @@ def prep_disk_for_backup(destination, disk, ticket_number): par['Image Name'] = par['Name'] if par['Name'] else 'Unknown' if 'IP' in destination: par['Image Path'] = r'\\{}\{}\{}'.format( - destination['IP'], destination['Share'], ticket_number) + destination['IP'], destination['Share'], backup_prefix) else: par['Image Path'] = r'{}:\{}'.format( - ticket_number, destination['Letter']) + destination['Letter'], backup_prefix) par['Image Path'] += r'\{}_{}.wim'.format( par['Number'], par['Image Name']) par['Image Path'] = fix_path(par['Image Path']) diff --git a/.bin/Scripts/functions/winpe_menus.py b/.bin/Scripts/functions/winpe_menus.py index 750f88a0..c5d6dd99 100644 --- a/.bin/Scripts/functions/winpe_menus.py +++ b/.bin/Scripts/functions/winpe_menus.py @@ -75,10 +75,14 @@ def menu_backup(): }} set_title('{}: Backup Menu'.format(KIT_NAME_FULL)) - # Set ticket Number + # Set backup prefix clear_screen() print_standard('{}\n'.format(global_vars['Title'])) ticket_number = get_ticket_number() + if ENABLED_TICKET_NUMBERS: + backup_prefix = ticket_number + else: + backup_prefix = get_simple_string(prompt='Enter backup name prefix') # Mount backup shares mount_backup_shares(read_write=True) @@ -107,12 +111,13 @@ def menu_backup(): raise GenericAbort # "Prep" disk - prep_disk_for_backup(destination, disk, ticket_number) + prep_disk_for_backup(destination, disk, dest_prefix) # Display details for backup task clear_screen() print_info('Create Backup - Details:\n') - show_data(message='Ticket:', data=ticket_number) + if ENABLED_TICKET_NUMBERS: + show_data(message='Ticket:', data=ticket_number) show_data( message = 'Source:', data = '[{}] ({}) {} {}'.format( @@ -293,7 +298,8 @@ def menu_setup(): # Display details for setup task clear_screen() print_info('Setup Windows - Details:\n') - show_data(message='Ticket:', data=ticket_number) + if ENABLED_TICKET_NUMBERS: + show_data(message='Ticket:', data=ticket_number) show_data(message='Installing:', data=windows_version['Name']) show_data( message = 'Boot Method:', From db09236b262ac35a7f7ea7a9f6a87123e624e791 Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 18:06:26 -0600 Subject: [PATCH 3/7] Updated user_data_transfer sections --- .bin/Scripts/functions/data.py | 40 +++++++++++++++--------------- .bin/Scripts/user_data_transfer.py | 15 +++++++++-- 2 files changed, 33 insertions(+), 22 deletions(-) diff --git a/.bin/Scripts/functions/data.py b/.bin/Scripts/functions/data.py index 93ab07ce..5e06a188 100644 --- a/.bin/Scripts/functions/data.py +++ b/.bin/Scripts/functions/data.py @@ -534,21 +534,21 @@ def select_destination(folder_path, prompt='Select destination'): return path -def select_source(ticket_number): - """Select backup from those found on the BACKUP_SERVERS for the ticket.""" +def select_source(backup_prefix): + """Select backup from those found on the BACKUP_SERVERS matching the prefix.""" selected_source = None local_sources = [] remote_sources = [] sources = [] mount_backup_shares(read_write=False) - # Check for ticket folders on servers + # Check for prefix folders on servers for server in BACKUP_SERVERS: if server['Mounted']: print_standard('Scanning {}...'.format(server['Name'])) for d in os.scandir(r'\\{IP}\{Share}'.format(**server)): if (d.is_dir() - and d.name.lower().startswith(ticket_number.lower())): + and d.name.lower().startswith(backup_prefix.lower())): # Add folder to remote_sources remote_sources.append({ 'Name': '{:9}| File-Based: [DIR] {}'.format( @@ -558,19 +558,19 @@ def select_source(ticket_number): 'Source': d}) # Check for images and subfolders - for ticket_path in remote_sources.copy(): - for item in os.scandir(ticket_path['Source'].path): + for prefix_path in remote_sources.copy(): + for item in os.scandir(prefix_path['Source'].path): if item.is_dir(): # Add folder to remote_sources remote_sources.append({ 'Name': r'{:9}| File-Based: [DIR] {}\{}'.format( - ticket_path['Server']['Name'], # Server - ticket_path['Source'].name, # Ticket folder + prefix_path['Server']['Name'], # Server + prefix_path['Source'].name, # Prefix folder item.name, # Sub-folder ), - 'Server': ticket_path['Server'], + 'Server': prefix_path['Server'], 'Sort': r'{}\{}'.format( - ticket_path['Source'].name, # Ticket folder + prefix_path['Source'].name, # Prefix folder item.name, # Sub-folder ), 'Source': item}) @@ -586,15 +586,15 @@ def select_source(ticket_number): remote_sources.append({ 'Disabled': bool(not is_valid_wim_file(subitem)), 'Name': r'{:9}| Image-Based: {:>7} {}\{}\{}'.format( - ticket_path['Server']['Name'], # Server + prefix_path['Server']['Name'], # Server size, # Size (duh) - ticket_path['Source'].name, # Ticket folder + prefix_path['Source'].name, # Prefix folder item.name, # Sub-folder subitem.name, # Image file ), - 'Server': ticket_path['Server'], + 'Server': prefix_path['Server'], 'Sort': r'{}\{}\{}'.format( - ticket_path['Source'].name, # Ticket folder + prefix_path['Source'].name, # Prefix folder item.name, # Sub-folder subitem.name, # Image file ), @@ -608,14 +608,14 @@ def select_source(ticket_number): remote_sources.append({ 'Disabled': bool(not is_valid_wim_file(item)), 'Name': r'{:9}| Image-Based: {:>7} {}\{}'.format( - ticket_path['Server']['Name'], # Server + prefix_path['Server']['Name'], # Server size, # Size (duh) - ticket_path['Source'].name, # Ticket folder + prefix_path['Source'].name, # Prefix folder item.name, # Image file ), - 'Server': ticket_path['Server'], + 'Server': prefix_path['Server'], 'Sort': r'{}\{}'.format( - ticket_path['Source'].name, # Ticket folder + prefix_path['Source'].name, # Prefix folder item.name, # Image file ), 'Source': item}) @@ -682,8 +682,8 @@ def select_source(ticket_number): else: selected_source = sources[int(selection)-1]['Source'] else: - print_error('ERROR: No backups found for ticket: {}.'.format( - ticket_number)) + print_error('ERROR: No backups found using prefix: {}.'.format( + backup_prefix)) umount_backup_shares() pause("Press Enter to exit...") exit_script() diff --git a/.bin/Scripts/user_data_transfer.py b/.bin/Scripts/user_data_transfer.py index 1c787cfe..e0339f20 100644 --- a/.bin/Scripts/user_data_transfer.py +++ b/.bin/Scripts/user_data_transfer.py @@ -18,17 +18,28 @@ if __name__ == '__main__': stay_awake() clear_screen() print_info('{}: User Data Transfer Tool\n'.format(KIT_NAME_FULL)) + + # Get backup name prefix ticket_number = get_ticket_number() + if ENABLED_TICKET_NUMBERS: + backup_prefix = ticket_number + else: + backup_prefix = get_simple_string(prompt='Enter backup name prefix') + + # Set destination folder_path = r'{}\Transfer'.format(KIT_NAME_SHORT) dest = select_destination(folder_path=folder_path, prompt='Which disk are we transferring to?') - source = select_source(ticket_number) + + # Set source items + source = select_source(backup_prefix) items = scan_source(source, dest) # Transfer clear_screen() print_info('Transfer Details:\n') - show_data('Ticket:', ticket_number) + if ENABLED_TICKET_NUMBERS: + show_data('Ticket:', ticket_number) show_data('Source:', source.path) show_data('Destination:', dest) From 02054a804711800fefb4fd09a7139f36d9329a6e Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 18:07:18 -0600 Subject: [PATCH 4/7] Updated system checklist and diagnostic scripts --- .bin/Scripts/system_checklist.py | 7 +++++-- .bin/Scripts/system_diagnostics.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/.bin/Scripts/system_checklist.py b/.bin/Scripts/system_checklist.py index e0f217ed..c24e9fb1 100644 --- a/.bin/Scripts/system_checklist.py +++ b/.bin/Scripts/system_checklist.py @@ -29,8 +29,11 @@ if __name__ == '__main__': 'FileNotFoundError': 'File not found', }, 'Warning': {}} - print_info('Starting System Checklist for Ticket #{}\n'.format( - ticket_number)) + if ENABLED_TICKET_NUMBERS: + print_info('Starting System Checklist for Ticket #{}\n'.format( + ticket_number)) + else: + print_info('Starting System Checklist\n') # Configure print_info('Configure') diff --git a/.bin/Scripts/system_diagnostics.py b/.bin/Scripts/system_diagnostics.py index 59e31465..518eee8c 100644 --- a/.bin/Scripts/system_diagnostics.py +++ b/.bin/Scripts/system_diagnostics.py @@ -31,8 +31,11 @@ if __name__ == '__main__': 'GenericRepair': 'Repaired', 'UnsupportedOSError': 'Unsupported OS', }} - print_info('Starting System Diagnostics for Ticket #{}\n'.format( - ticket_number)) + if ENABLED_TICKET_NUMBERS: + print_info('Starting System Diagnostics for Ticket #{}\n'.format( + ticket_number)) + else: + print_info('Starting System Diagnostics\n') # Sanitize Environment print_info('Sanitizing Environment') From b997a523853c62c2452d8281ade11b586fc2831f Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 18:17:36 -0600 Subject: [PATCH 5/7] Updated hw-diags section * Use date_time instead of ticket_number if disabled --- .bin/Scripts/functions/hw_diags.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.bin/Scripts/functions/hw_diags.py b/.bin/Scripts/functions/hw_diags.py index 1ece188c..63942ed2 100644 --- a/.bin/Scripts/functions/hw_diags.py +++ b/.bin/Scripts/functions/hw_diags.py @@ -125,9 +125,9 @@ def menu_diags(*args): if diag_modes[int(selection)-1]['Name'] != 'Quick drive test': # Save log for non-quick tests ticket_number = get_ticket_number() - global_vars['LogDir'] = '{}/Tickets/{}'.format( + global_vars['LogDir'] = '{}/Logs/{}'.format( global_vars['Env']['HOME'], - ticket_number) + ticket_number if ticket_number else global_vars['Date-Time']) os.makedirs(global_vars['LogDir'], exist_ok=True) global_vars['LogFile'] = '{}/Hardware Diagnostics.log'.format( global_vars['LogDir']) From 27953bde5acdb45283ec4dfcabce518415c7f761 Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 20:54:45 -0600 Subject: [PATCH 6/7] Update backup/restore code * Added support for local backups * Added volume label detection for local backups * Replace spaces in backup_prefix with underscores --- .bin/Scripts/functions/backup.py | 51 +++++++++++++++++++++++++-- .bin/Scripts/functions/winpe_menus.py | 13 ++++--- .bin/Scripts/user_data_transfer.py | 1 + 3 files changed, 58 insertions(+), 7 deletions(-) diff --git a/.bin/Scripts/functions/backup.py b/.bin/Scripts/functions/backup.py index a787a0ed..e33cea2c 100644 --- a/.bin/Scripts/functions/backup.py +++ b/.bin/Scripts/functions/backup.py @@ -1,5 +1,7 @@ # Wizard Kit: Functions - Backup +import ctypes + from functions.disk import * # Regex @@ -30,7 +32,39 @@ def backup_partition(disk, par): def fix_path(path): """Replace invalid filename characters with underscores.""" - return REGEX_BAD_PATH_NAMES.sub('_', path) + local_drive = path[1:2] == ':' + new_path = REGEX_BAD_PATH_NAMES.sub('_', path) + if local_drive: + new_path = '{}:{}'.format(new_path[0:1], new_path[2:]) + return new_path + +def get_volume_display_name(mountpoint): + """Get display name from volume mountpoint and label, returns str.""" + name = mountpoint + try: + kernel32 = ctypes.windll.kernel32 + vol_name_buffer = ctypes.create_unicode_buffer(1024) + fs_name_buffer = ctypes.create_unicode_buffer(1024) + serial_number = None + max_component_length = None + file_system_flags = None + + vol_info = kernel32.GetVolumeInformationW( + ctypes.c_wchar_p(mountpoint), + vol_name_buffer, + ctypes.sizeof(vol_name_buffer), + serial_number, + max_component_length, + file_system_flags, + fs_name_buffer, + ctypes.sizeof(fs_name_buffer) + ) + + name = '{} "{}"'.format(name, vol_name_buffer.value) + except: + pass + + return name def prep_disk_for_backup(destination, disk, backup_prefix): """Gather details about the disk and its partitions. @@ -114,7 +148,20 @@ def select_backup_destination(auto_select=True): actions = [ {'Name': 'Main Menu', 'Letter': 'M'}, ] - + + # Add local disks + for d in psutil.disk_partitions(): + if re.search(r'^{}'.format(global_vars['Env']['SYSTEMDRIVE']), d.mountpoint, re.IGNORECASE): + # Skip current OS drive + pass + elif 'fixed' in d.opts: + # Skip DVD, etc + destinations.append({ + 'Name': 'Local Disk - {}'.format( + get_volume_display_name(d.mountpoint)), + 'Letter': re.sub(r'^(\w):\\.*$', r'\1', d.mountpoint), + }) + # Size check for dest in destinations: if 'IP' in dest: diff --git a/.bin/Scripts/functions/winpe_menus.py b/.bin/Scripts/functions/winpe_menus.py index c5d6dd99..6d8d10bc 100644 --- a/.bin/Scripts/functions/winpe_menus.py +++ b/.bin/Scripts/functions/winpe_menus.py @@ -83,7 +83,14 @@ def menu_backup(): backup_prefix = ticket_number else: backup_prefix = get_simple_string(prompt='Enter backup name prefix') + backup_prefix = backup_prefix.replace(' ', '_') + # Assign drive letters + try_and_print( + message = 'Assigning letters...', + function = assign_volume_letters, + other_results = other_results) + # Mount backup shares mount_backup_shares(read_write=True) @@ -91,10 +98,6 @@ def menu_backup(): destination = select_backup_destination(auto_select=False) # Scan disks - try_and_print( - message = 'Assigning letters...', - function = assign_volume_letters, - other_results = other_results) result = try_and_print( message = 'Getting disk info...', function = scan_disks, @@ -111,7 +114,7 @@ def menu_backup(): raise GenericAbort # "Prep" disk - prep_disk_for_backup(destination, disk, dest_prefix) + prep_disk_for_backup(destination, disk, backup_prefix) # Display details for backup task clear_screen() diff --git a/.bin/Scripts/user_data_transfer.py b/.bin/Scripts/user_data_transfer.py index e0339f20..ce572f69 100644 --- a/.bin/Scripts/user_data_transfer.py +++ b/.bin/Scripts/user_data_transfer.py @@ -25,6 +25,7 @@ if __name__ == '__main__': backup_prefix = ticket_number else: backup_prefix = get_simple_string(prompt='Enter backup name prefix') + backup_prefix = backup_prefix.replace(' ', '_') # Set destination folder_path = r'{}\Transfer'.format(KIT_NAME_SHORT) From e305b3e60f91ef1d860d0e1a45fbc3756330dd4f Mon Sep 17 00:00:00 2001 From: 2Shirt <1923621+2Shirt@users.noreply.github.com> Date: Thu, 17 May 2018 20:56:55 -0600 Subject: [PATCH 7/7] Add warning for missing archive in extract_item() --- .bin/Scripts/functions/common.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.bin/Scripts/functions/common.py b/.bin/Scripts/functions/common.py index ed1d1938..f5a2ef95 100644 --- a/.bin/Scripts/functions/common.py +++ b/.bin/Scripts/functions/common.py @@ -190,6 +190,9 @@ def extract_item(item, filter='', silent=False): print_standard('Extracting "{item}"...'.format(item=item)) try: run_program(cmd) + except FileNotFoundError: + if not silent: + print_warning('WARNING: Archive not found') except subprocess.CalledProcessError: if not silent: print_warning('WARNING: Errors encountered while exctracting data')