From 446867b611216a58ed256caef7e331e4aff023f2 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 13:39:00 -0800 Subject: [PATCH 01/23] Started Python refactoring # Plan * Use current WizardKit scripts as the new base * Split functions into multiple files under Scripts\functions\ * Review menus and menu-flow * Log everything and open log at end of menu-flows * (before returning to the root/main menu) --- .gitignore | 3 +- Scripts/functions.py | 1006 ----------------- Scripts/functions/backup.py | 486 ++++++++ Scripts/functions/common.py | 682 +++++++++++ Scripts/functions/data.py | 607 ++++++++++ Scripts/{ => functions}/partition_uids.py | 2 +- Scripts/functions/windows_setup.py | 273 +++++ Scripts/{menu.py => functions/winpe_menus.py} | 208 ++-- Scripts/settings/main.py | 68 ++ Scripts/settings/tools.py | 55 + Scripts/winpe_root_menu.py | 21 + 11 files changed, 2317 insertions(+), 1094 deletions(-) delete mode 100644 Scripts/functions.py create mode 100644 Scripts/functions/backup.py create mode 100644 Scripts/functions/common.py create mode 100644 Scripts/functions/data.py rename Scripts/{ => functions}/partition_uids.py (98%) create mode 100644 Scripts/functions/windows_setup.py rename Scripts/{menu.py => functions/winpe_menus.py} (66%) create mode 100644 Scripts/settings/main.py create mode 100644 Scripts/settings/tools.py create mode 100644 Scripts/winpe_root_menu.py diff --git a/.gitignore b/.gitignore index ed8a9f58..6e041ca5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,11 +1,10 @@ +**/__pycache__/* *.bak *.iso -.bin/Scripts/__pycache__ .bin/tmp Drivers Logs Mount PEFiles -Scripts/__pycache__ WK/amd64/ WK/x86/ diff --git a/Scripts/functions.py b/Scripts/functions.py deleted file mode 100644 index 03f03bf3..00000000 --- a/Scripts/functions.py +++ /dev/null @@ -1,1006 +0,0 @@ -# WK WinPE Functions - -# Init -import os -import re -import shutil -import subprocess -import sys -import time -import winreg -os.chdir(os.path.dirname(os.path.realpath(__file__))) -bin = os.path.abspath('..\\') -sys.path.append(os.getcwd()) -from functions import * -import partition_uids - -# Init -BACKUP_SERVERS = [ - { 'IP': '10.0.0.10', - 'Mounted': False, - 'Name': 'ServerOne', - 'Share': 'Backups', - 'User': 'backup', - 'Pass': 'Abracadabra', - }, - { 'IP': '10.0.0.11', - 'Name': 'ServerTwo', - 'Mounted': False, - 'Share': 'Backups', - 'User': 'backup', - 'Pass': 'Abracadabra', - }, -] -WINDOWS_SERVER = { - 'IP': '10.0.0.10', - 'Name': 'ServerOne', - 'Mounted': False, - 'Share': 'Windows', - 'User': 'backup', # Using these credentials in case both the windows source and backup shares are mounted. - 'Pass': 'Abracadabra', # This is because Windows only allows one set of credentials to be used per server at a time. -} -WINDOWS_VERSIONS = [ - {'Name': 'Windows 7 Home Basic', - 'Image File': 'Win7', - 'Image Name': 'Windows 7 HOMEBASIC', - 'Family': '7'}, - {'Name': 'Windows 7 Home Premium', - 'Image File': 'Win7', - 'Image Name': 'Windows 7 HOMEPREMIUM', - 'Family': '7'}, - {'Name': 'Windows 7 Professional', - 'Image File': 'Win7', - 'Image Name': 'Windows 7 PROFESSIONAL', - 'Family': '7'}, - {'Name': 'Windows 7 Ultimate', - 'Image File': 'Win7', - 'Image Name': 'Windows 7 ULTIMATE', - 'Family': '7'}, - - {'Name': 'Windows 8.1', - 'Image File': 'Win8', - 'Image Name': 'Windows 8.1', - 'Family': '8', - 'CRLF': True}, - {'Name': 'Windows 8.1 Pro', - 'Image File': 'Win8', - 'Image Name': 'Windows 8.1 Pro', - 'Family': '8'}, - - {'Name': 'Windows 10 Home', - 'Image File': 'Win10', - 'Image Name': 'Windows 10 Home', - 'Family': '10', - 'CRLF': True}, - {'Name': 'Windows 10 Pro', - 'Image File': 'Win10', - 'Image Name': 'Windows 10 Pro', - 'Family': '10'}, -] -diskpart_script = '{tmp}\\diskpart.script'.format(tmp=os.environ['TMP']) - -## Colors -COLORS = { - 'CLEAR': '\033[0m', - 'RED': '\033[31m', - 'GREEN': '\033[32m', - 'YELLOW': '\033[33m', - 'BLUE': '\033[34m'} - -class AbortError(Exception): - pass - -class BackupError(Exception): - pass - -class SetupError(Exception): - pass - -def abort_to_main_menu(message='Returning to main menu...'): - print_warning(message) - pause('Press Enter to return to main menu... ') - raise AbortError - -def ask(prompt='Kotaero'): - answer = None - prompt = prompt + ' [Y/N]: ' - while answer is None: - tmp = input(prompt) - if re.search(r'^y(es|)$', tmp, re.IGNORECASE): - answer = True - elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE): - answer = False - return answer - -def assign_volume_letters(): - try: - # Run script - with open(diskpart_script, 'w') as script: - for vol in get_volumes(): - script.write('select volume {Number}\n'.format(**vol)) - script.write('assign\n') - run_program('diskpart /s {script}'.format(script=diskpart_script)) - except subprocess.CalledProcessError: - pass - -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.') - - 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 - -def convert_to_bytes(size): - size = str(size) - tmp = re.search(r'(\d+)\s+([KMGT]B)', size.upper()) - if tmp: - size = int(tmp.group(1)) - units = tmp.group(2) - if units == 'TB': - size *= 1099511627776 - elif units == 'GB': - size *= 1073741824 - elif units == 'MB': - size *= 1048576 - elif units == 'KB': - size *= 1024 - else: - return -1 - - return size - -def is_valid_image(bin=None, filename=None, imagename=None): - # Bail early - if bin is None: - raise Exception('bin not specified.') - if filename is None: - raise Exception('Filename not specified.') - if imagename is None: - raise Exception('Image Name not specified.') - - cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename) - try: - run_program(cmd) - except subprocess.CalledProcessError: - print_error('Invalid image: {filename}'.format(filename=filename)) - return False - - return True - -def find_windows_image(bin, windows_version=None): - """Search for a Windows source image file on local drives and network drives (in that order)""" - image = {} - - # Bail early - if windows_version is None: - raise Exception('Windows version not specified.') - imagefile = windows_version['Image File'] - - # Search local source - process_return = run_program('mountvol') - for tmp in re.findall(r'.*([A-Za-z]):\\', process_return.stdout.decode()): - for ext in ['esd', 'wim', 'swm']: - filename = '{drive}:\\images\\{imagefile}'.format(drive=tmp[0], imagefile=imagefile) - filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) - if os.path.isfile(filename_ext): - if is_valid_image(bin, filename_ext, windows_version['Image Name']): - image['Ext'] = ext - image['File'] = filename - image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' - image['Source'] = tmp[0] - break - - # Check for network source (if necessary) - if not any(image): - if not WINDOWS_SERVER['Mounted']: - mount_windows_share() - for ext in ['esd', 'wim', 'swm']: - filename = '\\\\{IP}\\{Share}\\images\\{imagefile}'.format(imagefile=imagefile, **WINDOWS_SERVER) - filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) - if os.path.isfile(filename_ext): - if is_valid_image(bin, filename_ext, windows_version['Image Name']): - image['Ext'] = ext - image['File'] = filename - image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' - image['Source'] = None - break - - # Display image to be used (if any) and return - if any(image): - print_info('Using image: {File}.{Ext}'.format(**image)) - return image - else: - print_error('Failed to find Windows source image for {winver}'.format(winver=windows_version['Name'])) - abort_to_main_menu('Aborting Windows setup') - -def format_gpt(disk=None, windows_family=None): - """Format disk for use as a Windows OS drive using the GPT (UEFI) layout.""" - - # Bail early - if disk is None: - raise Exception('No disk provided.') - if windows_family is None: - raise Exception('No Windows family provided.') - - # Format drive - # print_info('Drive will use a GPT (UEFI) layout.') - with open(diskpart_script, 'w') as script: - # Partition table - script.write('select disk {number}\n'.format(number=disk['Number'])) - script.write('clean\n') - script.write('convert gpt\n') - - # System partition - script.write('create partition efi size=260\n') # NOTE: Allows for Advanced Format 4K drives - script.write('format quick fs=fat32 label="System"\n') - script.write('assign letter="S"\n') - - # Microsoft Reserved (MSR) partition - script.write('create partition msr size=128\n') - - # Windows partition - script.write('create partition primary\n') - script.write('format quick fs=ntfs label="Windows"\n') - script.write('assign letter="W"\n') - - # Recovery Tools partition (Windows 8+) - if re.search(r'^(8|10)', windows_family): - script.write('shrink minimum=500\n') - script.write('create partition primary\n') - script.write('format quick fs=ntfs label="Recovery Tools"\n') - script.write('assign letter="T"\n') - script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n') - script.write('gpt attributes=0x8000000000000001\n') - - # Run script - run_program('diskpart /s {script}'.format(script=diskpart_script)) - time.sleep(2) - -def format_mbr(disk=None, windows_family=None): - """Format disk for use as a Windows OS drive using the MBR (legacy) layout.""" - - # Bail early - if disk is None: - raise Exception('No disk provided.') - if windows_family is None: - raise Exception('No Windows family provided.') - - # Format drive - # print_info('Drive will use a MBR (legacy) layout.') - with open(diskpart_script, 'w') as script: - # Partition table - script.write('select disk {number}\n'.format(number=disk['Number'])) - script.write('clean\n') - - # System partition - script.write('create partition primary size=100\n') - script.write('format fs=ntfs quick label="System Reserved"\n') - script.write('active\n') - script.write('assign letter="S"\n') - - # Windows partition - script.write('create partition primary\n') - script.write('format fs=ntfs quick label="Windows"\n') - script.write('assign letter="W"\n') - - # Recovery Tools partition (Windows 8+) - if re.search(r'^(8|10)', windows_family): - script.write('shrink minimum=500\n') - script.write('create partition primary\n') - script.write('format quick fs=ntfs label="Recovery"\n') - script.write('assign letter="T"\n') - script.write('set id=27\n') - - # Run script - run_program('diskpart /s {script}'.format(script=diskpart_script)) - time.sleep(2) - -def get_attached_disk_info(): - """Get details about the attached disks""" - disks = [] - print_info('Getting drive info...') - - # Assign all the letters - assign_volume_letters() - - # 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: - reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 'System\\CurrentControlSet\\Control') - reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] - if reg_value == 2: - boot_mode = 'UEFI' - except: - boot_mode = 'Unknown' - - return boot_mode - -def get_disk_details(disk=None): - details = {} - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - - try: - # Run script - with open(diskpart_script, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('detail disk\n') - process_return = run_program('diskpart /s {script}'.format(script=diskpart_script)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Remove empty lines - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Set disk name - details['Name'] = tmp[4] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair - tmp = [s.split(':') for s in tmp if ':' in s] - - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) - - return details - -def get_disks(): - disks = [] - - try: - # Run script - with open(diskpart_script, 'w') as script: - script.write('list disk\n') - process_return = run_program('diskpart /s {script}'.format(script=diskpart_script)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Append disk numbers - for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return): - _num = tmp[0] - _size = human_readable_size(tmp[1]) - disks.append({'Number': _num, 'Size': _size}) - - return disks - -def get_partition_details(disk=None, par=None): - details = {} - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - if par is None: - raise Exception('Partition not specified.') - - # Diskpart details - try: - # Run script - with open(diskpart_script, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('select partition {Number}\n'.format(**par)) - script.write('detail partition\n') - process_return = run_program('diskpart /s {script}'.format(script=diskpart_script)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Get volume letter or RAW status - tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return) - if tmp: - if tmp.group(1).upper() == 'RAW': - details['FileSystem'] = RAW - else: - details['Letter'] = tmp.group(1) - - # Remove empty lines from process_return - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair - tmp = [s.split(':') for s in tmp if ':' in s] - - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) - - # Get MBR type / GPT GUID for extra details on "Unknown" partitions - guid = partition_uids.lookup_guid(details['Type']) - if guid is not None: - details.update({ - 'Description': guid.get('Description', ''), - 'OS': guid.get('OS', '')}) - - if 'Letter' in details: - # Disk usage - tmp = shutil.disk_usage('{Letter}:\\'.format(**details)) - details['Used Space'] = human_readable_size(tmp.used) - - # fsutil details - try: - process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Remove empty lines from process_return - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Add "Feature" lines - details['File System Features'] = [s.strip() for s in tmp if ':' not in s] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair - tmp = [s.split(':') for s in tmp if ':' in s] - - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) - - # Set Volume Name - details['Name'] = details.get('Volume Name', '') - - # Set FileSystem Type - if details.get('FileSystem', '') != 'RAW': - details['FileSystem'] = details.get('File System Name', 'Unknown') - - return details - -def get_partitions(disk=None): - partitions = [] - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - - try: - # Run script - with open(diskpart_script, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('list partition\n') - process_return = run_program('diskpart /s {script}'.format(script=diskpart_script)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Append partition numbers - for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE): - _num = tmp[0] - _size = human_readable_size(tmp[1]) - partitions.append({'Number': _num, 'Size': _size}) - - return partitions - -def get_table_type(disk=None): - _type = 'Unknown' - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - - try: - with open(diskpart_script, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('uniqueid disk\n') - process_return = run_program('diskpart /s {script}'.format(script=diskpart_script)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - if re.findall(r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', process_return, re.IGNORECASE): - _type = 'GPT' - elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE): - _type = 'RAW' - elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE): - _type = 'MBR' - - return _type - -def get_ticket_id(): - ticket_id = None - while ticket_id is None: - tmp = input('Enter ticket number: ') - if re.match(r'^([0-9]+([\-_]*\w+|))$', tmp): - ticket_id = tmp - - return ticket_id - -def get_volumes(): - vols = [] - - try: - # Run script - with open(diskpart_script, 'w') as script: - script.write('list volume\n') - process_return = run_program('diskpart /s {script}'.format(script=diskpart_script)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Append volume numbers - for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return): - vols.append({'Number': tmp[0], 'Letter': tmp[1]}) - - return vols - -def human_readable_size(size, decimals=0): - # Prep string formatting - width = 3+decimals - if decimals > 0: - width += 1 - human_format = '>{width}.{decimals}f'.format(width=width, decimals=decimals) - tmp = '' - - # Convert size to int - try: - size = int(size) - except ValueError: - size = convert_to_bytes(size) - - # Verify we have a valid size - if size <= 0: - return '{size:>{width}} b'.format(size='???', width=width) - - # Format string - if size >= 1099511627776: - size /= 1099511627776 - tmp = '{size:{human_format}} Tb'.format(size=size, human_format=human_format) - elif size >= 1073741824: - size /= 1073741824 - tmp = '{size:{human_format}} Gb'.format(size=size, human_format=human_format) - elif size >= 1048576: - size /= 1048576 - tmp = '{size:{human_format}} Mb'.format(size=size, human_format=human_format) - elif size >= 1024: - size /= 1024 - tmp = '{size:{human_format}} Kb'.format(size=size, human_format=human_format) - else: - tmp = '{size:{human_format}} b'.format(size=size, human_format=human_format) - - # Return - return tmp - -def menu_select(title='~ Untitled Menu ~', main_entries=[], action_entries=[], prompt='Please make a selection', secret_exit=False): - """Display options in a menu for user selection""" - - # Bail early - if (len(main_entries) + len(action_entries) == 0): - raise Exception("MenuError: No items given") - - # Build menu - menu_splash = '{title}\n\n'.format(title=title) - valid_answers = [] - if (secret_exit): - valid_answers.append('Q') - - # Add main entries - if (len(main_entries) > 0): - for i in range(len(main_entries)): - entry = main_entries[i] - # Add Spacer - if ('CRLF' in entry): - menu_splash += '\n' - valid_answers.append(str(i+1)) - menu_splash += '{number:>{mwidth}}: {name}\n'.format(number=i+1, mwidth=len(str(len(main_entries))), name=entry.get('Display Name', entry['Name'])) - menu_splash += '\n' - - # Add action entries - if (len(action_entries) > 0): - for entry in action_entries: - # Add Spacer - if ('CRLF' in entry): - menu_splash += '\n' - valid_answers.append(entry['Letter']) - menu_splash += '{letter:>{mwidth}}: {name}\n'.format(letter=entry['Letter'].upper(), mwidth=len(str(len(action_entries))), name=entry['Name']) - menu_splash += '\n' - - answer = '' - - while (answer.upper() not in valid_answers): - os.system('cls') - print(menu_splash) - answer = input('{prompt}: '.format(prompt=prompt)) - - return answer.upper() - -def mount_backup_shares(): - """Mount the backup shares for use as destinations for backup image creation""" - - # Attempt to mount share(s) - for server in BACKUP_SERVERS: - # Blindly skip if we mounted earlier - if server['Mounted']: - continue - else: - try: - # Test connection - run_program('ping -w 800 -n 2 {IP}'.format(**server)) - - # Mount - run_program('net use \\\\{IP}\\{Share} /user:{User} {Pass}'.format(**server)) - print_info('Mounted {Name}'.format(**server)) - server['Mounted'] = True - except subprocess.CalledProcessError: - print_error('Failed to mount \\\\{Name}\\{Share}, {IP} unreachable.'.format(**server)) - time.sleep(1) - except: - print_warning('Failed to mount \\\\{Name}\\{Share} ({IP})'.format(**server)) - time.sleep(1) - -def mount_windows_share(): - """Mount the Windows images share for use in Windows setup""" - - # Blindly skip if we mounted earlier - if WINDOWS_SERVER['Mounted']: - return None - else: - try: - # Test connection - run_program('ping -w 800 -n 2 {IP}'.format(**WINDOWS_SERVER)) - # Mount - run_program('net use \\\\{IP}\\{Share} /user:{User} {Pass}'.format(**WINDOWS_SERVER)) - print_info('Mounted {Name}'.format(**WINDOWS_SERVER)) - WINDOWS_SERVER['Mounted'] = True - except subprocess.CalledProcessError: - print_error('Failed to mount \\\\{Name}\\{Share}, {IP} unreachable.'.format(**WINDOWS_SERVER)) - time.sleep(1) - except: - print_warning('Failed to mount \\\\{Name}\\{Share}'.format(**WINDOWS_SERVER)) - time.sleep(1) - -def pause(prompt='Press Enter to continue... '): - input(prompt) - -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 prep_disk_for_formatting(disk=None): - disk['Format Warnings'] = '\n' - width = len(str(len(disk['Partitions']))) - - # Bail early - if disk is None: - raise Exception('Disk not provided.') - - # Set boot method and partition table type - disk['Use GPT'] = True - if (get_boot_mode() == 'UEFI'): - if (not ask("Setup Windows to use UEFI booting?")): - disk['Use GPT'] = False - else: - if (ask("Setup Windows to use BIOS/Legacy booting?")): - disk['Use GPT'] = False - - # Set Display and Warning Strings - if len(disk['Partitions']) == 0: - disk['Format Warnings'] += 'No partitions found\n' - for par in disk['Partitions']: - if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): - # FileSystem not accessible to WinPE. List partition type / OS info for technician - par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format( - width=width, - q='"' if par['Name'] != '' else '', - **par) - else: - # FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician - par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format( - width=width, - q='"' if par['Name'] != '' else '', - **par) - -def print_error(message='Generic error', **kwargs): - print('{RED}{message}{CLEAR}'.format(message=message, **COLORS, **kwargs)) - -def print_info(message='Generic info', **kwargs): - print('{BLUE}{message}{CLEAR}'.format(message=message, **COLORS, **kwargs)) - -def print_success(message='Generic success', **kwargs): - print('{GREEN}{message}{CLEAR}'.format(message=message, **COLORS, **kwargs)) - -def print_warning(message='Generic warning', **kwargs): - print('{YELLOW}{message}{CLEAR}'.format(message=message, **COLORS, **kwargs)) - -def remove_volume_letters(keep=''): - if keep is None: - keep = '' - try: - # Run script - with open(diskpart_script, 'w') as script: - for vol in get_volumes(): - if vol['Letter'].upper() != keep.upper(): - script.write('select volume {Number}\n'.format(**vol)) - script.write('remove noerr\n') - run_program('diskpart /s {script}'.format(script=diskpart_script)) - except subprocess.CalledProcessError: - pass - -def run_program(cmd=None, args=[], check=True): - if cmd is None: - raise Exception('No program passed.') - - if len(args) > 0: - args = [cmd] + args - process_return = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=check) - else: - process_return = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=check) - - return process_return - -def select_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 select_disk(prompt='Which disk?'): - """Select a disk from the attached disks""" - disks = get_attached_disk_info() - - # Build menu - disk_options = [] - for disk in disks: - display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) - if len(disk['Partitions']) > 0: - pwidth=len(str(len(disk['Partitions']))) - for par in disk['Partitions']: - # Show unsupported partition(s) in RED - par_skip = False - if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): - par_skip = True - if par_skip: - display_name += COLORS['YELLOW'] - - # Main text - display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par) - if par['Name'] != '': - display_name += '\t"{Name}"'.format(**par) - - # Clear color (if set above) - if par_skip: - display_name += COLORS['CLEAR'] - else: - display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS) - disk_options.append({'Name': display_name, 'Disk': disk}) - actions = [ - {'Name': 'Main Menu', 'Letter': 'M'}, - ] - - # Menu loop - selection = menu_select(prompt, disk_options, actions) - - if (selection.isnumeric()): - return disk_options[int(selection)-1]['Disk'] - elif (selection == 'M'): - abort_to_main_menu() - -def select_minidump_path(): - dumps = [] - - # Assign volume letters first - assign_volume_letters() - - # Search for minidumps - tmp = run_program('mountvol') - tmp = [d for d in re.findall(r'.*([A-Za-z]):\\', tmp.stdout.decode())] - # Remove RAMDisk letter - if 'X' in tmp: - tmp.remove('X') - for drive in tmp: - if os.path.exists('{drive}:\\Windows\\MiniDump'.format(drive=drive)): - dumps.append({'Name': '{drive}:\\Windows\\MiniDump'.format(drive=drive)}) - - # Check results before showing menu - if len(dumps) == 0: - print_error(' No BSoD / MiniDump paths found') - time.sleep(2) - return None - - # Menu - selection = menu_select('Which BSoD / MiniDump path are we scanning?', dumps, []) - return dumps[int(selection) - 1]['Name'] - -def select_windows_version(): - actions = [{'Name': 'Main Menu', 'Letter': 'M'},] - - # Menu loop - selection = menu_select('Which version of Windows are we installing?', WINDOWS_VERSIONS, actions) - - if selection.isnumeric(): - return WINDOWS_VERSIONS[int(selection)-1] - elif selection == 'M': - abort_to_main_menu() - -def setup_windows(bin=None, windows_image=None, windows_version=None): - # Bail early - if bin is None: - raise Exception('bin path not specified.') - if windows_image is None: - raise Exception('Windows image not specified.') - if windows_version is None: - raise Exception('Windows version not specified.') - - # Apply image - cmd = '{bin}\\wimlib\\wimlib-imagex apply "{File}.{Ext}" "{Image Name}" W:\\ {Glob}'.format(bin=bin, **windows_image, **windows_version) - run_program(cmd) - -def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T'): - # Bail early - if windows_version is None: - raise Exception('Windows version not specified.') - - _win = '{win}:\\Windows'.format(win=windows_letter) - _winre = '{win}\\System32\\Recovery\\WinRE.wim'.format(win=_win) - _dest = '{tools}:\\Recovery\\WindowsRE'.format(tools=tools_letter) - - if re.search(r'^(8|10)', windows_version['Family']): - # Copy WinRE.wim - os.makedirs(_dest, exist_ok=True) - shutil.copy(_winre, '{dest}\\WinRE.wim'.format(dest=_dest)) - - # Set location - run_program('{win}\\System32\\reagentc /setreimage /path {dest} /target {win}'.format(dest=_dest, win=_win)) - else: - # Only supported on Windows 8 and above - raise SetupError - -def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'): - run_program('bcdboot {win}:\\Windows /s {sys}: /f {mode}'.format(win=windows_letter, sys=system_letter, mode=mode)) - -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.") diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py new file mode 100644 index 00000000..94f3b103 --- /dev/null +++ b/Scripts/functions/backup.py @@ -0,0 +1,486 @@ +# Wizard Kit PE: Functions - Backup + +from functions.common import * +import partition_uids + +def assign_volume_letters(): + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + for vol in get_volumes(): + script.write('select volume {Number}\n'.format(**vol)) + script.write('assign\n') + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + except subprocess.CalledProcessError: + pass + +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.') + + 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 + +def get_attached_disk_info(): + """Get details about the attached disks""" + disks = [] + print_info('Getting drive info...') + + # Assign all the letters + assign_volume_letters() + + # 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_disk_details(disk=None): + details = {} + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('detail disk\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Remove empty lines + tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] + + # Set disk name + details['Name'] = tmp[4] + + # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + tmp = [s.split(':') for s in tmp if ':' in s] + + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + return details + +def get_disks(): + disks = [] + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('list disk\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Append disk numbers + for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return): + _num = tmp[0] + _size = human_readable_size(tmp[1]) + disks.append({'Number': _num, 'Size': _size}) + + return disks + +def get_partition_details(disk=None, par=None): + details = {} + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + if par is None: + raise Exception('Partition not specified.') + + # Diskpart details + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('select partition {Number}\n'.format(**par)) + script.write('detail partition\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Get volume letter or RAW status + tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return) + if tmp: + if tmp.group(1).upper() == 'RAW': + details['FileSystem'] = RAW + else: + details['Letter'] = tmp.group(1) + + # Remove empty lines from process_return + tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] + + # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + tmp = [s.split(':') for s in tmp if ':' in s] + + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + # Get MBR type / GPT GUID for extra details on "Unknown" partitions + guid = partition_uids.lookup_guid(details['Type']) + if guid is not None: + details.update({ + 'Description': guid.get('Description', ''), + 'OS': guid.get('OS', '')}) + + if 'Letter' in details: + # Disk usage + tmp = shutil.disk_usage('{Letter}:\\'.format(**details)) + details['Used Space'] = human_readable_size(tmp.used) + + # fsutil details + try: + process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Remove empty lines from process_return + tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] + + # Add "Feature" lines + details['File System Features'] = [s.strip() for s in tmp if ':' not in s] + + # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + tmp = [s.split(':') for s in tmp if ':' in s] + + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + # Set Volume Name + details['Name'] = details.get('Volume Name', '') + + # Set FileSystem Type + if details.get('FileSystem', '') != 'RAW': + details['FileSystem'] = details.get('File System Name', 'Unknown') + + return details + +def get_partitions(disk=None): + partitions = [] + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('list partition\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Append partition numbers + for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE): + _num = tmp[0] + _size = human_readable_size(tmp[1]) + partitions.append({'Number': _num, 'Size': _size}) + + return partitions + +def get_table_type(disk=None): + _type = 'Unknown' + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + + try: + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('uniqueid disk\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + if re.findall(r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', process_return, re.IGNORECASE): + _type = 'GPT' + elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE): + _type = 'RAW' + elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE): + _type = 'MBR' + + return _type + +def get_volumes(): + vols = [] + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('list volume\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Append volume numbers + for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return): + vols.append({'Number': tmp[0], 'Letter': tmp[1]}) + + return vols + +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 prep_disk_for_formatting(disk=None): + disk['Format Warnings'] = '\n' + width = len(str(len(disk['Partitions']))) + + # Bail early + if disk is None: + raise Exception('Disk not provided.') + + # Set boot method and partition table type + disk['Use GPT'] = True + if (get_boot_mode() == 'UEFI'): + if (not ask("Setup Windows to use UEFI booting?")): + disk['Use GPT'] = False + else: + if (ask("Setup Windows to use BIOS/Legacy booting?")): + disk['Use GPT'] = False + + # Set Display and Warning Strings + if len(disk['Partitions']) == 0: + disk['Format Warnings'] += 'No partitions found\n' + for par in disk['Partitions']: + if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): + # FileSystem not accessible to WinPE. List partition type / OS info for technician + par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format( + width=width, + q='"' if par['Name'] != '' else '', + **par) + else: + # FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician + par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format( + width=width, + q='"' if par['Name'] != '' else '', + **par) + +def remove_volume_letters(keep=''): + if keep is None: + keep = '' + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + for vol in get_volumes(): + if vol['Letter'].upper() != keep.upper(): + script.write('select volume {Number}\n'.format(**vol)) + script.write('remove noerr\n') + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + except subprocess.CalledProcessError: + pass + +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 select_disk(prompt='Which disk?'): + """Select a disk from the attached disks""" + disks = get_attached_disk_info() + + # Build menu + disk_options = [] + for disk in disks: + display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) + if len(disk['Partitions']) > 0: + pwidth=len(str(len(disk['Partitions']))) + for par in disk['Partitions']: + # Show unsupported partition(s) in RED + par_skip = False + if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): + par_skip = True + if par_skip: + display_name += COLORS['YELLOW'] + + # Main text + display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par) + if par['Name'] != '': + display_name += '\t"{Name}"'.format(**par) + + # Clear color (if set above) + if par_skip: + display_name += COLORS['CLEAR'] + else: + display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS) + disk_options.append({'Name': display_name, 'Disk': disk}) + actions = [ + {'Name': 'Main Menu', 'Letter': 'M'}, + ] + + # Menu loop + selection = menu_select(prompt, disk_options, actions) + + if (selection.isnumeric()): + return disk_options[int(selection)-1]['Disk'] + elif (selection == 'M'): + abort_to_main_menu() + +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.") diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py new file mode 100644 index 00000000..a0453a95 --- /dev/null +++ b/Scripts/functions/common.py @@ -0,0 +1,682 @@ +# Wizard Kit PE: Functions - Common + +import os +import psutil +import re +import shutil +import subprocess +import sys +import time +import traceback +import winreg + +from subprocess import CalledProcessError + +from settings.main import * +from settings.tools import * + +# Global variables +global_vars = {} + +# STATIC VARIABLES +COLORS = { + 'CLEAR': '\033[0m', + 'RED': '\033[31m', + 'GREEN': '\033[32m', + 'YELLOW': '\033[33m', + 'BLUE': '\033[34m' +} +HKU = winreg.HKEY_USERS +HKCU = winreg.HKEY_CURRENT_USER +HKLM = winreg.HKEY_LOCAL_MACHINE + +# Error Classes +class BIOSKeyNotFoundError(Exception): + pass + +class BinNotFoundError(Exception): + pass + +class GenericError(Exception): + pass + +class GenericRepair(Exception): + pass + +class MultipleInstallationsError(Exception): + pass + +class NotInstalledError(Exception): + pass + +class NoProfilesError(Exception): + pass + +class PathNotFoundException(Exception): + pass + +class UnsupportedOSError(Exception): + pass + +# General functions +def abort(): + """Abort script.""" + print_warning('Aborted.') + sleep(5) + exit_script() + +def ask(prompt='Kotaero!'): + """Prompt the user with a Y/N question, log answer, and return a bool.""" + answer = None + prompt = '{} [Y/N]: '.format(prompt) + while answer is None: + tmp = input(prompt) + if re.search(r'^y(es|)$', tmp, re.IGNORECASE): + answer = True + elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE): + answer = False + message = '{prompt}{answer_text}'.format( + prompt = prompt, + answer_text = 'Yes' if answer else 'No') + print_log(message=message) + return answer + +def convert_to_bytes(size): + """Convert human-readable size str to bytes and return an int.""" + size = str(size) + tmp = re.search(r'(\d+)\s+([KMGT]B)', size.upper()) + if tmp: + size = int(tmp.group(1)) + units = tmp.group(2) + if units == 'TB': + size *= 1099511627776 + elif units == 'GB': + size *= 1073741824 + elif units == 'MB': + size *= 1048576 + elif units == 'KB': + size *= 1024 + else: + return -1 + + return size + +def exit_script(return_value=0): + """Exits the script after some cleanup and opens the log (if set).""" + # Remove dirs (if empty) + for dir in ['BackupDir', 'LogDir', 'TmpDir']: + try: + dir = global_vars[dir] + os.rmdir(dir) + except Exception: + pass + + # Open Log (if it exists) + log = global_vars.get('LogFile', '') + if log and os.path.exists(log): + try: + extract_item('NotepadPlusPlus', silent=True) + popen_program( + [global_vars['Tools']['NotepadPlusPlus'], + global_vars['LogFile']]) + except Exception: + print_error('ERROR: Failed to extract Notepad++ and open log.') + pause('Press Enter to exit...') + + # Kill Caffeine if still running + kill_process('caffeine.exe') + + # Exit + sys.exit(return_value) + +def extract_item(item, filter='', silent=False): + """Extract item from .cbin into .bin.""" + cmd = [ + global_vars['Tools']['SevenZip'], 'x', '-aos', '-bso0', '-bse0', + '-p{ArchivePassword}'.format(**global_vars), + r'-o{BinDir}\{item}'.format(item=item, **global_vars), + r'{CBinDir}\{item}.7z'.format(item=item, **global_vars), + filter] + if not silent: + print_standard('Extracting "{item}"...'.format(item=item)) + try: + run_program(cmd) + except subprocess.CalledProcessError: + if not silent: + print_warning('WARNING: Errors encountered while exctracting data') + +def get_ticket_number(): + """Get TicketNumber from user, save in LogDir, and return as str.""" + ticket_number = None + while ticket_number is None: + _input = input('Enter ticket number: ') + if re.match(r'^([0-9]+([-_]?\w+|))$', _input): + ticket_number = _input + with open(r'{LogDir}\TicketNumber'.format(**global_vars), 'w') as f: + f.write(ticket_number) + return ticket_number + +def human_readable_size(size, decimals=0): + """Convert size in bytes to a human-readable format and return a str.""" + # Prep string formatting + width = 3+decimals + if decimals > 0: + width += 1 + + # Convert size to int + try: + size = int(size) + except ValueError: + size = convert_to_bytes(size) + + # Verify we have a valid size + if size <= 0: + return '{size:>{width}} b'.format(size='???', width=width) + + # Convert to sensible units + if size >= 1099511627776: + size /= 1099511627776 + units = 'Tb' + elif size >= 1073741824: + size /= 1073741824 + units = 'Gb' + elif size >= 1048576: + size /= 1048576 + units = 'Mb' + elif size >= 1024: + size /= 1024 + units = 'Kb' + else: + units = ' b' + + # Return + return '{size:>{width}.{decimals}f} {units}'.format( + size=size, width=width, decimals=decimals, units=units) + +def kill_process(name): + """Kill any running caffeine.exe processes.""" + for proc in psutil.process_iter(): + if proc.name() == name: + proc.kill() + +def major_exception(): + """Display traceback and exit""" + print_error('Major exception') + print_warning(SUPPORT_MESSAGE) + print(traceback.format_exc()) + print_log(traceback.format_exc()) + sleep(30) + pause('Press Enter to exit...') + exit_script(1) + +def menu_select(title='~ Untitled Menu ~', + prompt='Please make a selection', secret_exit=False, + main_entries=[], action_entries=[], disabled_label='DISABLED'): + """Display options in a menu and return selected option as a str.""" + # Bail early + if not main_entries and not action_entries: + raise Exception("MenuError: No items given") + + # Build menu + menu_splash = '{}\n\n'.format(title) + width = len(str(len(main_entries))) + valid_answers = [] + if (secret_exit): + valid_answers.append('Q') + + # Add main entries + for i in range(len(main_entries)): + entry = main_entries[i] + # Add Spacer + if ('CRLF' in entry): + menu_splash += '\n' + entry_str = '{number:>{width}}: {name}'.format( + number = i+1, + width = width, + name = entry.get('Display Name', entry['Name'])) + if entry.get('Disabled', False): + entry_str = '{YELLOW}{entry_str} ({disabled}){CLEAR}'.format( + entry_str = entry_str, + disabled = disabled_label, + **COLORS) + else: + valid_answers.append(str(i+1)) + menu_splash += '{}\n'.format(entry_str) + menu_splash += '\n' + + # Add action entries + for entry in action_entries: + # Add Spacer + if ('CRLF' in entry): + menu_splash += '\n' + valid_answers.append(entry['Letter']) + menu_splash += '{letter:>{width}}: {name}\n'.format( + letter = entry['Letter'].upper(), + width = len(str(len(action_entries))), + name = entry['Name']) + menu_splash += '\n' + + answer = '' + + while (answer.upper() not in valid_answers): + os.system('cls') + print(menu_splash) + answer = input('{}: '.format(prompt)) + + return answer.upper() + +def non_clobber_rename(full_path): + """Append suffix to path, if necessary, to avoid clobbering path""" + new_path = full_path + _i = 1; + while os.path.exists(new_path): + new_path = '{path}_{i}'.format(i=_i, path=full_path) + _i += 1 + + return new_path + +def pause(prompt='Press Enter to continue... '): + """Simple pause implementation.""" + input(prompt) + +def ping(addr='google.com'): + """Attempt to ping addr.""" + cmd = ['ping', '-n', '2', addr] + run_program(cmd) + +def popen_program(cmd, pipe=False, minimized=False, shell=False, **kwargs): + """Run program and return a subprocess.Popen object.""" + startupinfo=None + if minimized: + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW + startupinfo.wShowWindow = 6 + + if pipe: + popen_obj = subprocess.Popen(cmd, shell=shell, startupinfo=startupinfo, + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + else: + popen_obj = subprocess.Popen(cmd, shell=shell, startupinfo=startupinfo) + + return popen_obj + +def print_error(*args, **kwargs): + """Prints message to screen in RED.""" + print_standard(*args, color=COLORS['RED'], **kwargs) + +def print_info(*args, **kwargs): + """Prints message to screen in BLUE.""" + print_standard(*args, color=COLORS['BLUE'], **kwargs) + +def print_standard(message='Generic info', + color=None, end='\n', timestamp=True, **kwargs): + """Prints message to screen and log (if set).""" + display_message = message + if color: + display_message = color + message + COLORS['CLEAR'] + # **COLORS is used below to support non-"standard" color printing + print(display_message.format(**COLORS), end=end, **kwargs) + print_log(message, end, timestamp) + +def print_success(*args, **kwargs): + """Prints message to screen in GREEN.""" + print_standard(*args, color=COLORS['GREEN'], **kwargs) + +def print_warning(*args, **kwargs): + """Prints message to screen in YELLOW.""" + print_standard(*args, color=COLORS['YELLOW'], **kwargs) + +def print_log(message='', end='\n', timestamp=True): + time_str = time.strftime("%Y-%m-%d %H%M%z: ") if timestamp else '' + if 'LogFile' in global_vars and global_vars['LogFile'] is not None: + with open(global_vars['LogFile'], 'a') as f: + for line in message.splitlines(): + f.write('{timestamp}{line}{end}'.format( + timestamp = time_str, + line = line, + end = end)) + +def run_program(cmd, args=[], check=True, pipe=True, shell=False): + """Run program and return a subprocess.CompletedProcess object.""" + if args: + # Deprecated so let's raise an exception to find & fix all occurances + print_error('ERROR: Using args is no longer supported.') + raise Exception + cmd = [c for c in cmd if c] + if shell: + cmd = ' '.join(cmd) + + if pipe: + process_return = subprocess.run(cmd, check=check, shell=shell, + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + else: + process_return = subprocess.run(cmd, check=check, shell=shell) + + return process_return + +def show_info(message='~Some message~', info='~Some info~', indent=8, width=32): + """Display info with formatting.""" + print_standard('{indent}{message:<{width}}{info}'.format( + indent=' '*indent, width=width, message=message, info=info)) + +def sleep(seconds=2): + """Wait for a while.""" + time.sleep(seconds) + +def stay_awake(): + """Prevent the system from sleeping or hibernating.""" + # Bail if caffeine is already running + for proc in psutil.process_iter(): + if proc.name() == 'caffeine.exe': + return + # Extract and run + extract_item('Caffeine', silent=True) + try: + popen_program(global_vars['Tools']['Caffeine']) + except Exception: + print_error('ERROR: No caffeine available.') + print_warning('Please set the power setting to High Performance.') + +def get_exception(s): + """Get exception by name, returns Exception object.""" + return getattr(sys.modules[__name__], s) + +def try_and_print(message='Trying...', + function=None, cs='CS', ns='NS', other_results={}, + catch_all=True, print_return=False, silent_function=True, + indent=8, width=32, *args, **kwargs): + """Run function, print if successful or not, and return dict. + + other_results is in the form of + { + 'Warning': {'ExceptionClassName': 'Result Message'}, + 'Error': {'ExceptionClassName': 'Result Message'} + } + The the ExceptionClassNames will be excepted conditions + and the result string will be printed in the correct color. + catch_all=False will result in unspecified exceptions being re-raised.""" + err = 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() + e_exceptions = tuple(get_exception(e) for e in e_exceptions) + w_results = other_results.get('Warning', {}) + e_results = other_results.get('Error', {}) + + # Run function and catch errors + print_standard('{indent}{message:<{width}}'.format( + indent=' '*indent, message=message, width=width), end='', flush=True) + try: + out = function(*args, **kwargs) + if print_return: + print_standard(out[0], timestamp=False) + for item in out[1:]: + print_standard('{indent}{item}'.format( + indent=' '*(indent+width), item=item)) + elif silent_function: + print_success(cs, timestamp=False) + except w_exceptions as e: + _result = w_results.get(e.__class__.__name__, 'Warning') + print_warning(_result, timestamp=False) + err = e + except e_exceptions as e: + _result = e_results.get(e.__class__.__name__, 'Error') + print_error(_result, timestamp=False) + err = e + except Exception: + print_error(ns, timestamp=False) + err = traceback.format_exc() + + # Return or raise? + if bool(err) and not catch_all: + raise + else: + return {'CS': not bool(err), 'Error': err} + +def upload_data(path, file): + """Add CLIENT_INFO_SERVER to authorized connections and upload file.""" + if not ENABLED_UPLOAD_DATA: + raise GenericError('Feature disabled.') + + extract_item('PuTTY', filter='wizkit.ppk psftp.exe', silent=True) + + # Authorize connection to the server + winreg.CreateKey(HKCU, r'Software\SimonTatham\PuTTY\SshHostKeys') + with winreg.OpenKey(HKCU, r'Software\SimonTatham\PuTTY\SshHostKeys', + access=winreg.KEY_WRITE) as key: + winreg.SetValueEx(key, + 'rsa2@22:{IP}'.format(**CLIENT_INFO_SERVER), 0, + winreg.REG_SZ, CLIENT_INFO_SERVER['RegEntry']) + + # Write batch file + with open(r'{TmpDir}\psftp.batch'.format(**global_vars), + 'w', encoding='ascii') as f: + f.write('lcd "{path}"\n'.format(path=path)) + f.write('cd "{Share}"\n'.format(**CLIENT_INFO_SERVER)) + f.write('mkdir {TicketNumber}\n'.format(**global_vars)) + f.write('cd {TicketNumber}\n'.format(**global_vars)) + f.write('put "{file}"\n'.format(file=file)) + + # Upload Info + cmd = [ + global_vars['Tools']['PuTTY-PSFTP'], + '-noagent', + '-i', r'{BinDir}\PuTTY\wizkit.ppk'.format(**global_vars), + '{User}@{IP}'.format(**CLIENT_INFO_SERVER), + '-b', r'{TmpDir}\psftp.batch'.format(**global_vars)] + run_program(cmd) + +def upload_info(): + """Upload compressed Info file to the NAS as set in settings.main.py.""" + if not ENABLED_UPLOAD_DATA: + raise GenericError('Feature disabled.') + + path = '{ClientDir}'.format(**global_vars) + file = 'Info_{Date-Time}.7z'.format(**global_vars) + upload_data(path, file) + +def compress_info(): + """Compress ClientDir info folders with 7-Zip for upload_info().""" + path = '{ClientDir}'.format(**global_vars) + file = 'Info_{Date-Time}.7z'.format(**global_vars) + _cmd = [ + global_vars['Tools']['SevenZip'], + 'a', '-t7z', '-mx=9', '-bso0', '-bse0', + r'{}\{}'.format(path, file), + r'{ClientDir}\Info'.format(**global_vars)] + run_program(_cmd) + +def wait_for_process(name, poll_rate=3): + """Wait for process by name.""" + running = True + while running: + sleep(poll_rate) + running = False + for proc in psutil.process_iter(): + if re.search(r'^{}'.format(name), proc.name(), re.IGNORECASE): + running = True + sleep(1) + +# global_vars functions +def init_global_vars(): + """Sets global variables based on system info.""" + print_info('Initializing') + os.system('title Wizard Kit') + init_functions = [ + ['Checking .bin...', find_bin], + ['Checking environment...', set_common_vars], + ['Checking OS...', check_os], + ['Checking tools...', check_tools], + ['Creating folders...', make_tmp_dirs], + ['Clearing collisions...', clean_env_vars], + ] + try: + for f in init_functions: + try_and_print( + message=f[0], function=f[1], + cs='Done', ns='Error', catch_all=False) + except: + major_exception() + +def check_os(): + """Set OS specific variables.""" + tmp = {} + + # Query registry + _reg_path = winreg.OpenKey( + HKLM, r'SOFTWARE\Microsoft\Windows NT\CurrentVersion') + for key in ['CSDVersion', 'CurrentBuild', 'CurrentBuildNumber', + 'CurrentVersion', 'ProductName']: + try: + tmp[key] = winreg.QueryValueEx(_reg_path, key)[0] + if key in ['CurrentBuild', 'CurrentBuildNumber']: + tmp[key] = int(tmp[key]) + except ValueError: + # Couldn't convert Build to int so this should be interesting... + tmp[key] = 0 + except Exception: + tmp[key] = 'Unknown' + + # Determine OS bit depth + tmp['Arch'] = 32 + if 'PROGRAMFILES(X86)' in global_vars['Env']: + tmp['Arch'] = 64 + + # Determine OS Name + tmp['Name'] = '{ProductName} {CSDVersion}'.format(**tmp) + if tmp['CurrentBuild'] == 9600: + tmp['Name'] += ' Update' # Win 8.1u + if tmp['CurrentBuild'] == 10240: + tmp['Name'] += ' Release 1507 "Threshold 1"' + if tmp['CurrentBuild'] == 10586: + tmp['Name'] += ' Release 1511 "Threshold 2"' + if tmp['CurrentBuild'] == 14393: + tmp['Name'] += ' Release 1607 "Redstone 1" / "Anniversary Update"' + if tmp['CurrentBuild'] == 15063: + tmp['Name'] += ' Release 1703 "Redstone 2" / "Creators Update"' + if tmp['CurrentBuild'] == 16299: + tmp['Name'] += ' Release 1709 "Redstone 3" / "Fall Creators Update"' + tmp['Name'] = tmp['Name'].replace('Service Pack ', 'SP') + tmp['Name'] = tmp['Name'].replace('Unknown Release', 'Release') + tmp['Name'] = re.sub(r'\s+', ' ', tmp['Name']) + + # Determine OS version + name = '{Name} x{Arch}'.format(**tmp) + if tmp['CurrentVersion'] == '6.0': + tmp['Version'] = 'Vista' + name += ' (very outdated)' + elif tmp['CurrentVersion'] == '6.1': + tmp['Version'] = '7' + if tmp['CSDVersion'] == 'Service Pack 1': + name += ' (outdated)' + else: + name += ' (very outdated)' + elif tmp['CurrentVersion'] in ['6.2', '6.3']: + if int(tmp['CurrentBuildNumber']) <= 9600: + tmp['Version'] = '8' + elif int(tmp['CurrentBuildNumber']) >= 10240: + tmp['Version'] = '10' + if tmp['CurrentBuild'] in [9200, 10240, 10586]: + name += ' (very outdated)' + elif tmp['CurrentBuild'] in [9600, 14393, 15063]: + name += ' (outdated)' + elif tmp['CurrentBuild'] == 16299: + pass # Current Win10 + else: + name += ' (unrecognized)' + tmp['DisplayName'] = name + + # == vista == + # 6.0.6000 + # 6.0.6001 + # 6.0.6002 + # ==== 7 ==== + # 6.1.7600 + # 6.1.7601 + # 6.1.7602 + # ==== 8 ==== + # 6.2.9200 + # === 8.1 === + # 6.3.9200 + # === 8.1u == + # 6.3.9600 + # === 10 v1507 "Threshold 1" == + # 6.3.10240 + # === 10 v1511 "Threshold 2" == + # 6.3.10586 + # === 10 v1607 "Redstone 1" "Anniversary Update" == + # 6.3.14393 + # === 10 v1703 "Redstone 2" "Creators Update" == + # 6.3.15063 + # === 10 v1709 "Redstone 3" "Fall Creators Update" == + # 6.3.16299 + global_vars['OS'] = tmp + +def check_tools(): + """Set tool variables based on OS bit-depth and tool availability.""" + if global_vars['OS'].get('Arch', 32) == 64: + global_vars['Tools'] = { + k: v.get('64', v.get('32')) for (k, v) in TOOLS.items()} + else: + global_vars['Tools'] = {k: v.get('32') for (k, v) in TOOLS.items()} + + # Fix paths + global_vars['Tools'] = {k: os.path.join(global_vars['BinDir'], v) + for (k, v) in global_vars['Tools'].items()} + +def clean_env_vars(): + """Remove conflicting global_vars and env variables. + + This fixes an issue where both global_vars and + global_vars['Env'] are expanded at the same time.""" + for key in global_vars.keys(): + global_vars['Env'].pop(key, None) + +def find_bin(): + """Find .bin folder in the cwd or it's parents.""" + wd = os.getcwd() + base = None + while base is None: + if os.path.exists('.bin'): + base = os.getcwd() + break + if re.fullmatch(r'\w:\\', os.getcwd()): + break + os.chdir('..') + os.chdir(wd) + if base is None: + raise BinNotFoundError + global_vars['BaseDir'] = base + +def make_tmp_dirs(): + """Make temp directories.""" + os.makedirs(global_vars['BackupDir'], exist_ok=True) + os.makedirs(global_vars['LogDir'], exist_ok=True) + os.makedirs(global_vars['TmpDir'], exist_ok=True) + +def set_common_vars(): + """Set common variables.""" + global_vars['Date'] = time.strftime("%Y-%m-%d") + global_vars['Date-Time'] = time.strftime("%Y-%m-%d_%H%M_%z") + global_vars['Env'] = os.environ.copy() + + global_vars['ArchivePassword'] = ARCHIVE_PASSWORD + global_vars['BinDir'] = r'{BaseDir}\.bin'.format( + **global_vars) + global_vars['CBinDir'] = r'{BaseDir}\.cbin'.format( + **global_vars) + global_vars['ClientDir'] = r'{SYSTEMDRIVE}\{prefix}'.format( + prefix=KIT_NAME_SHORT, **global_vars['Env']) + global_vars['BackupDir'] = r'{ClientDir}\Backups\{Date}'.format( + **global_vars) + global_vars['LogDir'] = r'{ClientDir}\Info\{Date}'.format( + **global_vars) + global_vars['ProgBackupDir'] = r'{ClientDir}\Backups'.format( + **global_vars) + global_vars['QuarantineDir'] = r'{ClientDir}\Quarantine'.format( + **global_vars) + global_vars['TmpDir'] = r'{BinDir}\tmp'.format( + **global_vars) + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/Scripts/functions/data.py b/Scripts/functions/data.py new file mode 100644 index 00000000..66797329 --- /dev/null +++ b/Scripts/functions/data.py @@ -0,0 +1,607 @@ +# Wizard Kit PE: Functions - Data + +import ctypes + +from operator import itemgetter + +from functions.common import * + +# Classes +class LocalDisk(): + def __init__(self, disk): + self.disk = disk + self.name = disk.mountpoint.upper() + self.path = self.name + def is_dir(self): + # Should always be true + return True + +# Regex +REGEX_EXCL_ITEMS = re.compile( + r'^(\.(AppleDB|AppleDesktop|AppleDouble' + r'|com\.apple\.timemachine\.supported|dbfseventsd' + r'|DocumentRevisions-V100.*|DS_Store|fseventsd|PKInstallSandboxManager' + r'|Spotlight.*|SymAV.*|symSchedScanLockxz|TemporaryItems|Trash.*' + r'|vol|VolumeIcon\.icns)|desktop\.(ini|.*DB|.*DF)' + r'|(hiberfil|pagefile)\.sys|lost\+found|Network\.*Trash\.*Folder' + r'|Recycle[dr]|System\.*Volume\.*Information|Temporary\.*Items' + r'|Thumbs\.db)$', + re.IGNORECASE) +REGEX_EXCL_ROOT_ITEMS = re.compile( + r'^\\?(boot(mgr|nxt)$|Config.msi' + r'|(eula|globdata|install|vc_?red)' + r'|.*.sys$|System Volume Information|RECYCLER?|\$Recycle\.bin' + r'|\$?Win(dows(.old.*|\.~BT|)$|RE_)|\$GetCurrent|Windows10Upgrade' + r'|PerfLogs|Program Files|SYSTEM.SAV' + r'|.*\.(esd|swm|wim|dd|map|dmg|image)$)', + re.IGNORECASE) +REGEX_INCL_ROOT_ITEMS = re.compile( + r'^\\?(AdwCleaner|(My\s*|)(Doc(uments?( and Settings|)|s?)|Downloads' + r'|Media|Music|Pic(ture|)s?|Vid(eo|)s?)' + r'|{prefix}(-?Info|-?Transfer|)' + r'|(ProgramData|Recovery|Temp.*|Users)$' + r'|.*\.(log|txt|rtf|qb\w*|avi|m4a|m4v|mp4|mkv|jpg|png|tiff?)$)' + r''.format(prefix=KIT_NAME_SHORT), + re.IGNORECASE) +REGEX_WIM_FILE = re.compile( + r'\.wim$', + re.IGNORECASE) +REGEX_WINDOWS_OLD = re.compile( + r'^\\Win(dows|)\.old', + re.IGNORECASE) + +# STATIC VARIABLES +FAST_COPY_EXCLUDES = [ + r'\*.esd', + r'\*.swm', + r'\*.wim', + r'\*.dd', + r'\*.dd.tgz', + r'\*.dd.txz', + r'\*.map', + r'\*.dmg', + r'\*.image', + r'$RECYCLE.BIN', + r'$Recycle.Bin', + r'.AppleDB', + r'.AppleDesktop', + r'.AppleDouble', + r'.com.apple.timemachine.supported', + r'.dbfseventsd', + r'.DocumentRevisions-V100*', + r'.DS_Store', + r'.fseventsd', + r'.PKInstallSandboxManager', + r'.Spotlight*', + r'.SymAV*', + r'.symSchedScanLockxz', + r'.TemporaryItems', + r'.Trash*', + r'.vol', + r'.VolumeIcon.icns', + r'desktop.ini', + r'Desktop?DB', + r'Desktop?DF', + r'hiberfil.sys', + r'lost+found', + r'Network?Trash?Folder', + r'pagefile.sys', + r'Recycled', + r'RECYCLER', + r'System?Volume?Information', + r'Temporary?Items', + r'Thumbs.db', + ] +FAST_COPY_ARGS = [ + '/cmd=noexist_only', + '/utf8', + '/skip_empty_dir', + '/linkdest', + '/no_ui', + '/auto_close', + '/exclude={}'.format(';'.join(FAST_COPY_EXCLUDES)), + ] +# Code borrowed from: https://stackoverflow.com/a/29075319 +SEM_NORMAL = ctypes.c_uint() +SEM_FAILCRITICALERRORS = 1 +SEM_NOOPENFILEERRORBOX = 0x8000 +SEM_FAIL = SEM_NOOPENFILEERRORBOX | SEM_FAILCRITICALERRORS + +def cleanup_transfer(dest_path): + """Fix attributes and move extraneous items outside the Transfer folder.""" + try: + # Remove dest_path if empty + os.rmdir(dest_path) + except OSError: + pass + if not os.path.exists(dest_path): + # Bail if dest_path was empty and removed + raise Exception + + # Fix attributes + cmd = ['attrib', '-a', '-h', '-r', '-s', dest_path] + run_program(cmd, check=False) + + for root, dirs, files in os.walk(dest_path, topdown=False): + for name in dirs: + # Remove empty directories and junction points + try: + os.rmdir(os.path.join(root, name)) + except OSError: + pass + for name in files: + # "Remove" files based on exclusion regex + if REGEX_EXCL_ITEMS.search(name): + # Make dest folder + dest_name = root.replace(dest_path, dest_path+'.Removed') + os.makedirs(dest_name, exist_ok=True) + # Set dest filename + dest_name = os.path.join(dest_name, name) + dest_name = non_clobber_rename(dest_name) + source_name = os.path.join(root, name) + try: + shutil.move(source_name, dest_name) + except Exception: + pass + +def is_valid_wim_file(item): + """Checks if the provided os.DirEntry is a valid WIM file, returns bool.""" + valid = bool(item.is_file() and REGEX_WIM_FILE.search(item.name)) + if valid: + extract_item('wimlib', silent=True) + cmd = [global_vars['Tools']['wimlib-imagex'], 'info', item.path] + try: + run_program(cmd) + except subprocess.CalledProcessError: + valid = False + print_log('WARNING: Image "{}" damaged.'.format(item.name)) + return valid + +def mount_backup_shares(): + """Mount the backup shares unless labeled as already mounted.""" + for server in BACKUP_SERVERS: + # Blindly skip if we mounted earlier + if server['Mounted']: + continue + + mount_network_share(server) + +def mount_network_share(server): + """Mount a network share defined by server.""" + # Test connection + try: + ping(server['IP']) + except subprocess.CalledProcessError: + print_error( + r'Failed to mount \\{Name}\{Share}, {IP} unreachable.'.format( + **server)) + sleep(1) + return False + + # Mount + cmd = r'net use \\{IP}\{Share} /user:{User} {Pass}'.format(**server) + cmd = cmd.split(' ') + try: + run_program(cmd) + except Exception: + print_warning(r'Failed to mount \\{Name}\{Share} ({IP})'.format( + **server)) + sleep(1) + else: + print_info('Mounted {Name}'.format(**server)) + server['Mounted'] = True + +def run_fast_copy(items, dest): + """Copy items to dest using FastCopy.""" + if not items: + raise Exception + + cmd = [global_vars['Tools']['FastCopy'], *FAST_COPY_ARGS] + if 'LogFile' in global_vars: + cmd.append('/logfile={LogFile}'.format(**global_vars)) + cmd.extend(items) + cmd.append('/to={}\\'.format(dest)) + + run_program(cmd) + +def run_wimextract(source, items, dest): + """Extract items from source WIM to dest folder.""" + if not items: + raise Exception + extract_item('wimlib', silent=True) + + # Write files.txt + with open(r'{TmpDir}\wim_files.txt'.format(**global_vars), 'w') as f: + # Defaults + for item in items: + f.write('{item}\n'.format(item=item)) + sleep(1) # For safety? + + # Extract files + cmd = [ + global_vars['Tools']['wimlib-imagex'], + 'extract', + source, '1', + r'@{TmpDir}\wim_files.txt'.format(**global_vars), + '--dest-dir={}\\'.format(dest), + '--no-acls', + '--nullglob'] + run_program(cmd) + +def scan_source(source_obj, dest_path): + """Scan source for files/folders to transfer.""" + selected_items = [] + + if source_obj.is_dir(): + # File-Based + print_standard('Scanning source (folder): {}'.format(source_obj.path)) + selected_items = scan_source_path(source_obj.path, dest_path) + else: + # Image-Based + if REGEX_WIM_FILE.search(source_obj.name): + print_standard('Scanning source (image): {}'.format( + source_obj.path)) + selected_items = scan_source_wim(source_obj.path, dest_path) + else: + print_error('ERROR: Unsupported image: {}'.format( + source_obj.path)) + raise GenericError + + return selected_items + +def scan_source_path(source_path, dest_path, rel_path=None, interactive=True): + """Scan source folder for files/folders to transfer, returns list. + + This will scan the root and (recursively) any Windows.old folders.""" + rel_path = '\\' + rel_path if rel_path else '' + if rel_path: + dest_path = dest_path + rel_path + selected_items = [] + win_olds = [] + + # Root items + root_items = [] + for item in os.scandir(source_path): + if REGEX_INCL_ROOT_ITEMS.search(item.name): + root_items.append(item.path) + elif not REGEX_EXCL_ROOT_ITEMS.search(item.name): + if (not interactive + or ask('Copy: "{}{}" ?'.format(rel_path, item.name))): + root_items.append(item.path) + if REGEX_WINDOWS_OLD.search(item.name): + win_olds.append(item) + if root_items: + selected_items.append({ + 'Message': '{}Root Items...'.format(rel_path), + 'Items': root_items.copy(), + 'Destination': dest_path}) + + # Fonts + if os.path.exists(r'{}\Windows\Fonts'.format(source_path)): + selected_items.append({ + 'Message': '{}Fonts...'.format(rel_path), + 'Items': [r'{}\Windows\Fonts'.format(rel_path)], + 'Destination': r'{}\Windows'.format(dest_path)}) + + # Registry + registry_items = [] + for folder in ['config', 'OEM']: + folder = r'Windows\System32\{}'.format(folder) + folder = os.path.join(source_path, folder) + if os.path.exists(folder): + registry_items.append(folder) + if registry_items: + selected_items.append({ + 'Message': '{}Registry...'.format(rel_path), + 'Items': registry_items.copy(), + 'Destination': r'{}\Windows\System32'.format(dest_path)}) + + # Windows.old(s) + for old in win_olds: + selected_items.append( + scan_source_path( + old.path, dest_path, rel_path=old.name, interactive=False)) + + # Done + return selected_items + +def scan_source_wim(source_wim, dest_path, rel_path=None, interactive=True): + """Scan source WIM file for files/folders to transfer, returns list. + + This will scan the root and (recursively) any Windows.old folders.""" + rel_path = '\\' + rel_path if rel_path else '' + selected_items = [] + win_olds = [] + + # Scan source + extract_item('wimlib', silent=True) + cmd = [ + global_vars['Tools']['wimlib-imagex'], 'dir', + source_wim, '1'] + try: + file_list = run_program(cmd) + except subprocess.CalledProcessError: + print_error('ERROR: Failed to get file list.') + raise + + # Root Items + file_list = [i.strip() + for i in file_list.stdout.decode('utf-8', 'ignore').splitlines() + if i.count('\\') == 1 and i.strip() != '\\'] + root_items = [] + if rel_path: + file_list = [i.replace(rel_path, '') for i in file_list] + for item in file_list: + if REGEX_INCL_ROOT_ITEMS.search(item): + root_items.append(item) + elif not REGEX_EXCL_ROOT_ITEMS.search(item): + if (not interactive + or ask('Extract: "{}{}" ?'.format(rel_path, item))): + root_items.append('{}{}'.format(rel_path, item)) + if REGEX_WINDOWS_OLD.search(item): + win_olds.append(item) + if root_items: + selected_items.append({ + 'Message': '{}Root Items...'.format(rel_path), + 'Items': root_items.copy(), + 'Destination': dest_path}) + + # Fonts + if wim_contains(source_wim, r'{}Windows\Fonts'.format(rel_path)): + selected_items.append({ + 'Message': '{}Fonts...'.format(rel_path), + 'Items': [r'{}\Windows\Fonts'.format(rel_path)], + 'Destination': dest_path}) + + # Registry + registry_items = [] + for folder in ['config', 'OEM']: + folder = r'{}Windows\System32\{}'.format(rel_path, folder) + if wim_contains(source_wim, folder): + registry_items.append(folder) + if registry_items: + selected_items.append({ + 'Message': '{}Registry...'.format(rel_path), + 'Items': registry_items.copy(), + 'Destination': dest_path}) + + # Windows.old(s) + for old in win_olds: + scan_source_wim(source_wim, dest_path, rel_path=old, interactive=False) + + # Done + return selected_items + +def select_destination(folder_path, prompt='Select destination'): + """Select destination drive, returns path as string.""" + disk = select_disk(prompt) + if 'fixed' not in disk['Disk'].opts: + folder_path = folder_path.replace('\\', '-') + path = '{disk}{folder_path}_{Date}'.format( + disk = disk['Disk'].mountpoint, + folder_path = folder_path, + **global_vars) + + # Avoid merging with existing folder + path = non_clobber_rename(path) + os.makedirs(path, exist_ok=True) + + return path + +def select_disk(title='Select disk', auto_select=True): + """Select disk from attached disks. returns dict.""" + actions = [{'Name': 'Quit', 'Letter': 'Q'}] + disks = [] + + # Build list of disks + set_thread_error_mode(silent=True) # Prevents "No disk" popups + for d in psutil.disk_partitions(): + info = { + 'Disk': d, + 'Name': d.mountpoint} + try: + usage = psutil.disk_usage(d.device) + free = '{free} / {total} available'.format( + free = human_readable_size(usage.free, 2), + total = human_readable_size(usage.total, 2)) + except Exception: + # Meh, leaving unsupported destinations out + pass + # free = 'Unknown' + # info['Disabled'] = True + else: + info['Display Name'] = '{} ({})'.format(info['Name'], free) + disks.append(info) + set_thread_error_mode(silent=False) # Return to normal + + # Skip menu? + if len(disks) == 1 and auto_select: + return disks[0] + + # Show menu + selection = menu_select(title, main_entries=disks, action_entries=actions) + if selection == 'Q': + exit_script() + else: + return disks[int(selection)-1] + +def select_source(ticket_number): + """Select backup from those found on the BACKUP_SERVERS for the ticket.""" + selected_source = None + sources = [] + mount_backup_shares() + + # Check for ticket 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())): + # Add folder to sources + sources.append({ + 'Name': '{:9}| File-Based: [DIR] {}'.format( + server['Name'], d.name), + 'Server': server, + 'Source': d}) + + # Check for images and subfolders + for ticket_path in sources.copy(): + for item in os.scandir(ticket_path['Source'].path): + if item.is_dir(): + # Add folder to sources + sources.append({ + 'Name': r'{:9}| File-Based: [DIR] {}\{}'.format( + ticket_path['Server']['Name'], # Server + ticket_path['Source'].name, # Ticket folder + item.name, # Sub-folder + ), + 'Server': ticket_path['Server'], + 'Source': item}) + + # Check for images in folder + for subitem in os.scandir(item.path): + if REGEX_WIM_FILE.search(item.name): + # Add image to sources + try: + size = human_readable_size(item.stat().st_size) + except Exception: + size = ' ? ?' # unknown + sources.append({ + 'Disabled': bool(not is_valid_wim_file(subitem)), + 'Name': r'{:9}| Image-Based: {:>7} {}\{}\{}'.format( + ticket_path['Server']['Name'], # Server + size, # Size (duh) + ticket_path['Source'].name, # Ticket folder + item.name, # Sub-folder + subitem.name, # Image file + ), + 'Server': ticket_path['Server'], + 'Source': subitem}) + elif REGEX_WIM_FILE.search(item.name): + # Add image to sources + try: + size = human_readable_size(item.stat().st_size) + except Exception: + size = ' ? ?' # unknown + sources.append({ + 'Disabled': bool(not is_valid_wim_file(item)), + 'Name': r'{:9}| Image-Based: {:>7} {}\{}'.format( + ticket_path['Server']['Name'], # Server + size, # Size (duh) + ticket_path['Source'].name, # Ticket folder + item.name, # Image file + ), + 'Server': ticket_path['Server'], + 'Source': item}) + # Check for local sources + print_standard('Scanning for local sources...') + set_thread_error_mode(silent=True) # Prevents "No disk" popups + sys_drive = global_vars['Env']['SYSTEMDRIVE'] + for d in psutil.disk_partitions(): + if re.search(r'^{}'.format(sys_drive), d.mountpoint, re.IGNORECASE): + # Skip current OS drive + continue + if 'fixed' in d.opts: + # Skip DVD, etc + sources.append({ + 'Name': '{:9}| File-Based: [DISK] {}'.format( + ' Local', d.mountpoint), + 'Source': LocalDisk(d)}) + set_thread_error_mode(silent=False) # Return to normal + + # Build Menu + sources.sort(key=itemgetter('Name')) + actions = [{'Name': 'Quit', 'Letter': 'Q'}] + + # Select backup from sources + if len(sources) > 0: + selection = menu_select( + 'Which backup are we using?', + main_entries=sources, + action_entries=actions, + disabled_label='DAMAGED') + if selection == 'Q': + umount_backup_shares() + exit_script() + else: + selected_source = sources[int(selection)-1]['Source'] + else: + print_error('ERROR: No backups found for ticket: {}.'.format( + ticket_number)) + umount_backup_shares() + pause("Press Enter to exit...") + exit_script() + + # Done + return selected_source + +def set_thread_error_mode(silent=True): + """Disable or Enable Windows error message dialogs. + + Disable when scanning for disks to avoid popups for empty cardreaders, etc + """ + # Code borrowed from: https://stackoverflow.com/a/29075319 + kernel32 = ctypes.WinDLL('kernel32') + + if silent: + kernel32.SetThreadErrorMode(SEM_FAIL, ctypes.byref(SEM_NORMAL)) + else: + kernel32.SetThreadErrorMode(SEM_NORMAL, ctypes.byref(SEM_NORMAL)) + +def transfer_source(source_obj, dest_path, selected_items): + """Transfer, or extract, files/folders from source to destination.""" + if source_obj.is_dir(): + # Run FastCopy for each selection "group" + for group in selected_items: + try_and_print(message=group['Message'], + function=run_fast_copy, cs='Done', + items=group['Items'], + dest=group['Destination']) + else: + if REGEX_WIM_FILE.search(source_obj.name): + # Extract files from WIM + for group in selected_items: + try_and_print(message=group['Message'], + function=run_wimextract, cs='Done', + source=source_obj.path, + items=group['Items'], + dest=group['Destination']) + else: + print_error('ERROR: Unsupported image: {}'.format(source_obj.path)) + raise GenericError + +def umount_backup_shares(): + """Unnount the backup shares regardless of current status.""" + for server in BACKUP_SERVERS: + umount_network_share(server) + +def umount_network_share(server): + """Unnount a network share defined by server.""" + cmd = r'net use \\{IP}\{Share} /delete'.format(**server) + cmd = cmd.split(' ') + try: + run_program(cmd) + except Exception: + print_error(r'Failed to umount \\{Name}\{Share}.'.format(**server)) + sleep(1) + else: + print_info('Umounted {Name}'.format(**server)) + server['Mounted'] = False + +def wim_contains(source_path, file_path): + """Check if the WIM contains a file or folder.""" + _cmd = [ + global_vars['Tools']['wimlib-imagex'], 'dir', + source_path, '1', + '--path={}'.format(file_path), + '--one-file-only'] + try: + run_program(_cmd) + except subprocess.CalledProcessError: + return False + else: + return True + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/Scripts/partition_uids.py b/Scripts/functions/partition_uids.py similarity index 98% rename from Scripts/partition_uids.py rename to Scripts/functions/partition_uids.py index 1484e2af..e31117db 100644 --- a/Scripts/partition_uids.py +++ b/Scripts/functions/partition_uids.py @@ -1,4 +1,4 @@ -# WK WinPE - PARTITION UIDs +# Wizard Kit PE: Functions - PARTITION UIDs # sources: https://en.wikipedia.org/wiki/GUID_Partition_Table # https://en.wikipedia.org/wiki/Partition_type diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py new file mode 100644 index 00000000..e70e88cf --- /dev/null +++ b/Scripts/functions/windows_setup.py @@ -0,0 +1,273 @@ +# Wizard Kit PE: Functions - Windows Setup + +from functions.common import * + +# STATIC VARIABLES +DISKPART_SCRIPT = '{tmp}\\diskpart.script'.format(tmp=os.environ['TMP']) +WINDOWS_VERSIONS = [ + {'Name': 'Windows 7 Home Basic', + 'Image File': 'Win7', + 'Image Name': 'Windows 7 HOMEBASIC', + 'Family': '7'}, + {'Name': 'Windows 7 Home Premium', + 'Image File': 'Win7', + 'Image Name': 'Windows 7 HOMEPREMIUM', + 'Family': '7'}, + {'Name': 'Windows 7 Professional', + 'Image File': 'Win7', + 'Image Name': 'Windows 7 PROFESSIONAL', + 'Family': '7'}, + {'Name': 'Windows 7 Ultimate', + 'Image File': 'Win7', + 'Image Name': 'Windows 7 ULTIMATE', + 'Family': '7'}, + + {'Name': 'Windows 8.1', + 'Image File': 'Win8', + 'Image Name': 'Windows 8.1', + 'Family': '8', + 'CRLF': True}, + {'Name': 'Windows 8.1 Pro', + 'Image File': 'Win8', + 'Image Name': 'Windows 8.1 Pro', + 'Family': '8'}, + + {'Name': 'Windows 10 Home', + 'Image File': 'Win10', + 'Image Name': 'Windows 10 Home', + 'Family': '10', + 'CRLF': True}, + {'Name': 'Windows 10 Pro', + 'Image File': 'Win10', + 'Image Name': 'Windows 10 Pro', + 'Family': '10'}, + ] + +def is_index_in_image(bin=None, filename=None, imagename=None): + # Bail early + if bin is None: + raise Exception('bin not specified.') + if filename is None: + raise Exception('Filename not specified.') + if imagename is None: + raise Exception('Image Name not specified.') + + cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename) + try: + run_program(cmd) + except subprocess.CalledProcessError: + print_error('Invalid image: {filename}'.format(filename=filename)) + return False + + return True + +def find_windows_image(bin, windows_version=None): + """Search for a Windows source image file on local drives and network drives (in that order)""" + image = {} + + # Bail early + if windows_version is None: + raise Exception('Windows version not specified.') + imagefile = windows_version['Image File'] + + # Search local source + process_return = run_program('mountvol') + for tmp in re.findall(r'.*([A-Za-z]):\\', process_return.stdout.decode()): + for ext in ['esd', 'wim', 'swm']: + filename = '{drive}:\\images\\{imagefile}'.format(drive=tmp[0], imagefile=imagefile) + filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) + if os.path.isfile(filename_ext): + if is_index_in_image(bin, filename_ext, windows_version['Image Name']): + image['Ext'] = ext + image['File'] = filename + image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' + image['Source'] = tmp[0] + break + + # Check for network source (if necessary) + if not any(image): + if not WINDOWS_SERVER['Mounted']: + mount_windows_share() + for ext in ['esd', 'wim', 'swm']: + filename = '\\\\{IP}\\{Share}\\images\\{imagefile}'.format(imagefile=imagefile, **WINDOWS_SERVER) + filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) + if os.path.isfile(filename_ext): + if is_index_in_image(bin, filename_ext, windows_version['Image Name']): + image['Ext'] = ext + image['File'] = filename + image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' + image['Source'] = None + break + + # Display image to be used (if any) and return + if any(image): + print_info('Using image: {File}.{Ext}'.format(**image)) + return image + else: + print_error('Failed to find Windows source image for {winver}'.format(winver=windows_version['Name'])) + abort_to_main_menu('Aborting Windows setup') + +def format_gpt(disk=None, windows_family=None): + """Format disk for use as a Windows OS drive using the GPT (UEFI) layout.""" + + # Bail early + if disk is None: + raise Exception('No disk provided.') + if windows_family is None: + raise Exception('No Windows family provided.') + + # Format drive + # print_info('Drive will use a GPT (UEFI) layout.') + with open(DISKPART_SCRIPT, 'w') as script: + # Partition table + script.write('select disk {number}\n'.format(number=disk['Number'])) + script.write('clean\n') + script.write('convert gpt\n') + + # System partition + script.write('create partition efi size=260\n') # NOTE: Allows for Advanced Format 4K drives + script.write('format quick fs=fat32 label="System"\n') + script.write('assign letter="S"\n') + + # Microsoft Reserved (MSR) partition + script.write('create partition msr size=128\n') + + # Windows partition + script.write('create partition primary\n') + script.write('format quick fs=ntfs label="Windows"\n') + script.write('assign letter="W"\n') + + # Recovery Tools partition (Windows 8+) + if re.search(r'^(8|10)', windows_family): + script.write('shrink minimum=500\n') + script.write('create partition primary\n') + script.write('format quick fs=ntfs label="Recovery Tools"\n') + script.write('assign letter="T"\n') + script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n') + script.write('gpt attributes=0x8000000000000001\n') + + # Run script + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + time.sleep(2) + +def format_mbr(disk=None, windows_family=None): + """Format disk for use as a Windows OS drive using the MBR (legacy) layout.""" + + # Bail early + if disk is None: + raise Exception('No disk provided.') + if windows_family is None: + raise Exception('No Windows family provided.') + + # Format drive + # print_info('Drive will use a MBR (legacy) layout.') + with open(DISKPART_SCRIPT, 'w') as script: + # Partition table + script.write('select disk {number}\n'.format(number=disk['Number'])) + script.write('clean\n') + + # System partition + script.write('create partition primary size=100\n') + script.write('format fs=ntfs quick label="System Reserved"\n') + script.write('active\n') + script.write('assign letter="S"\n') + + # Windows partition + script.write('create partition primary\n') + script.write('format fs=ntfs quick label="Windows"\n') + script.write('assign letter="W"\n') + + # Recovery Tools partition (Windows 8+) + if re.search(r'^(8|10)', windows_family): + script.write('shrink minimum=500\n') + script.write('create partition primary\n') + script.write('format quick fs=ntfs label="Recovery"\n') + script.write('assign letter="T"\n') + script.write('set id=27\n') + + # Run script + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + time.sleep(2) + +def get_boot_mode(): + boot_mode = 'Legacy' + try: + reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 'System\\CurrentControlSet\\Control') + reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] + if reg_value == 2: + boot_mode = 'UEFI' + except: + boot_mode = 'Unknown' + + return boot_mode + +def mount_windows_share(): + """Mount the Windows images share for use in Windows setup""" + + # Blindly skip if we mounted earlier + if WINDOWS_SERVER['Mounted']: + return None + else: + try: + # Test connection + run_program('ping -w 800 -n 2 {IP}'.format(**WINDOWS_SERVER)) + # Mount + run_program('net use \\\\{IP}\\{Share} /user:{User} {Pass}'.format(**WINDOWS_SERVER)) + print_info('Mounted {Name}'.format(**WINDOWS_SERVER)) + WINDOWS_SERVER['Mounted'] = True + except subprocess.CalledProcessError: + print_error('Failed to mount \\\\{Name}\\{Share}, {IP} unreachable.'.format(**WINDOWS_SERVER)) + time.sleep(1) + except: + print_warning('Failed to mount \\\\{Name}\\{Share}'.format(**WINDOWS_SERVER)) + time.sleep(1) + +def select_windows_version(): + actions = [{'Name': 'Main Menu', 'Letter': 'M'},] + + # Menu loop + selection = menu_select('Which version of Windows are we installing?', WINDOWS_VERSIONS, actions) + + if selection.isnumeric(): + return WINDOWS_VERSIONS[int(selection)-1] + elif selection == 'M': + abort_to_main_menu() + +def setup_windows(bin=None, windows_image=None, windows_version=None): + # Bail early + if bin is None: + raise Exception('bin path not specified.') + if windows_image is None: + raise Exception('Windows image not specified.') + if windows_version is None: + raise Exception('Windows version not specified.') + + # Apply image + cmd = '{bin}\\wimlib\\wimlib-imagex apply "{File}.{Ext}" "{Image Name}" W:\\ {Glob}'.format(bin=bin, **windows_image, **windows_version) + run_program(cmd) + +def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T'): + # Bail early + if windows_version is None: + raise Exception('Windows version not specified.') + + _win = '{win}:\\Windows'.format(win=windows_letter) + _winre = '{win}\\System32\\Recovery\\WinRE.wim'.format(win=_win) + _dest = '{tools}:\\Recovery\\WindowsRE'.format(tools=tools_letter) + + if re.search(r'^(8|10)', windows_version['Family']): + # Copy WinRE.wim + os.makedirs(_dest, exist_ok=True) + shutil.copy(_winre, '{dest}\\WinRE.wim'.format(dest=_dest)) + + # Set location + run_program('{win}\\System32\\reagentc /setreimage /path {dest} /target {win}'.format(dest=_dest, win=_win)) + else: + # Only supported on Windows 8 and above + raise SetupError + +def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'): + run_program('bcdboot {win}:\\Windows /s {sys}: /f {mode}'.format(win=windows_letter, sys=system_letter, mode=mode)) + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/Scripts/menu.py b/Scripts/functions/winpe_menus.py similarity index 66% rename from Scripts/menu.py rename to Scripts/functions/winpe_menus.py index 7a72f729..21d2e11b 100644 --- a/Scripts/menu.py +++ b/Scripts/functions/winpe_menus.py @@ -1,38 +1,66 @@ -# WK WinPE Menu +# Wizard Kit PE: Menus -# Init -import os -import re -import subprocess -import sys -import time -import traceback -os.chdir(os.path.dirname(os.path.realpath(__file__))) -bin = os.path.abspath('..\\') -sys.path.append(os.getcwd()) -from functions import * +from functions.data import * -## Colors -COLORS = { - 'CLEAR': '\033[0m', - 'RED': '\033[31m', - 'GREEN': '\033[32m', - 'YELLOW': '\033[33m', - 'BLUE': '\033[34m'} +# STATIC VARIABLES +FAST_COPY_ARGS = [ + '/cmd=noexist_only', + '/utf8', + '/skip_empty_dir', + '/linkdest', + '/no_ui', + '/auto_close', + '/exclude={}'.format(';'.join(FAST_COPY_EXCLUDES)), + ] +PE_TOOLS = { + 'BlueScreenView': { + 'Path': r'BlueScreenView\BlueScreenView.exe', + }, + 'FastCopy': { + 'Path': r'FastCopy\FastCopy.exe', + 'Args': FAST_COPY_ARGS, + }, + 'HWiNFO': { + 'Path': r'HWiNFO\HWiNFO.exe', + }, + 'NT Password Editor': { + 'Path': r'NT Password Editor\ntpwedit.exe', + }, + 'Notepad++': { + 'Path': r'NotepadPlusPlus\NotepadPlusPlus.exe', + }, + 'PhotoRec': { + 'Path': r'TestDisk\photorec_win.exe', + 'Args': ['-new_console:n'], + }, + 'Prime95': { + 'Path': r'Prime95\prime95.exe', + }, + 'ProduKey': { + 'Path': r'ProduKey\ProduKey.exe', + }, + 'Q-Dir': { + 'Path': r'Q-Dir\Q-Dir.exe', + }, + 'TestDisk': { + 'Path': r'TestDisk\testdisk_win.exe', + 'Args': ['-new_console:n'], + }, + } -def menu_backup_imaging(): +def menu_backup(): """Take backup images of partition(s) in the WIM format and save them to a backup share""" errors = False # Set ticket ID os.system('cls') - ticket_id = get_ticket_id() + ticket_id = get_ticket_number() # Mount backup shares mount_backup_shares() # Select destination - dest = select_destination() + dest = select_backup_destination() if dest is None: abort_to_main_menu('Aborting Backup Creation') @@ -90,13 +118,48 @@ def menu_backup_imaging(): time.sleep(5) pause('\nPress Enter to return to main menu... ') -def menu_windows_setup(): +def menu_root(): + title = '{}: Main Menu'.format(KIT_NAME_FULL) + menus = [ + {'Name': 'Create Backups', 'Menu': menu_backup}, + {'Name': 'Setup Windows', 'Menu': menu_setup}, + {'Name': 'Misc Tools', 'Menu': menu_tools}, + ] + actions = [ + {'Name': 'Command Prompt', 'Letter': 'C'}, + {'Name': 'Reboot', 'Letter': 'R'}, + {'Name': 'Shutdown', 'Letter': 'S'}, + ] + + # Main loop + while True: + selection = menu_select( + title=title, + main_entries=menus, + action_entries=actions, + secret_exit=True) + + if (selection.isnumeric()): + try: + menus[int(selection)-1]['Menu']() + except AbortError: + pass + elif (selection == 'C'): + run_program(['cmd', '-new_console:n'], check=False) + elif (selection == 'R'): + run_program(['wpeutil', 'reboot']) + elif (selection == 'S'): + run_program(['wpeutil', 'shutdown']) + else: + exit_script() + +def menu_setup(): """Format a drive, partition for MBR or GPT, apply a Windows image, and rebuild the boot files""" errors = False # Set ticket ID os.system('cls') - ticket_id = get_ticket_id() + ticket_id = get_ticket_number() # Select the version of Windows to apply windows_version = select_windows_version() @@ -197,80 +260,55 @@ def menu_windows_setup(): pause('\nPress Enter to return to main menu... ') def menu_tools(): - tools = [ - {'Name': 'Blue Screen View', 'Folder': 'BlueScreenView', 'File': 'BlueScreenView.exe'}, - {'Name': 'Fast Copy', 'Folder': 'FastCopy', 'File': 'FastCopy.exe', 'Args': ['/log', '/logfile=X:\WK\Info\FastCopy.log', '/cmd=noexist_only', '/utf8', '/skip_empty_dir', '/linkdest', '/open_window', '/balloon=FALSE', r'/exclude=$RECYCLE.BIN;$Recycle.Bin;.AppleDB;.AppleDesktop;.AppleDouble;.com.apple.timemachine.supported;.dbfseventsd;.DocumentRevisions-V100*;.DS_Store;.fseventsd;.PKInstallSandboxManager;.Spotlight*;.SymAV*;.symSchedScanLockxz;.TemporaryItems;.Trash*;.vol;.VolumeIcon.icns;desktop.ini;Desktop?DB;Desktop?DF;hiberfil.sys;lost+found;Network?Trash?Folder;pagefile.sys;Recycled;RECYCLER;System?Volume?Information;Temporary?Items;Thumbs.db']}, - {'Name': 'HWiNFO', 'Folder': 'HWiNFO', 'File': 'HWiNFO.exe'}, - {'Name': 'NT Password Editor', 'Folder': 'NT Password Editor', 'File': 'ntpwedit.exe'}, - {'Name': 'Notepad++', 'Folder': 'NotepadPlusPlus', 'File': 'notepadplusplus.exe'}, - {'Name': 'PhotoRec', 'Folder': 'TestDisk', 'File': 'photorec_win.exe', 'Args': ['-new_console:n']}, - {'Name': 'Prime95', 'Folder': 'Prime95', 'File': 'prime95.exe'}, - {'Name': 'ProduKey', 'Folder': 'ProduKey', 'File': 'ProduKey.exe', 'Args': ['/external', '/ExtractEdition:1']}, - {'Name': 'Q-Dir', 'Folder': 'Q-Dir', 'File': 'Q-Dir.exe'}, - {'Name': 'TestDisk', 'Folder': 'TestDisk', 'File': 'testdisk_win.exe', 'Args': ['-new_console:n']}, - ] - actions = [ - {'Name': 'Main Menu', 'Letter': 'M'}, - ] + title = '{}: Tools Menu'.format(KIT_NAME_FULL) + tools = [k for k in sorted(PE_TOOLS.keys())] + actions = [{'Name': 'Main Menu', 'Letter': 'M'},] # Menu loop while True: - selection = menu_select('Tools Menu', tools, actions) - + selection = menu_select(title, tools, actions) if (selection.isnumeric()): tool = tools[int(selection)-1] - cmd = ['{bin}\\{folder}\\{file}'.format(bin=bin, folder=tool['Folder'], file=tool['File'])] - if tool['Name'] == 'Blue Screen View': + cmd = [PE_TOOLS[tool]['Path']] + PE_TOOLS[tool].get('Args', []) + if tool == 'Blue Screen View': # Select path to scan minidump_path = select_minidump_path() - if minidump_path is not None: - tool['Args'] = ['/MiniDumpFolder', minidump_path] - if 'Args' in tool: - cmd += tool['Args'] + if minidump_path: + cmd.extend(['/MiniDumpFolder', minidump_path]) try: - subprocess.Popen(cmd) - except: + popen_program(cmd) + except Exception: print_error('Failed to run {prog}'.format(prog=tool['Name'])) - time.sleep(2) + time.sleep(2) + pause() elif (selection == 'M'): break -def menu_main(): - menus = [ - {'Name': 'Create Backups', 'Menu': menu_backup_imaging}, - {'Name': 'Setup Windows', 'Menu': menu_windows_setup}, - {'Name': 'Misc Tools', 'Menu': menu_tools}, - ] - actions = [ - {'Name': 'Command Prompt', 'Letter': 'C'}, - {'Name': 'Reboot', 'Letter': 'R'}, - {'Name': 'Shutdown', 'Letter': 'S'}, - ] +def select_minidump_path(): + dumps = [] - # Main loop - while True: - selection = menu_select('Main Menu', menus, actions, secret_exit=True) + # Assign volume letters first + assign_volume_letters() - if (selection.isnumeric()): - try: - menus[int(selection)-1]['Menu']() - except AbortError: - pass - except: - print_error('Major exception in: {menu}'.format(menu=menus[int(selection)-1]['Name'])) - print_warning(' Please let The Wizard know and he\'ll look into it (Please include the details below).') - print(traceback.print_exc()) - time.sleep(5) - print_info(' You can retry but if this crashes again an alternative approach may be required.') - pause('\nPress enter to return to the main menu') - elif (selection == 'C'): - run_program(['cmd', '-new_console:n'], check=False) - elif (selection == 'R'): - run_program(['wpeutil', 'reboot']) - elif (selection == 'S'): - run_program(['wpeutil', 'shutdown']) - else: - sys.exit(0) + # Search for minidumps + tmp = run_program('mountvol') + tmp = [d for d in re.findall(r'.*([A-Za-z]):\\', tmp.stdout.decode())] + # Remove RAMDisk letter + if 'X' in tmp: + tmp.remove('X') + for drive in tmp: + if os.path.exists('{drive}:\\Windows\\MiniDump'.format(drive=drive)): + dumps.append({'Name': '{drive}:\\Windows\\MiniDump'.format(drive=drive)}) + + # Check results before showing menu + if len(dumps) == 0: + print_error(' No BSoD / MiniDump paths found') + time.sleep(2) + return None + + # Menu + selection = menu_select('Which BSoD / MiniDump path are we scanning?', dumps, []) + return dumps[int(selection) - 1]['Name'] if __name__ == '__main__': - menu_main() \ No newline at end of file + print("This file is not meant to be called directly.") diff --git a/Scripts/settings/main.py b/Scripts/settings/main.py new file mode 100644 index 00000000..4b1d9818 --- /dev/null +++ b/Scripts/settings/main.py @@ -0,0 +1,68 @@ +# Wizard Kit PE: Settings - Main / Branding + +# Features +ENABLED_UPLOAD_DATA = False + +# STATIC VARIABLES (also used by .cmd files) +## Not using spaces aroung '=' for easier .cmd substrings +ARCHIVE_PASSWORD='Abracadabra' +KIT_NAME_FULL='Wizard Kit PE' +KIT_NAME_SHORT='WKPE' +OFFICE_SERVER_IP='10.0.0.10' +QUICKBOOKS_SERVER_IP='10.0.0.10' +SUPPORT_MESSAGE='Please let 2Shirt know by opening an issue on GitHub' +TIME_ZONE='Pacific Standard Time' # Always use "Standard Time" (DST is applied correctly) + +# SERVER VARIABLES +## NOTE: Windows can only use one user per server. This means that if +## one server serves multiple shares then you have to use the same +## user/password for all of those shares. +BACKUP_SERVERS = [ + { 'IP': '10.0.0.10', + 'Name': 'ServerOne', + 'Mounted': False, + 'Share': 'Backups', + 'User': 'restore', + 'Pass': 'Abracadabra', + }, + { 'IP': '10.0.0.11', + 'Name': 'ServerTwo', + 'Mounted': False, + 'Share': 'Backups', + 'User': 'restore', + 'Pass': 'Abracadabra', + }, +] +CLIENT_INFO_SERVER = { + 'IP': '10.0.0.10', + 'RegEntry': r'0x10001,0x00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000', + 'Share': '/srv/ClientInfo', + 'User': 'upload', +} +OFFICE_SERVER = { + 'IP': OFFICE_SERVER_IP, + 'Name': 'ServerOne', + 'Mounted': False, + 'Share': 'Office', + 'User': 'restore', + 'Pass': 'Abracadabra', +} +QUICKBOOKS_SERVER = { + 'IP': QUICKBOOKS_SERVER_IP, + 'Name': 'ServerOne', + 'Mounted': False, + 'Share': 'QuickBooks', + 'User': 'restore', + 'Pass': 'Abracadabra', +} +WINDOWS_SERVER = { + 'IP': '10.0.0.10', + 'Name': 'ServerOne', + 'Mounted': False, + 'Share': 'Windows', + 'User': 'restore', + 'Pass': 'Abracadabra', +} + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/Scripts/settings/tools.py b/Scripts/settings/tools.py new file mode 100644 index 00000000..ab8be9e1 --- /dev/null +++ b/Scripts/settings/tools.py @@ -0,0 +1,55 @@ +# Wizard Kit PE: Settings - Tools + +TOOLS = { + # NOTE: BinDir will be prepended to these paths at runtime + 'AIDA64': { + '32': r'AIDA64\aida64.exe'}, + 'AutoRuns': { + '32': r'Autoruns\autoruns.exe', + '64': r'Autoruns\autoruns64.exe'}, + 'BleachBit': { + '32': r'BleachBit\bleachbit_console.exe'}, + 'Caffeine': { + '32': r'Caffeine\caffeine.exe'}, + 'Du': { + '32': r'Du\du.exe', + '64': r'Du\du64.exe'}, + 'ERUNT': { + '32': r'ERUNT\ERUNT.EXE'}, + 'Everything': { + '32': r'Everything\Everything.exe', + '64': r'Everything\Everything64.exe'}, + 'FastCopy': { + '32': r'FastCopy\FastCopy.exe', + '64': r'FastCopy\FastCopy64.exe'}, + 'HitmanPro': { + '32': r'HitmanPro\HitmanPro.exe', + '64': r'HitmanPro\HitmanPro64.exe'}, + 'HWiNFO': { + '32': r'HWiNFO\HWiNFO.exe', + '64': r'HWiNFO\HWiNFO64.exe'}, + 'KVRT': { + '32': r'KVRT\KVRT.exe'}, + 'NotepadPlusPlus': { + '32': r'NotepadPlusPlus\notepadplusplus.exe'}, + 'ProduKey': { + '32': r'ProduKey\ProduKey.exe', + '64': r'ProduKey\ProduKey64.exe'}, + 'PuTTY-PSFTP': { + '32': r'PuTTY\PSFTP.EXE'}, + 'RKill': { + '32': r'RKill\RKill.exe'}, + 'SevenZip': { + '32': r'7-Zip\7za.exe', + '64': r'7-Zip\7za64.exe'}, + 'TDSSKiller': { + '32': r'TDSSKiller\TDSSKiller.exe'}, + 'wimlib-imagex': { + '32': r'wimlib\x32\wimlib-imagex.exe', + '64': r'wimlib\x64\wimlib-imagex.exe'}, + 'XMPlay': { + '32': r'XMPlay\xmplay.exe'}, + } + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py new file mode 100644 index 00000000..370186f0 --- /dev/null +++ b/Scripts/winpe_root_menu.py @@ -0,0 +1,21 @@ +# Wizard Kit PE: Root Menu + +import os +import sys + +# Init +os.chdir(os.path.dirname(os.path.realpath(__file__))) +sys.path.append(os.getcwd()) +from functions.common import * +from functions.winpe_menus import * +init_global_vars() +os.system('title {}: Root Menu'.format(KIT_NAME_FULL)) +global_vars['LogFile'] = r'{LogDir}\WinPE.log'.format(**global_vars) + +if __name__ == '__main__': + try: + menu_root() + except SystemExit: + pass + except: + major_exception() From dd20cdd36e04433fbdf2275423cd83c373e65921 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 13:40:22 -0800 Subject: [PATCH 02/23] Renamed X:\WK to X:\.bin * Allows functions\common.py to be used as-is --- .bin/Scripts/build_pe.ps1 | 14 +++++++------- System32/Winpeshl.ini | 4 ++-- System32/menu.cmd | 2 +- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.bin/Scripts/build_pe.ps1 b/.bin/Scripts/build_pe.ps1 index 631d6d9e..afcc6fe7 100644 --- a/.bin/Scripts/build_pe.ps1 +++ b/.bin/Scripts/build_pe.ps1 @@ -501,16 +501,16 @@ if ($MyInvocation.InvocationName -ne ".") { # Add WK tools Write-Host "Copying tools..." - Copy-Item -Path "$Root\WK\$Arch" -Destination "$Mount\WK" -Recurse -Force - Copy-Item -Path "$Root\WK\_include\*" -Destination "$Mount\WK" -Recurse -Force + Copy-Item -Path "$Root\WK\$Arch" -Destination "$Mount\.bin" -Recurse -Force + Copy-Item -Path "$Root\WK\_include\*" -Destination "$Mount\.bin" -Recurse -Force if ($Arch -eq "amd64") { - $DestIni = "$Mount\WK\HWiNFO\HWiNFO64.INI" + $DestIni = "$Mount\.bin\HWiNFO\HWiNFO64.INI" } else { - $DestIni = "$Mount\WK\HWiNFO\HWiNFO32.INI" + $DestIni = "$Mount\.bin\HWiNFO\HWiNFO32.INI" } - Move-Item -Path "$Mount\WK\HWiNFO\HWiNFO.INI" -Destination $DestIni -Force - Copy-Item -Path "$Root\WinPE.jpg" -Destination "$Mount\WK\ConEmu\ConEmu.jpg" -Recurse -Force - Copy-Item -Path "$Root\Scripts" -Destination "$Mount\WK\Scripts" -Recurse -Force + Move-Item -Path "$Mount\.bin\HWiNFO\HWiNFO.INI" -Destination $DestIni -Force + Copy-Item -Path "$Root\WinPE.jpg" -Destination "$Mount\.bin\ConEmu\ConEmu.jpg" -Recurse -Force + Copy-Item -Path "$Root\Scripts" -Destination "$Mount\.bin\Scripts" -Recurse -Force # Add System32 items $HostSystem32 = "{0}\System32" -f $Env:SystemRoot diff --git a/System32/Winpeshl.ini b/System32/Winpeshl.ini index 716d7d6e..d272162f 100644 --- a/System32/Winpeshl.ini +++ b/System32/Winpeshl.ini @@ -2,5 +2,5 @@ [LaunchApps] wpeinit wpeutil updatebootinfo -cd /d "%SystemDrive%\WK" -"%SystemDrive%\WK\ConEmu\ConEmu.exe", /cmd cmd /k cd "%SystemDrive%\WK" & python "%SystemDrive%\WK\Scripts\menu.py" +cd /d "%SystemDrive%\.bin" +"%SystemDrive%\.bin\ConEmu\ConEmu.exe", /cmd cmd /k cd "%SystemDrive%\.bin" & python "%SystemDrive%\.bin\Scripts\menu.py" diff --git a/System32/menu.cmd b/System32/menu.cmd index 771dc372..b8e5509a 100644 --- a/System32/menu.cmd +++ b/System32/menu.cmd @@ -1,2 +1,2 @@ @echo off -python "%SystemDrive%\WK\Scripts\menu.py" \ No newline at end of file +python "%SystemDrive%\.bin\Scripts\menu.py" \ No newline at end of file From e136283a71e43820f572b2eaa366cb6160a5c3b8 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 14:40:44 -0800 Subject: [PATCH 03/23] Added clear_screen() --- Scripts/functions/common.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py index a0453a95..d20a9ea6 100644 --- a/Scripts/functions/common.py +++ b/Scripts/functions/common.py @@ -81,6 +81,10 @@ def ask(prompt='Kotaero!'): print_log(message=message) return answer +def clear_screen(): + """Simple wrapper for cls.""" + os.system('cls') + def convert_to_bytes(size): """Convert human-readable size str to bytes and return an int.""" size = str(size) From b5b03e4dfef7e5d1e69d3e7a4f4b5a15c540ea34 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 15:19:43 -0800 Subject: [PATCH 04/23] Fixed notepad replacement --- .bin/Scripts/build_pe.ps1 | 13 ++++++------ WK/_include/NotepadPlusPlus/npp.cmd | 3 +++ WK/_include/NotepadPlusPlus/npp.vbs | 33 ----------------------------- 3 files changed, 9 insertions(+), 40 deletions(-) create mode 100644 WK/_include/NotepadPlusPlus/npp.cmd delete mode 100644 WK/_include/NotepadPlusPlus/npp.vbs diff --git a/.bin/Scripts/build_pe.ps1 b/.bin/Scripts/build_pe.ps1 index afcc6fe7..4276b152 100644 --- a/.bin/Scripts/build_pe.ps1 +++ b/.bin/Scripts/build_pe.ps1 @@ -499,7 +499,7 @@ if ($MyInvocation.InvocationName -ne ".") { ) Start-Process -FilePath $DISM -ArgumentList $ArgumentList -NoNewWindow -Wait - # Add WK tools + # Add tools Write-Host "Copying tools..." Copy-Item -Path "$Root\WK\$Arch" -Destination "$Mount\.bin" -Recurse -Force Copy-Item -Path "$Root\WK\_include\*" -Destination "$Mount\.bin" -Recurse -Force @@ -538,17 +538,16 @@ if ($MyInvocation.InvocationName -ne ".") { $RegKey = $Hive.OpenSubKey($RegPath) $CurValue = $RegKey.GetValue( "Path", $false, [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames) - $NewValue = "$CurValue;%SystemDrive%\WK\7-Zip;%SystemDrive%\WK\python;%SystemDrive%\WK\wimlib" + $NewValue = "$CurValue;%SystemDrive%\.bin\7-Zip;%SystemDrive%\.bin\python;%SystemDrive%\.bin\wimlib" Set-ItemProperty -Path "HKLM:\$RegPath" -Name "Path" -Value $NewValue -Force | Out-Null $Hive.close() $RegKey.close() # Replace Notepad - ## Currently broken ## - # $RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe" - # $NewValue = 'wscript "X:\WK\NotepadPlusPlus\npp.vbs"' - # New-Item -Path $RegPath -Force | Out-Null - # New-ItemProperty -Path $RegPath -Name "Debugger" -Value $NewValue -Force | Out-Null + $RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe" + $NewValue = 'cmd /c "%SystemDrive%\.bin\NotepadPlusPlus\npp.cmd"' + New-Item -Path $RegPath -Force | Out-Null + New-ItemProperty -Path $RegPath -Name "Debugger" -Value $NewValue -Force | Out-Null # Run garbage collection to release potential stale handles ## Credit: https://jrich523.wordpress.com/2012/03/06/powershell-loading-and-unloading-registry-hives/ diff --git a/WK/_include/NotepadPlusPlus/npp.cmd b/WK/_include/NotepadPlusPlus/npp.cmd new file mode 100644 index 00000000..d387df55 --- /dev/null +++ b/WK/_include/NotepadPlusPlus/npp.cmd @@ -0,0 +1,3 @@ +@echo off + +start "" %SystemDrive%\.bin\NotepadPlusPlus\notepadplusplus.exe %2 %3 %4 %5 %6 %7 %8 %9 \ No newline at end of file diff --git a/WK/_include/NotepadPlusPlus/npp.vbs b/WK/_include/NotepadPlusPlus/npp.vbs deleted file mode 100644 index 3fb12a41..00000000 --- a/WK/_include/NotepadPlusPlus/npp.vbs +++ /dev/null @@ -1,33 +0,0 @@ -'// DISCLAIMER -'// THIS COMES WITH NO WARRANTY, IMPLIED OR OTHERWISE. USE AT YOUR OWN RISK -'// IF YOU ARE NOT COMFORTABLE EDITING THE REGISTRY THEN DO NOT USE THIS SCRIPT -'// -'// NOTES: -'// This affects all users. -'// This will prevent ANY executable named notepad.exe from running located anywhere on this computer!! -'// -'// Save this text to your notepad++ folder as a text file named npp.vbs (some AV don't like vbs, get a different AV :-P ) -'// -'// USAGE -'// 1) -'// Navigate to registry key HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\ -'// -' // 2) -'// Add new subkey called notepad.exe -'// This step is what tells windows to use the notepad++ exe, to undo simply delete this key -'// -'// 3) -'// Create new Sting Value called Debugger -'// -'// 4) -'// Modify value and enter wscript.exe "path to npp.vbs" e.g. wscript.exe "C:\Program Files\Notepad++\npp.vbs" - -Option Explicit -Dim sCmd, x -sCmd = """" & LeftB(WScript.ScriptFullName, LenB(WScript.ScriptFullName) - LenB(WScript.ScriptName)) & "notepad++.exe" & """ """ -For x = 1 To WScript.Arguments.Count - 1 - sCmd = sCmd & WScript.Arguments(x) & " " -Next -sCmd = sCmd & """" -CreateObject("WScript.Shell").Run sCmd, 1, True -WScript.Quit \ No newline at end of file From 9a3234c82221524b0dd68772e6f8a91a60e6e4fd Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 16:07:51 -0800 Subject: [PATCH 05/23] Simplified Windows share sections --- Scripts/functions/common.py | 2 +- Scripts/functions/windows_setup.py | 34 ++++++++---------------------- 2 files changed, 10 insertions(+), 26 deletions(-) diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py index d20a9ea6..fcb2dd51 100644 --- a/Scripts/functions/common.py +++ b/Scripts/functions/common.py @@ -174,7 +174,7 @@ def human_readable_size(size, decimals=0): size = convert_to_bytes(size) # Verify we have a valid size - if size <= 0: + if size < 0: return '{size:>{width}} b'.format(size='???', width=width) # Convert to sensible units diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py index e70e88cf..8de76f76 100644 --- a/Scripts/functions/windows_setup.py +++ b/Scripts/functions/windows_setup.py @@ -61,13 +61,9 @@ def is_index_in_image(bin=None, filename=None, imagename=None): return True -def find_windows_image(bin, windows_version=None): +def find_windows_image(bin, windows_version): """Search for a Windows source image file on local drives and network drives (in that order)""" image = {} - - # Bail early - if windows_version is None: - raise Exception('Windows version not specified.') imagefile = windows_version['Image File'] # Search local source @@ -85,9 +81,10 @@ def find_windows_image(bin, windows_version=None): break # Check for network source (if necessary) - if not any(image): + if not image: + mount_windows_share() if not WINDOWS_SERVER['Mounted']: - mount_windows_share() + return None for ext in ['esd', 'wim', 'swm']: filename = '\\\\{IP}\\{Share}\\images\\{imagefile}'.format(imagefile=imagefile, **WINDOWS_SERVER) filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) @@ -202,25 +199,12 @@ def get_boot_mode(): return boot_mode def mount_windows_share(): - """Mount the Windows images share for use in Windows setup""" - - # Blindly skip if we mounted earlier + """Mount the Windows images share unless labeled as already mounted.""" if WINDOWS_SERVER['Mounted']: - return None - else: - try: - # Test connection - run_program('ping -w 800 -n 2 {IP}'.format(**WINDOWS_SERVER)) - # Mount - run_program('net use \\\\{IP}\\{Share} /user:{User} {Pass}'.format(**WINDOWS_SERVER)) - print_info('Mounted {Name}'.format(**WINDOWS_SERVER)) - WINDOWS_SERVER['Mounted'] = True - except subprocess.CalledProcessError: - print_error('Failed to mount \\\\{Name}\\{Share}, {IP} unreachable.'.format(**WINDOWS_SERVER)) - time.sleep(1) - except: - print_warning('Failed to mount \\\\{Name}\\{Share}'.format(**WINDOWS_SERVER)) - time.sleep(1) + # Blindly skip if we mounted earlier + continue + + mount_network_share(WINDOWS_SERVER) def select_windows_version(): actions = [{'Name': 'Main Menu', 'Letter': 'M'},] From 80cb9b8ceab845003b3145fae347339383b7d8b6 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 16:40:40 -0800 Subject: [PATCH 06/23] Function separation done. * Split backup.py into disk & backup * disk.py is for lower level disk management * Renamed functions\data.py's select_disk() to select_volume() * Avoid name collision with functions\disk.py's select_disk() * --- Scripts/functions/backup.py | 327 --------------------------- Scripts/functions/data.py | 4 +- Scripts/functions/disk.py | 345 +++++++++++++++++++++++++++++ Scripts/functions/windows_setup.py | 56 ++--- Scripts/functions/winpe_menus.py | 4 +- Scripts/winpe_root_menu.py | 1 - 6 files changed, 372 insertions(+), 365 deletions(-) create mode 100644 Scripts/functions/disk.py diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py index 94f3b103..e33004ba 100644 --- a/Scripts/functions/backup.py +++ b/Scripts/functions/backup.py @@ -1,18 +1,6 @@ # Wizard Kit PE: Functions - Backup from functions.common import * -import partition_uids - -def assign_volume_letters(): - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - for vol in get_volumes(): - script.write('select volume {Number}\n'.format(**vol)) - script.write('assign\n') - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - except subprocess.CalledProcessError: - pass def backup_partition(bin=None, disk=None, par=None): # Bail early @@ -40,232 +28,6 @@ def backup_partition(bin=None, disk=None, par=None): par['Error'] = err.stderr.decode().splitlines() raise BackupError -def get_attached_disk_info(): - """Get details about the attached disks""" - disks = [] - print_info('Getting drive info...') - - # Assign all the letters - assign_volume_letters() - - # 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_disk_details(disk=None): - details = {} - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('detail disk\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Remove empty lines - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Set disk name - details['Name'] = tmp[4] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair - tmp = [s.split(':') for s in tmp if ':' in s] - - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) - - return details - -def get_disks(): - disks = [] - - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('list disk\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Append disk numbers - for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return): - _num = tmp[0] - _size = human_readable_size(tmp[1]) - disks.append({'Number': _num, 'Size': _size}) - - return disks - -def get_partition_details(disk=None, par=None): - details = {} - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - if par is None: - raise Exception('Partition not specified.') - - # Diskpart details - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('select partition {Number}\n'.format(**par)) - script.write('detail partition\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Get volume letter or RAW status - tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return) - if tmp: - if tmp.group(1).upper() == 'RAW': - details['FileSystem'] = RAW - else: - details['Letter'] = tmp.group(1) - - # Remove empty lines from process_return - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair - tmp = [s.split(':') for s in tmp if ':' in s] - - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) - - # Get MBR type / GPT GUID for extra details on "Unknown" partitions - guid = partition_uids.lookup_guid(details['Type']) - if guid is not None: - details.update({ - 'Description': guid.get('Description', ''), - 'OS': guid.get('OS', '')}) - - if 'Letter' in details: - # Disk usage - tmp = shutil.disk_usage('{Letter}:\\'.format(**details)) - details['Used Space'] = human_readable_size(tmp.used) - - # fsutil details - try: - process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Remove empty lines from process_return - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Add "Feature" lines - details['File System Features'] = [s.strip() for s in tmp if ':' not in s] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair - tmp = [s.split(':') for s in tmp if ':' in s] - - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) - - # Set Volume Name - details['Name'] = details.get('Volume Name', '') - - # Set FileSystem Type - if details.get('FileSystem', '') != 'RAW': - details['FileSystem'] = details.get('File System Name', 'Unknown') - - return details - -def get_partitions(disk=None): - partitions = [] - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('list partition\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Append partition numbers - for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE): - _num = tmp[0] - _size = human_readable_size(tmp[1]) - partitions.append({'Number': _num, 'Size': _size}) - - return partitions - -def get_table_type(disk=None): - _type = 'Unknown' - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - - try: - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('uniqueid disk\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - if re.findall(r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', process_return, re.IGNORECASE): - _type = 'GPT' - elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE): - _type = 'RAW' - elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE): - _type = 'MBR' - - return _type - -def get_volumes(): - vols = [] - - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('list volume\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() - except subprocess.CalledProcessError: - pass - else: - # Append volume numbers - for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return): - vols.append({'Number': tmp[0], 'Letter': tmp[1]}) - - return vols - def prep_disk_for_backup(dest=None, disk=None, ticket_id=None): disk['Backup Warnings'] = '\n' disk['Clobber Risk'] = [] @@ -342,54 +104,6 @@ def prep_disk_for_backup(dest=None, disk=None, ticket_id=None): 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 prep_disk_for_formatting(disk=None): - disk['Format Warnings'] = '\n' - width = len(str(len(disk['Partitions']))) - - # Bail early - if disk is None: - raise Exception('Disk not provided.') - - # Set boot method and partition table type - disk['Use GPT'] = True - if (get_boot_mode() == 'UEFI'): - if (not ask("Setup Windows to use UEFI booting?")): - disk['Use GPT'] = False - else: - if (ask("Setup Windows to use BIOS/Legacy booting?")): - disk['Use GPT'] = False - - # Set Display and Warning Strings - if len(disk['Partitions']) == 0: - disk['Format Warnings'] += 'No partitions found\n' - for par in disk['Partitions']: - if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): - # FileSystem not accessible to WinPE. List partition type / OS info for technician - par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format( - width=width, - q='"' if par['Name'] != '' else '', - **par) - else: - # FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician - par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format( - width=width, - q='"' if par['Name'] != '' else '', - **par) - -def remove_volume_letters(keep=''): - if keep is None: - keep = '' - try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - for vol in get_volumes(): - if vol['Letter'].upper() != keep.upper(): - script.write('select volume {Number}\n'.format(**vol)) - script.write('remove noerr\n') - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - except subprocess.CalledProcessError: - pass - def select_backup_destination(): # Build menu dests = [] @@ -420,47 +134,6 @@ def select_backup_destination(): print_warning('No backup destinations found.') return None -def select_disk(prompt='Which disk?'): - """Select a disk from the attached disks""" - disks = get_attached_disk_info() - - # Build menu - disk_options = [] - for disk in disks: - display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) - if len(disk['Partitions']) > 0: - pwidth=len(str(len(disk['Partitions']))) - for par in disk['Partitions']: - # Show unsupported partition(s) in RED - par_skip = False - if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): - par_skip = True - if par_skip: - display_name += COLORS['YELLOW'] - - # Main text - display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par) - if par['Name'] != '': - display_name += '\t"{Name}"'.format(**par) - - # Clear color (if set above) - if par_skip: - display_name += COLORS['CLEAR'] - else: - display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS) - disk_options.append({'Name': display_name, 'Disk': disk}) - actions = [ - {'Name': 'Main Menu', 'Letter': 'M'}, - ] - - # Menu loop - selection = menu_select(prompt, disk_options, actions) - - if (selection.isnumeric()): - return disk_options[int(selection)-1]['Disk'] - elif (selection == 'M'): - abort_to_main_menu() - def verify_wim_backup(bin=None, par=None): # Bail early if bin is None: diff --git a/Scripts/functions/data.py b/Scripts/functions/data.py index 66797329..63825293 100644 --- a/Scripts/functions/data.py +++ b/Scripts/functions/data.py @@ -374,7 +374,7 @@ def scan_source_wim(source_wim, dest_path, rel_path=None, interactive=True): def select_destination(folder_path, prompt='Select destination'): """Select destination drive, returns path as string.""" - disk = select_disk(prompt) + disk = select_volume(prompt) if 'fixed' not in disk['Disk'].opts: folder_path = folder_path.replace('\\', '-') path = '{disk}{folder_path}_{Date}'.format( @@ -388,7 +388,7 @@ def select_destination(folder_path, prompt='Select destination'): return path -def select_disk(title='Select disk', auto_select=True): +def select_volume(title='Select disk', auto_select=True): """Select disk from attached disks. returns dict.""" actions = [{'Name': 'Quit', 'Letter': 'Q'}] disks = [] diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py new file mode 100644 index 00000000..3b7730c0 --- /dev/null +++ b/Scripts/functions/disk.py @@ -0,0 +1,345 @@ +# Wizard Kit PE: Functions - Disk + +from functions.common import * +import partition_uids + +def assign_volume_letters(): + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + for vol in get_volumes(): + script.write('select volume {Number}\n'.format(**vol)) + script.write('assign\n') + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + except subprocess.CalledProcessError: + pass + +def get_attached_disk_info(): + """Get details about the attached disks""" + disks = [] + print_info('Getting drive info...') + + # Assign all the letters + assign_volume_letters() + + # 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: + reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 'System\\CurrentControlSet\\Control') + reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] + if reg_value == 2: + boot_mode = 'UEFI' + except: + boot_mode = 'Unknown' + + return boot_mode + +def get_disk_details(disk=None): + details = {} + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('detail disk\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Remove empty lines + tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] + + # Set disk name + details['Name'] = tmp[4] + + # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + tmp = [s.split(':') for s in tmp if ':' in s] + + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + return details + +def get_disks(): + disks = [] + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('list disk\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Append disk numbers + for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return): + _num = tmp[0] + _size = human_readable_size(tmp[1]) + disks.append({'Number': _num, 'Size': _size}) + + return disks + +def get_partition_details(disk=None, par=None): + details = {} + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + if par is None: + raise Exception('Partition not specified.') + + # Diskpart details + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('select partition {Number}\n'.format(**par)) + script.write('detail partition\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Get volume letter or RAW status + tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return) + if tmp: + if tmp.group(1).upper() == 'RAW': + details['FileSystem'] = RAW + else: + details['Letter'] = tmp.group(1) + + # Remove empty lines from process_return + tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] + + # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + tmp = [s.split(':') for s in tmp if ':' in s] + + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + # Get MBR type / GPT GUID for extra details on "Unknown" partitions + guid = partition_uids.lookup_guid(details['Type']) + if guid is not None: + details.update({ + 'Description': guid.get('Description', ''), + 'OS': guid.get('OS', '')}) + + if 'Letter' in details: + # Disk usage + tmp = shutil.disk_usage('{Letter}:\\'.format(**details)) + details['Used Space'] = human_readable_size(tmp.used) + + # fsutil details + try: + process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Remove empty lines from process_return + tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] + + # Add "Feature" lines + details['File System Features'] = [s.strip() for s in tmp if ':' not in s] + + # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + tmp = [s.split(':') for s in tmp if ':' in s] + + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + # Set Volume Name + details['Name'] = details.get('Volume Name', '') + + # Set FileSystem Type + if details.get('FileSystem', '') != 'RAW': + details['FileSystem'] = details.get('File System Name', 'Unknown') + + return details + +def get_partitions(disk=None): + partitions = [] + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('list partition\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Append partition numbers + for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE): + _num = tmp[0] + _size = human_readable_size(tmp[1]) + partitions.append({'Number': _num, 'Size': _size}) + + return partitions + +def get_table_type(disk=None): + _type = 'Unknown' + + # Bail early + if disk is None: + raise Exception('Disk not specified.') + + try: + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {Number}\n'.format(**disk)) + script.write('uniqueid disk\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + if re.findall(r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', process_return, re.IGNORECASE): + _type = 'GPT' + elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE): + _type = 'RAW' + elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE): + _type = 'MBR' + + return _type + +def get_volumes(): + vols = [] + + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('list volume\n') + process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + process_return = process_return.stdout.decode().strip() + except subprocess.CalledProcessError: + pass + else: + # Append volume numbers + for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return): + vols.append({'Number': tmp[0], 'Letter': tmp[1]}) + + return vols + +def prep_disk_for_formatting(disk=None): + disk['Format Warnings'] = '\n' + width = len(str(len(disk['Partitions']))) + + # Bail early + if disk is None: + raise Exception('Disk not provided.') + + # Set boot method and partition table type + disk['Use GPT'] = True + if (get_boot_mode() == 'UEFI'): + if (not ask("Setup Windows to use UEFI booting?")): + disk['Use GPT'] = False + else: + if (ask("Setup Windows to use BIOS/Legacy booting?")): + disk['Use GPT'] = False + + # Set Display and Warning Strings + if len(disk['Partitions']) == 0: + disk['Format Warnings'] += 'No partitions found\n' + for par in disk['Partitions']: + if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): + # FileSystem not accessible to WinPE. List partition type / OS info for technician + par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format( + width=width, + q='"' if par['Name'] != '' else '', + **par) + else: + # FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician + par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format( + width=width, + q='"' if par['Name'] != '' else '', + **par) + +def remove_volume_letters(keep=''): + if keep is None: + keep = '' + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + for vol in get_volumes(): + if vol['Letter'].upper() != keep.upper(): + script.write('select volume {Number}\n'.format(**vol)) + script.write('remove noerr\n') + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + except subprocess.CalledProcessError: + pass + +def select_disk(prompt='Which disk?'): + """Select a disk from the attached disks""" + disks = get_attached_disk_info() + + # Build menu + disk_options = [] + for disk in disks: + display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) + if len(disk['Partitions']) > 0: + pwidth=len(str(len(disk['Partitions']))) + for par in disk['Partitions']: + # Show unsupported partition(s) in RED + par_skip = False + if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): + par_skip = True + if par_skip: + display_name += COLORS['YELLOW'] + + # Main text + display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par) + if par['Name'] != '': + display_name += '\t"{Name}"'.format(**par) + + # Clear color (if set above) + if par_skip: + display_name += COLORS['CLEAR'] + else: + display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS) + disk_options.append({'Name': display_name, 'Disk': disk}) + actions = [ + {'Name': 'Main Menu', 'Letter': 'M'}, + ] + + # Menu loop + selection = menu_select(prompt, disk_options, actions) + + if (selection.isnumeric()): + return disk_options[int(selection)-1]['Disk'] + elif (selection == 'M'): + abort_to_main_menu() + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py index 8de76f76..199999a4 100644 --- a/Scripts/functions/windows_setup.py +++ b/Scripts/functions/windows_setup.py @@ -1,9 +1,9 @@ # Wizard Kit PE: Functions - Windows Setup -from functions.common import * +from functions.data import * # STATIC VARIABLES -DISKPART_SCRIPT = '{tmp}\\diskpart.script'.format(tmp=os.environ['TMP']) +DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP']) WINDOWS_VERSIONS = [ {'Name': 'Windows 7 Home Basic', 'Image File': 'Win7', @@ -43,24 +43,6 @@ WINDOWS_VERSIONS = [ 'Family': '10'}, ] -def is_index_in_image(bin=None, filename=None, imagename=None): - # Bail early - if bin is None: - raise Exception('bin not specified.') - if filename is None: - raise Exception('Filename not specified.') - if imagename is None: - raise Exception('Image Name not specified.') - - cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename) - try: - run_program(cmd) - except subprocess.CalledProcessError: - print_error('Invalid image: {filename}'.format(filename=filename)) - return False - - return True - def find_windows_image(bin, windows_version): """Search for a Windows source image file on local drives and network drives (in that order)""" image = {} @@ -73,7 +55,7 @@ def find_windows_image(bin, windows_version): filename = '{drive}:\\images\\{imagefile}'.format(drive=tmp[0], imagefile=imagefile) filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) if os.path.isfile(filename_ext): - if is_index_in_image(bin, filename_ext, windows_version['Image Name']): + if wim_contains_image(bin, filename_ext, windows_version['Image Name']): image['Ext'] = ext image['File'] = filename image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' @@ -89,7 +71,7 @@ def find_windows_image(bin, windows_version): filename = '\\\\{IP}\\{Share}\\images\\{imagefile}'.format(imagefile=imagefile, **WINDOWS_SERVER) filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) if os.path.isfile(filename_ext): - if is_index_in_image(bin, filename_ext, windows_version['Image Name']): + if wim_contains_image(bin, filename_ext, windows_version['Image Name']): image['Ext'] = ext image['File'] = filename image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' @@ -186,18 +168,6 @@ def format_mbr(disk=None, windows_family=None): run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) time.sleep(2) -def get_boot_mode(): - boot_mode = 'Legacy' - try: - reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 'System\\CurrentControlSet\\Control') - reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] - if reg_value == 2: - boot_mode = 'UEFI' - except: - boot_mode = 'Unknown' - - return boot_mode - def mount_windows_share(): """Mount the Windows images share unless labeled as already mounted.""" if WINDOWS_SERVER['Mounted']: @@ -253,5 +223,23 @@ def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T') def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'): run_program('bcdboot {win}:\\Windows /s {sys}: /f {mode}'.format(win=windows_letter, sys=system_letter, mode=mode)) +def wim_contains_image(bin=None, filename=None, imagename=None): + # Bail early + if bin is None: + raise Exception('bin not specified.') + if filename is None: + raise Exception('Filename not specified.') + if imagename is None: + raise Exception('Image Name not specified.') + + cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename) + try: + run_program(cmd) + except subprocess.CalledProcessError: + print_error('Invalid image: {filename}'.format(filename=filename)) + return False + + return True + if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index 21d2e11b..b03802b4 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -1,6 +1,8 @@ # Wizard Kit PE: Menus -from functions.data import * +from functions.backup import * +from functions.disk import * +from functions.windows_setup import * # STATIC VARIABLES FAST_COPY_ARGS = [ diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py index 370186f0..542114a5 100644 --- a/Scripts/winpe_root_menu.py +++ b/Scripts/winpe_root_menu.py @@ -6,7 +6,6 @@ import sys # Init os.chdir(os.path.dirname(os.path.realpath(__file__))) sys.path.append(os.getcwd()) -from functions.common import * from functions.winpe_menus import * init_global_vars() os.system('title {}: Root Menu'.format(KIT_NAME_FULL)) 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 07/23] 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: From 67f08c50429d88198503dff6b1ed7837bceca4d3 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 19:45:59 -0800 Subject: [PATCH 08/23] updated prep_disk_for_backup() * Partition['Image Path'] is set to the full destination path * (i.e. ['Image Path'] + '\\' + ['Image File'] * Partition['Image File'] variable has been removed * Simplified ['Backup Warnings'] section * Added fix_path() * Replaces unsupported characters/strings with underscores --- Scripts/functions/backup.py | 127 +++++++++++++++++-------------- Scripts/functions/winpe_menus.py | 55 +++++++------ Scripts/winpe_root_menu.py | 2 + 3 files changed, 103 insertions(+), 81 deletions(-) diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py index f91c692f..62e45d19 100644 --- a/Scripts/functions/backup.py +++ b/Scripts/functions/backup.py @@ -2,6 +2,15 @@ from functions.common import * +# Regex +REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE) +REGEX_BAD_PATH_NAMES = re.compile( + r'([<>:"/\\\|\?\*]' + r'|^(CON|PRN|AUX|NUL|COM\d*|LPT\d*)$)' + r'|^\s+' + r'|[\s\.]+$', + re.IGNORECASE) + def backup_partition(disk, partition): if par['Image Exists'] or par['Number'] in disk['Bad Partitions']: raise GenericAbort @@ -10,89 +19,91 @@ def backup_partition(disk, partition): global_vars['Tools']['wimlib-imagex'], 'capture' '{}:\\'.format(par['Letter']), - r'{}\{}'.format(par['Image Path'], par['Image File']), + par['Image Path'], par['Image Name'], # Image name par['Image Name'], # Image description ' --compress=none', ] - os.makedirs(par['Image Path'], exist_ok=True) + dest_dir = re.sub(r'(.*)\\.*$', r'\1', par['Image Path'], re.IGNORECASE) + os.makedirs(dest_dir, exist_ok=True) run_program(cmd) -def prep_disk_for_backup(dest=None, disk=None, ticket_id=None): - disk['Backup Warnings'] = '\n' +def fix_path(path): + return REGEX_BAD_PATH_NAMES.sub('_', path) + +def prep_disk_for_backup(destination, disk, ticket_number): 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) + disk['Bad Partitions'] = [par['Number'] for par in disk['Partitions'] + if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem'])] + num_valid_partitions = len(disk['Partitions']) - len(disk['Bad Partitions']) + disk['Valid Partitions'] = num_valid_partitions if disk['Valid Partitions'] <= 0: - abort_to_main_menu(' No partitions can be imaged for the selected drive') + print_error('ERROR: No partitions can be backed up for this disk') + raise GenericAbort # Prep partitions for par in disk['Partitions']: + display = 'Partition {num:>{width}}:\t{size} {fs}'.format( + num = par['Number'], + width = width, + size = par['Size'], + fs = par['FileSystem']) + 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) + # Set display string using partition description & OS type + display = ' * {display}\t\t{q}{name}{q}\t{desc} ({os})'.format( + display = display, + q = '"' if par['Name'] != '' else '', + name = par['Name'], + desc = par['Description'], + os = par['OS']) + display = '{YELLOW}{display}{CLEAR}'.format( + display=display, **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) + 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) 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'])) + par['Image Path'] = r'{}:\{}'.format( + ticket_number, destination['Letter']) + par['Image Path'] += r'\{}_{}.wim'.format( + par['Number'], par['Image Name']) + par['Image Path'] = fix_path(par['Image Path']) # Check for existing backups - par['Image Exists'] = False - if os.path.exists('{Image Path}\\{Image File}'.format(**par)): - par['Image Exists'] = True + par['Image Exists'] = os.path.exists(par['Image Path']) + if par['Image Exists']: disk['Clobber Risk'].append(par['Number']) - par['Display String'] = '{BLUE} + '.format(**COLORS) + display = '{} + {}'.format(COLORS['BLUE'], display) else: - par['Display String'] = '{CLEAR} '.format(**COLORS) + display = '{} {}'.format(COLORS['CLEAR'], display) # 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, + display += ' (Used: {used})\t{q}{name}{q}{CLEAR}'.format( + used = par['Used Space'], + q = '"' if par['Name'] != '' else '', + name = par['Name'], **COLORS) + # For all partitions + par['Display String'] = display # 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) + warnings = '\n' + if disk['Bad Partitions']: + warnings += '{} * Unsupported filesystem{}\n'.format( + COLORS['YELLOW'], COLORS['CLEAR']) + if disk['Clobber Risk']: + warnings += '{} + Backup exists on {}{}\n'.format( + COLORS['BLUE'], destination['Name'], COLORS['CLEAR']) + if disk['Bad Partitions'] or disk['Clobber Risk']: + warnings += '\n{}Marked partition(s) will NOT be backed up.{}\n'.format( + COLORS['YELLOW'], COLORS['CLEAR']) + disk['Backup Warnings'] = warnings def select_backup_destination(): # Build menu @@ -133,8 +144,8 @@ def verify_wim_backup(bin=None, par=None): # 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)): + cmd = '{bin}\\wimlib\\wimlib-imagex verify "{Image Path}" --nocheck'.format(bin=bin, **par) + if not os.path.exists('{Image Path}'.format(**par)): print_error('Missing.') else: try: diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index 2056b4fd..ddbcc313 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -62,40 +62,49 @@ def menu_backup(): 'GenericRepair': 'Repaired', }} - # Set ticket ID + # Set ticket Number os.system('cls') - ticket_id = get_ticket_number() + ticket_number = get_ticket_number() # Mount backup shares mount_backup_shares() # Select destination - dest = select_backup_destination() - if dest is None: - abort_to_main_menu('Aborting Backup Creation') + destination = select_backup_destination() + if not destination: + raise GenericAbort # Select disk to backup disk = select_disk('For which drive are we creating backups?') - if disk is None: - abort_to_main_menu() - prep_disk_for_backup(dest, disk, ticket_id) + if not disk: + raise GenericAbort + + # "Prep" disk? + prep_disk_for_backup(destination, disk, ticket_number) # Display details for backup task os.system('cls') - print('Create Backup - Details:\n') - print(' Ticket: \t{ticket_id}'.format(ticket_id=ticket_id)) - print(' Source: \t[{Table}] ({Type}) {Name} {Size}\n'.format(**disk)) - print(' Destination:\t{name}'.format(name=dest.get('Display Name', dest['Name']))) + print_info('Create Backup - Details:\n') + # def show_info(message='~Some message~', info='~Some info~', indent=8, width=32): + show_info(message='Ticket:', info=ticket_number) + show_info( + message = 'Source:', + info = '[{Table}] ({Type}) {Name} {Size}'.format(**disk), + ) + show_info( + message = 'Destination:', + info = destination.get('Display Name', destination['Name']), + ) for par in disk['Partitions']: - print(par['Display String']) - print(disk['Backup Warnings']) + show_info(message='', info=par['Display String'], width=20) + print_standard(disk['Backup Warnings']) # Ask to proceed if (not ask('Proceed with backup?')): - abort_to_main_menu('Aborting Backup Creation') + raise GenericAbort # Backup partition(s) - print('\n\nStarting task.\n') + print_info('\n\nStarting task.\n') for par in disk['Partitions']: message = 'Partition {} Backup...'.format(par['Number']) result = try_and_print(message=message, function=backup_partition, @@ -106,9 +115,9 @@ def menu_backup(): # Verify backup(s) if disk['Valid Partitions'] > 1: - print('\n\n Verifying backups\n') + print_info('\n\n Verifying backups\n') else: - print('\n\n Verifying backup\n') + print_info('\n\n Verifying backup\n') for par in disk['Partitions']: if par['Number'] in disk['Bad Partitions']: continue # Skip verification @@ -121,9 +130,9 @@ def menu_backup(): if errors: print_warning('\nErrors were encountered and are detailed below.') for par in [p for p in disk['Partitions'] if 'Error' in p]: - print(' Partition {Number} Error:'.format(**par)) - for line in [line.strip() for line in par['Error'] if line.strip() != '']: - print_error('\t{line}'.format(line=line)) + print_standard(' Partition {} Error:'.format(par['Number'])) + for line in [line.strip() for line in par['Error'] if line.strip()]: + print_error('\t{}'.format(line)) time.sleep(30) else: print_success('\nNo errors were encountered during imaging.') @@ -171,7 +180,7 @@ def menu_setup(): # Set ticket ID os.system('cls') - ticket_id = get_ticket_number() + ticket_number = get_ticket_number() # Select the version of Windows to apply windows_version = select_windows_version() @@ -188,7 +197,7 @@ def menu_setup(): # Display details for setup task os.system('cls') print('Setup Windows - Details:\n') - print(' Ticket: \t{ticket_id}'.format(ticket_id=ticket_id)) + print(' Ticket: \t{ticket_number}'.format(ticket_number=ticket_number)) print(' Installing: \t{winver}'.format(winver=windows_version['Name'])) print(' Boot Method:\t{_type}'.format( _type='UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)')) diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py index 542114a5..9ff1508d 100644 --- a/Scripts/winpe_root_menu.py +++ b/Scripts/winpe_root_menu.py @@ -14,6 +14,8 @@ global_vars['LogFile'] = r'{LogDir}\WinPE.log'.format(**global_vars) if __name__ == '__main__': try: menu_root() + except GenericAbort: + pause('Press Enter to return to main menu... ') except SystemExit: pass except: From c043c3398d7de8dbb14bf68d570f0bb5a8d2027d Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 19:56:42 -0800 Subject: [PATCH 09/23] Reordered functions --- Scripts/functions/data.py | 74 +++++++++++++++++++-------------------- 1 file changed, 37 insertions(+), 37 deletions(-) diff --git a/Scripts/functions/data.py b/Scripts/functions/data.py index 63825293..d180c699 100644 --- a/Scripts/functions/data.py +++ b/Scripts/functions/data.py @@ -388,43 +388,6 @@ def select_destination(folder_path, prompt='Select destination'): return path -def select_volume(title='Select disk', auto_select=True): - """Select disk from attached disks. returns dict.""" - actions = [{'Name': 'Quit', 'Letter': 'Q'}] - disks = [] - - # Build list of disks - set_thread_error_mode(silent=True) # Prevents "No disk" popups - for d in psutil.disk_partitions(): - info = { - 'Disk': d, - 'Name': d.mountpoint} - try: - usage = psutil.disk_usage(d.device) - free = '{free} / {total} available'.format( - free = human_readable_size(usage.free, 2), - total = human_readable_size(usage.total, 2)) - except Exception: - # Meh, leaving unsupported destinations out - pass - # free = 'Unknown' - # info['Disabled'] = True - else: - info['Display Name'] = '{} ({})'.format(info['Name'], free) - disks.append(info) - set_thread_error_mode(silent=False) # Return to normal - - # Skip menu? - if len(disks) == 1 and auto_select: - return disks[0] - - # Show menu - selection = menu_select(title, main_entries=disks, action_entries=actions) - if selection == 'Q': - exit_script() - else: - return disks[int(selection)-1] - def select_source(ticket_number): """Select backup from those found on the BACKUP_SERVERS for the ticket.""" selected_source = None @@ -536,6 +499,43 @@ def select_source(ticket_number): # Done return selected_source +def select_volume(title='Select disk', auto_select=True): + """Select disk from attached disks. returns dict.""" + actions = [{'Name': 'Quit', 'Letter': 'Q'}] + disks = [] + + # Build list of disks + set_thread_error_mode(silent=True) # Prevents "No disk" popups + for d in psutil.disk_partitions(): + info = { + 'Disk': d, + 'Name': d.mountpoint} + try: + usage = psutil.disk_usage(d.device) + free = '{free} / {total} available'.format( + free = human_readable_size(usage.free, 2), + total = human_readable_size(usage.total, 2)) + except Exception: + # Meh, leaving unsupported destinations out + pass + # free = 'Unknown' + # info['Disabled'] = True + else: + info['Display Name'] = '{} ({})'.format(info['Name'], free) + disks.append(info) + set_thread_error_mode(silent=False) # Return to normal + + # Skip menu? + if len(disks) == 1 and auto_select: + return disks[0] + + # Show menu + selection = menu_select(title, main_entries=disks, action_entries=actions) + if selection == 'Q': + exit_script() + else: + return disks[int(selection)-1] + def set_thread_error_mode(silent=True): """Disable or Enable Windows error message dialogs. 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 10/23] 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 62e45d19..4ab89a43 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 ddbcc313..2f405325 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 9ff1508d..ad19b34b 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: From b96e5f3be6d2bb2a7e2db8144809eda8f036df7d Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 20:31:37 -0800 Subject: [PATCH 11/23] updated verify_wim_backup() --- Scripts/functions/backup.py | 30 ++++++++++-------------------- 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py index 4ab89a43..c884638a 100644 --- a/Scripts/functions/backup.py +++ b/Scripts/functions/backup.py @@ -137,26 +137,16 @@ def select_backup_destination(auto_select=True): else: return destinations[int(selection)-1] -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}" --nocheck'.format(bin=bin, **par) - if not os.path.exists('{Image Path}'.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 +def verify_wim_backup(partition): + if not os.path.exists(partition['Image Path']): + raise PathNotFoundError + cmd = [ + global_vars['Tools']['wimlib-imagex'], + 'verify', + partition['Image Path'], + ' --nocheck', + ] + run_program(cmd) if __name__ == '__main__': print("This file is not meant to be called directly.") From 850a1fca738f747c5393a472513761fdd32be24c Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 20:31:46 -0800 Subject: [PATCH 12/23] updated menu_backup() --- Scripts/functions/winpe_menus.py | 42 ++++++++++++++++++++++---------- 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index 2f405325..e14c0a22 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -56,6 +56,7 @@ def menu_backup(): other_results = { 'Error': { 'CalledProcessError': 'Unknown Error', + 'PathNotFoundError': 'Missing', }, 'Warning': { 'GenericAbort': 'Skipped', @@ -112,25 +113,40 @@ def menu_backup(): par['Error'] = result['Error'] # Verify backup(s) - if disk['Valid Partitions'] > 1: - print_info('\n\n Verifying backups\n') - else: - print_info('\n\n Verifying backup\n') - for par in disk['Partitions']: - if par['Number'] in disk['Bad Partitions']: - continue # Skip verification - try: - verify_wim_backup(bin, par) - except BackupError: - errors = True + if disk['Valid Partitions']: + print_info('\n\n Verifying backup images(s)\n') + for par in disk['Partitions']: + if par['Number'] in disk['Bad Partitions']: + continue # Skip verification + message = 'Partition {} Image...'.format(par['Number']) + result = try_and_print(message=message, function=verify_wim_backup, + other_results=other_results, partition=par) + if not result['CS']: + errors = True + par['Error'] = result['Error'] # Print summary if errors: print_warning('\nErrors were encountered and are detailed below.') for par in [p for p in disk['Partitions'] if 'Error' in p]: print_standard(' Partition {} Error:'.format(par['Number'])) - for line in [line.strip() for line in par['Error'] if line.strip()]: - print_error('\t{}'.format(line)) + if hasattr(par['Error'], 'stderr'): + try: + par['Error'] = par['Error'].stderr.decode() + except: + # Deal with badly formatted error message + pass + if isinstance(par['Error'], basestring): + print_error('\t{}'.format(par['Error'])) + else: + try: + par['Error'] = par['Error'].splitlines() + par['Error'] = [line.strip() for line in par['Error']] + par['Error'] = [line for line in par['Error'] if line] + except: + pass + for line in par['Error']: + print_error('\t{}'.format(line)) time.sleep(30) else: print_success('\nNo errors were encountered during imaging.') From 4ed6d41d1065bdecbfde2f40f9f083c87ee4c8fe Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 21:08:52 -0800 Subject: [PATCH 13/23] Updated menu / title sections * Added set_title() which sets window title and global_vars['Title'] * menu_select() will now display global_vars['Title'] above title * (If it exists) * Bugfix: fixed a few bad calls of menu_select() --- Scripts/functions/backup.py | 4 +++- Scripts/functions/common.py | 11 +++++++++++ Scripts/functions/disk.py | 7 +++++-- Scripts/functions/windows_setup.py | 5 ++++- Scripts/functions/winpe_menus.py | 23 +++++++++++++++-------- Scripts/winpe_root_menu.py | 2 +- 6 files changed, 39 insertions(+), 13 deletions(-) diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py index c884638a..c0e31370 100644 --- a/Scripts/functions/backup.py +++ b/Scripts/functions/backup.py @@ -131,7 +131,9 @@ def select_backup_destination(auto_select=True): return destinations[0] selection = menu_select( - 'Where are we backing up to?', destinations, actions) + title = 'Where are we backing up to?', + main_entries = destinations, + action_entries = actions) if selection == 'M': raise GenericAbort else: diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py index 4cf147e7..5204b45f 100644 --- a/Scripts/functions/common.py +++ b/Scripts/functions/common.py @@ -224,6 +224,10 @@ def menu_select(title='~ Untitled Menu ~', if not main_entries and not action_entries: raise Exception("MenuError: No items given") + # Set title + if 'Title' in global_vars: + title = '{}\n\n{}'.format(global_vars['Title']) + # Build menu menu_splash = '{}\n\n'.format(title) width = len(str(len(main_entries))) @@ -361,6 +365,13 @@ def run_program(cmd, args=[], check=True, pipe=True, shell=False): return process_return +def set_title(title='~Some Title~'): + """Set title. + + Used for window title and menu titles.""" + global_vars['Title'] = title + os.system('title {}'.format(title)) + def show_info(message='~Some message~', info='~Some info~', indent=8, width=32): """Display info with formatting.""" print_standard('{indent}{message:<{width}}{info}'.format( diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index 3b7730c0..5ea4be45 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -300,7 +300,7 @@ def remove_volume_letters(keep=''): except subprocess.CalledProcessError: pass -def select_disk(prompt='Which disk?'): +def select_disk(title='Which disk?'): """Select a disk from the attached disks""" disks = get_attached_disk_info() @@ -334,7 +334,10 @@ def select_disk(prompt='Which disk?'): ] # Menu loop - selection = menu_select(prompt, disk_options, actions) + selection = menu_select( + title = title, + main_entries = disk_options, + action_entries = actions) if (selection.isnumeric()): return disk_options[int(selection)-1]['Disk'] diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py index 199999a4..aa1ced35 100644 --- a/Scripts/functions/windows_setup.py +++ b/Scripts/functions/windows_setup.py @@ -180,7 +180,10 @@ def select_windows_version(): actions = [{'Name': 'Main Menu', 'Letter': 'M'},] # Menu loop - selection = menu_select('Which version of Windows are we installing?', WINDOWS_VERSIONS, actions) + selection = menu_select( + title = 'Which version of Windows are we installing?', + main_entries = WINDOWS_VERSIONS, + action_entries = actions) if selection.isnumeric(): return WINDOWS_VERSIONS[int(selection)-1] diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index e14c0a22..419ef08e 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -62,6 +62,7 @@ def menu_backup(): 'GenericAbort': 'Skipped', 'GenericRepair': 'Repaired', }} + set_title('{}: Backup Menu'.format(KIT_NAME_FULL)) # Set ticket Number os.system('cls') @@ -154,7 +155,6 @@ def menu_backup(): pause('\nPress Enter to return to main menu... ') def menu_root(): - title = '{}: Main Menu'.format(KIT_NAME_FULL) menus = [ {'Name': 'Create Backups', 'Menu': menu_backup}, {'Name': 'Setup Windows', 'Menu': menu_setup}, @@ -168,11 +168,12 @@ def menu_root(): # Main loop while True: + set_title(KIT_NAME_FULL) selection = menu_select( - title=title, - main_entries=menus, - action_entries=actions, - secret_exit=True) + title = 'Main Menu', + main_entries = menus, + action_entries = actions, + secret_exit = True) if (selection.isnumeric()): try: @@ -191,6 +192,7 @@ def menu_root(): def menu_setup(): """Format a drive, partition for MBR or GPT, apply a Windows image, and rebuild the boot files""" errors = False + set_title('{}: Setup Menu'.format(KIT_NAME_FULL)) # Set ticket ID os.system('cls') @@ -295,13 +297,16 @@ def menu_setup(): pause('\nPress Enter to return to main menu... ') def menu_tools(): - title = '{}: Tools Menu'.format(KIT_NAME_FULL) tools = [k for k in sorted(PE_TOOLS.keys())] actions = [{'Name': 'Main Menu', 'Letter': 'M'},] + set_title(KIT_NAME_FULL) # Menu loop while True: - selection = menu_select(title, tools, actions) + selection = menu_select( + title = 'Tools Menu', + main_entries = tools, + action_entries = actions) if (selection.isnumeric()): tool = tools[int(selection)-1] cmd = [PE_TOOLS[tool]['Path']] + PE_TOOLS[tool].get('Args', []) @@ -342,7 +347,9 @@ def select_minidump_path(): return None # Menu - selection = menu_select('Which BSoD / MiniDump path are we scanning?', dumps, []) + selection = menu_select( + title = 'Which BSoD / MiniDump path are we scanning?', + main_entries = dumps) return dumps[int(selection) - 1]['Name'] if __name__ == '__main__': diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py index ad19b34b..4bd39aef 100644 --- a/Scripts/winpe_root_menu.py +++ b/Scripts/winpe_root_menu.py @@ -8,7 +8,7 @@ os.chdir(os.path.dirname(os.path.realpath(__file__))) sys.path.append(os.getcwd()) from functions.winpe_menus import * init_global_vars() -os.system('title {}: Root Menu'.format(KIT_NAME_FULL)) +set_title('{}: Root Menu'.format(KIT_NAME_FULL)) global_vars['LogFile'] = r'{LogDir}\WinPE.log'.format(**global_vars) if __name__ == '__main__': From deb1e8f4fd20b2b67ede84a41611a56851c748e5 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 21:33:55 -0800 Subject: [PATCH 14/23] Volume letter updates * Added reassign_letter() * Attempts to reassign a volume to better ensure predictable letters * i.e. Local Windows source volume letter == 'I' * Adjusted code to avoid the "hidden" assign_volume_letters() call in select_disk() --- Scripts/functions/disk.py | 24 +++++++++++++++++++----- Scripts/functions/winpe_menus.py | 15 +++++++++------ 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index 5ea4be45..f19bb81a 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -4,6 +4,7 @@ from functions.common import * import partition_uids def assign_volume_letters(): + remove_volume_letters() try: # Run script with open(DISKPART_SCRIPT, 'w') as script: @@ -19,9 +20,6 @@ def get_attached_disk_info(): disks = [] print_info('Getting drive info...') - # Assign all the letters - assign_volume_letters() - # Get disks disks = get_disks() @@ -286,8 +284,24 @@ def prep_disk_for_formatting(disk=None): q='"' if par['Name'] != '' else '', **par) -def remove_volume_letters(keep=''): - if keep is None: +def reassign_volume_letter(letter, new_letter='I'): + if not letter: + # Ignore + return None + try: + # Run script + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select volume {}\n'.format(letter)) + script.write('remove noerr\n') + script.write('assign letter={}\n'.format(new_letter)) + run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + except subprocess.CalledProcessError: + pass + else: + return new_letter + +def remove_volume_letters(keep=None): + if not keep: keep = '' try: # Run script diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index 419ef08e..ceda3aa1 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -75,6 +75,7 @@ def menu_backup(): destination = select_backup_destination() # Select disk to backup + assign_volume_letters() disk = select_disk('For which drive are we creating backups?') if not disk: raise GenericAbort @@ -200,16 +201,15 @@ def menu_setup(): # Select the version of Windows to apply windows_version = select_windows_version() - - # Select drive to use as the OS drive - dest_disk = select_disk('To which drive are we installing Windows?') - prep_disk_for_formatting(dest_disk) # Find Windows image - ## NOTE: Needs to happen AFTER select_disk() is called as there's a hidden assign_volume_letters(). - ## This changes the current letters thus preventing installing from a local source. windows_image = find_windows_image(bin, windows_version) + # Select drive to use as the OS drive + assign_volume_letters() + dest_disk = select_disk('To which drive are we installing Windows?') + prep_disk_for_formatting(dest_disk) + # Display details for setup task os.system('cls') print('Setup Windows - Details:\n') @@ -235,6 +235,9 @@ def menu_setup(): # Release currently used volume letters (ensures that the drives will get S, T, & W as needed below) remove_volume_letters(keep=windows_image['Source']) + new_letter = reassign_volume_letter(letter=windows_image['Source']) + if new_letter: + windows_image['Source'] = new_letter # Format and partition drive print('\n Formatting Drive... \t\t', end='', flush=True) From e9ff02375f97cea418c85f10f7fee88e2377ad68 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 21:52:51 -0800 Subject: [PATCH 15/23] More letter updates --- Scripts/functions/disk.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index f19bb81a..4a265300 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -4,13 +4,16 @@ from functions.common import * import partition_uids def assign_volume_letters(): + with open(DISKPART_SCRIPT, 'w') as script: + for vol in get_volumes(): + script.write('select volume {Number}\n'.format(**vol)) + script.write('assign\n') + + # Remove current letters remove_volume_letters() + + # Run script try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - for vol in get_volumes(): - script.write('select volume {Number}\n'.format(**vol)) - script.write('assign\n') run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) except subprocess.CalledProcessError: pass @@ -303,13 +306,14 @@ def reassign_volume_letter(letter, new_letter='I'): def remove_volume_letters(keep=None): if not keep: keep = '' + with open(DISKPART_SCRIPT, 'w') as script: + for vol in get_volumes(): + if vol['Letter'].upper() != keep.upper(): + script.write('select volume {Number}\n'.format(**vol)) + script.write('remove noerr\n') + + # Run script try: - # Run script - with open(DISKPART_SCRIPT, 'w') as script: - for vol in get_volumes(): - if vol['Letter'].upper() != keep.upper(): - script.write('select volume {Number}\n'.format(**vol)) - script.write('remove noerr\n') run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) except subprocess.CalledProcessError: pass 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 16/23] 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 From 6903078ee017ae602fae1d599be683b62fa4cadc Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 22:19:31 -0800 Subject: [PATCH 17/23] updated prep_disk_for_formatting() --- Scripts/functions/disk.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index 15343053..3b2ff179 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -248,18 +248,28 @@ def prep_disk_for_formatting(disk=None): if len(disk['Partitions']) == 0: disk['Format Warnings'] += 'No partitions found\n' for par in disk['Partitions']: - if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): - # FileSystem not accessible to WinPE. List partition type / OS info for technician - par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format( - width=width, - q='"' if par['Name'] != '' else '', - **par) + display = ' Partition {num:>{width}}:\t{size} {fs}'.format( + num = par['Number'], + width = width, + size = par['Size'], + fs = par['FileSystem']) + + if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']): + # Set display string using partition description & OS type + display += '\t\t{q}{name}{q}\t{desc} ({os})'.format( + display = display, + q = '"' if par['Name'] != '' else '', + name = par['Name'], + desc = par['Description'], + os = par['OS']) else: - # FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician - par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format( - width=width, - q='"' if par['Name'] != '' else '', - **par) + # List space used instead of partition description & OS type + display += ' (Used: {used})\t{q}{name}{q}'.format( + used = par['Used Space'], + q = '"' if par['Name'] != '' else '', + name = par['Name']) + # For all partitions + par['Display String'] = display def reassign_volume_letter(letter, new_letter='I'): if not letter: @@ -320,7 +330,7 @@ def select_disk(title='Which disk?', disks): disk_options = [] for disk in disks: display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) - if len(disk['Partitions']) > 0: + if disk['Partitions']: pwidth=len(str(len(disk['Partitions']))) for par in disk['Partitions']: # Show unsupported partition(s) in RED From a0460d6a82775179ec3596429508e92a4f999c52 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 22:35:16 -0800 Subject: [PATCH 18/23] updated select_disk() --- Scripts/functions/disk.py | 37 ++++++++++++++++++------------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index 3b2ff179..d216847e 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -330,26 +330,25 @@ def select_disk(title='Which disk?', disks): disk_options = [] for disk in disks: display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) - if disk['Partitions']: - pwidth=len(str(len(disk['Partitions']))) - for par in disk['Partitions']: - # Show unsupported partition(s) in RED - par_skip = False - if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE): - par_skip = True - if par_skip: - display_name += COLORS['YELLOW'] + pwidth=len(str(len(disk['Partitions']))) + for par in disk['Partitions']: + # Main text + p_name = 'Partition {num:>{width}}: {size} ({fs})'.format( + num = par['Number'], + width = pwidth, + size = par['Size'], + fs = par['FileSystem']) + if par['Name']: + p_name += '\t"{}"'.format(par['Name']) + + # Show unsupported partition(s) + if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']): + p_display_name = '{YELLOW}{display}{CLEAR}'.format(display=p_name, **COLORS) + + display_name += '\n\t\t\t{}'.format(display_name) + if not disk['Partitions']: + display_name += '\n\t\t\t{YELLOW}No partitions found.{CLEAR}'.format(**COLORS) - # Main text - display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par) - if par['Name'] != '': - display_name += '\t"{Name}"'.format(**par) - - # Clear color (if set above) - if par_skip: - display_name += COLORS['CLEAR'] - else: - display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS) disk_options.append({'Name': display_name, 'Disk': disk}) actions = [ {'Name': 'Main Menu', 'Letter': 'M'}, From c9513804823fc449df646d91e2caca3338b096f8 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 23:06:23 -0800 Subject: [PATCH 19/23] functions\disk.py done --- Scripts/functions/disk.py | 239 +++++++++++++++++++------------------- 1 file changed, 118 insertions(+), 121 deletions(-) diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index d216847e..50277ca2 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -3,10 +3,18 @@ from functions.common import * import partition_uids +# Regex +REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE) +REGEX_DISK_GPT = re.compile( + r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', + re.IGNORECASE) +REGEX_DISK_MBR = re.compile(r'Disk ID: [A-Z0-9]+', re.IGNORECASE) +REGEX_DISK_RAW = re.compile(r'Disk ID: 00000000', re.IGNORECASE) + def assign_volume_letters(): with open(DISKPART_SCRIPT, 'w') as script: for vol in get_volumes(): - script.write('select volume {Number}\n'.format(**vol)) + script.write('select volume {}\n'.format(vol['Number'])) script.write('assign\n') # Remove current letters @@ -14,14 +22,15 @@ def assign_volume_letters(): # Run script try: - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + run_program(['diskpart', '/s', DISKPART_SCRIPT]) except subprocess.CalledProcessError: pass def get_boot_mode(): boot_mode = 'Legacy' try: - reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 'System\\CurrentControlSet\\Control') + reg_key = winreg.OpenKey( + winreg.HKEY_LOCAL_MACHINE, r'System\CurrentControlSet\Control') reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] if reg_value == 2: boot_mode = 'UEFI' @@ -30,32 +39,25 @@ def get_boot_mode(): return boot_mode -def get_disk_details(disk=None): +def get_disk_details(disk): details = {} - - # Bail early - if disk is None: - raise Exception('Disk not specified.') + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {}\n'.format(disk['Number'])) + script.write('detail disk\n') try: # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('detail disk\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() + output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: # Remove empty lines - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - + tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] # Set disk name details['Name'] = tmp[4] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + # Split each line on ':' skipping those without ':' tmp = [s.split(':') for s in tmp if ':' in s] - # Add key/value pairs to the details variable and return dict details.update({key.strip(): value.strip() for (key, value) in tmp}) @@ -63,90 +65,85 @@ def get_disk_details(disk=None): def get_disks(): disks = [] + with open(DISKPART_SCRIPT, 'w') as script: + script.write('list disk\n') try: # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('list disk\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() + output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: # Append disk numbers - for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return): - _num = tmp[0] - _size = human_readable_size(tmp[1]) - disks.append({'Number': _num, 'Size': _size}) + for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output): + num = tmp[0] + size = human_readable_size(tmp[1]) + disks.append({'Number': num, 'Size': size}) return disks -def get_partition_details(disk=None, par=None): +def get_partition_details(disk, partition): details = {} - - # Bail early - if disk is None: - raise Exception('Disk not specified.') - if par is None: - raise Exception('Partition not specified.') + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {}\n'.format(disk['Number'])) + script.write('select partition {}\n'.format(partition['Number'])) + script.write('detail partition\n') # Diskpart details try: # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('select partition {Number}\n'.format(**par)) - script.write('detail partition\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() + output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: # Get volume letter or RAW status - tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return) + tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output) if tmp: if tmp.group(1).upper() == 'RAW': details['FileSystem'] = RAW else: details['Letter'] = tmp.group(1) - - # Remove empty lines from process_return - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + # Remove empty lines from output + tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] + # Split each line on ':' skipping those without ':' tmp = [s.split(':') for s in tmp if ':' in s] - # Add key/value pairs to the details variable and return dict details.update({key.strip(): value.strip() for (key, value) in tmp}) # Get MBR type / GPT GUID for extra details on "Unknown" partitions guid = partition_uids.lookup_guid(details['Type']) - if guid is not None: + if guid: details.update({ 'Description': guid.get('Description', ''), 'OS': guid.get('OS', '')}) if 'Letter' in details: # Disk usage - tmp = shutil.disk_usage('{Letter}:\\'.format(**details)) + tmp = shutil.disk_usage('{}:\\'.format(details['Letter'])) details['Used Space'] = human_readable_size(tmp.used) # fsutil details + cmd = [ + 'fsutil', + 'fsinfo', + 'volumeinfo', + '{}:'.format(details['Letter']) + ] try: - process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details)) - process_return = process_return.stdout.decode().strip() + output = run_program(cmd) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: - # Remove empty lines from process_return - tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] - + # Remove empty lines from output + tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] # Add "Feature" lines - details['File System Features'] = [s.strip() for s in tmp if ':' not in s] - - # Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair + details['File System Features'] = [s.strip() for s in tmp + if ':' not in s] + # Split each line on ':' skipping those without ':' tmp = [s.split(':') for s in tmp if ':' in s] - # Add key/value pairs to the details variable and return dict details.update({key.strip(): value.strip() for (key, value) in tmp}) @@ -159,74 +156,72 @@ def get_partition_details(disk=None, par=None): return details -def get_partitions(disk=None): +def get_partitions(disk): partitions = [] - - # Bail early - if disk is None: - raise Exception('Disk not specified.') + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {}\n'.format(disk['Number'])) + script.write('list partition\n') try: # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('list partition\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() + output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: # Append partition numbers - for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE): - _num = tmp[0] - _size = human_readable_size(tmp[1]) - partitions.append({'Number': _num, 'Size': _size}) + regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+' + for tmp in re.findall(regex, output, re.IGNORECASE): + num = tmp[0] + size = human_readable_size(tmp[1]) + partitions.append({'Number': num, 'Size': size}) return partitions -def get_table_type(disk=None): - _type = 'Unknown' - - # Bail early - if disk is None: - raise Exception('Disk not specified.') +def get_table_type(disk): + part_type = 'Unknown' + with open(DISKPART_SCRIPT, 'w') as script: + script.write('select disk {}\n'.format(disk['Number'])) + script.write('uniqueid disk\n') try: - with open(DISKPART_SCRIPT, 'w') as script: - script.write('select disk {Number}\n'.format(**disk)) - script.write('uniqueid disk\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() + output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: - if re.findall(r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', process_return, re.IGNORECASE): - _type = 'GPT' - elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE): - _type = 'RAW' - elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE): - _type = 'MBR' + if REGEX_DISK_GPT.search(output): + part_type = 'GPT' + elif REGEX_DISK_MBR.search(output): + part_type = 'MBR' + elif REGEX_DISK_RAW.search(output): + part_type = 'RAW' + else: + part_type = 'Unknown - return _type + return part_type def get_volumes(): vols = [] + with open(DISKPART_SCRIPT, 'w') as script: + script.write('list volume\n') try: # Run script - with open(DISKPART_SCRIPT, 'w') as script: - script.write('list volume\n') - process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) - process_return = process_return.stdout.decode().strip() + output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) + output = output.stdout.decode().strip() except subprocess.CalledProcessError: pass else: # Append volume numbers - for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return): + for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output): vols.append({'Number': tmp[0], 'Letter': tmp[1]}) return vols +def is_bad_partition(par): + return 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']) + def prep_disk_for_formatting(disk=None): disk['Format Warnings'] = '\n' width = len(str(len(disk['Partitions']))) @@ -247,29 +242,29 @@ def prep_disk_for_formatting(disk=None): # Set Display and Warning Strings if len(disk['Partitions']) == 0: disk['Format Warnings'] += 'No partitions found\n' - for par in disk['Partitions']: + for partition in disk['Partitions']: display = ' Partition {num:>{width}}:\t{size} {fs}'.format( - num = par['Number'], + num = partition['Number'], width = width, - size = par['Size'], - fs = par['FileSystem']) + size = partition['Size'], + fs = partition['FileSystem']) - if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']): + if is_bad_partition(partition): # Set display string using partition description & OS type display += '\t\t{q}{name}{q}\t{desc} ({os})'.format( display = display, - q = '"' if par['Name'] != '' else '', - name = par['Name'], - desc = par['Description'], - os = par['OS']) + q = '"' if partition['Name'] != '' else '', + name = partition['Name'], + desc = partition['Description'], + os = partition['OS']) else: # List space used instead of partition description & OS type display += ' (Used: {used})\t{q}{name}{q}'.format( - used = par['Used Space'], - q = '"' if par['Name'] != '' else '', - name = par['Name']) + used = partition['Used Space'], + q = '"' if partition['Name'] != '' else '', + name = partition['Name']) # For all partitions - par['Display String'] = display + partition['Display String'] = display def reassign_volume_letter(letter, new_letter='I'): if not letter: @@ -281,7 +276,7 @@ def reassign_volume_letter(letter, new_letter='I'): script.write('select volume {}\n'.format(letter)) script.write('remove noerr\n') script.write('assign letter={}\n'.format(new_letter)) - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + run_program(['diskpart', '/s', DISKPART_SCRIPT]) except subprocess.CalledProcessError: pass else: @@ -293,12 +288,12 @@ def remove_volume_letters(keep=None): with open(DISKPART_SCRIPT, 'w') as script: for vol in get_volumes(): if vol['Letter'].upper() != keep.upper(): - script.write('select volume {Number}\n'.format(**vol)) + script.write('select volume {}\n'.format(vol['Number'])) script.write('remove noerr\n') # Run script try: - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + run_program(['diskpart', '/s', DISKPART_SCRIPT]) except subprocess.CalledProcessError: pass @@ -317,9 +312,9 @@ def scan_disks(): # Get partition info for disk disk['Partitions'] = get_partitions(disk) - for par in disk['Partitions']: + for partition in disk['Partitions']: # Get partition details - par.update(get_partition_details(disk, par)) + partition.update(get_partition_details(disk, partition)) # Done return disks @@ -331,23 +326,25 @@ def select_disk(title='Which disk?', disks): for disk in disks: display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) pwidth=len(str(len(disk['Partitions']))) - for par in disk['Partitions']: + for partition in disk['Partitions']: # Main text p_name = 'Partition {num:>{width}}: {size} ({fs})'.format( - num = par['Number'], + num = partition['Number'], width = pwidth, - size = par['Size'], - fs = par['FileSystem']) - if par['Name']: - p_name += '\t"{}"'.format(par['Name']) + size = partition['Size'], + fs = partition['FileSystem']) + if partition['Name']: + p_name += '\t"{}"'.format(partition['Name']) # Show unsupported partition(s) - if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']): - p_display_name = '{YELLOW}{display}{CLEAR}'.format(display=p_name, **COLORS) + if is_bad_partition(partition): + p_display_name = '{YELLOW}{display}{CLEAR}'.format( + display=p_name, **COLORS) display_name += '\n\t\t\t{}'.format(display_name) if not disk['Partitions']: - display_name += '\n\t\t\t{YELLOW}No partitions found.{CLEAR}'.format(**COLORS) + display_name += '\n\t\t\t{}No partitions found.{}'.format( + COLORS['YELLOW'], COLORS['CLEAR']) disk_options.append({'Name': display_name, 'Disk': disk}) actions = [ From c09d7ab603ba206e1305bcfffb3f40a2c1dc1ff0 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Thu, 30 Nov 2017 23:06:41 -0800 Subject: [PATCH 20/23] updated functions\backup.py --- Scripts/functions/backup.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/Scripts/functions/backup.py b/Scripts/functions/backup.py index c0e31370..a562f4f5 100644 --- a/Scripts/functions/backup.py +++ b/Scripts/functions/backup.py @@ -1,9 +1,8 @@ # Wizard Kit PE: Functions - Backup -from functions.common import * +from functions.disk import * # Regex -REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE) REGEX_BAD_PATH_NAMES = re.compile( r'([<>:"/\\\|\?\*]' r'|^(CON|PRN|AUX|NUL|COM\d*|LPT\d*)$)' @@ -37,7 +36,7 @@ def prep_disk_for_backup(destination, disk, ticket_number): # Get partition totals disk['Bad Partitions'] = [par['Number'] for par in disk['Partitions'] - if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem'])] + if is_bad_partition(partition)] num_valid_partitions = len(disk['Partitions']) - len(disk['Bad Partitions']) disk['Valid Partitions'] = num_valid_partitions if disk['Valid Partitions'] <= 0: From 33924c183e08bcdc9c16e92682d18182fd8337e6 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Fri, 1 Dec 2017 09:59:19 -0800 Subject: [PATCH 21/23] Updated find_windows_image() & setup_windows() * Merged File and Ext dict entries * Using psutil instead of mountvol --- Scripts/functions/windows_setup.py | 91 ++++++++++++++---------------- Scripts/functions/winpe_menus.py | 6 +- 2 files changed, 45 insertions(+), 52 deletions(-) diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py index aa1ced35..008c649d 100644 --- a/Scripts/functions/windows_setup.py +++ b/Scripts/functions/windows_setup.py @@ -43,48 +43,48 @@ WINDOWS_VERSIONS = [ 'Family': '10'}, ] -def find_windows_image(bin, windows_version): - """Search for a Windows source image file on local drives and network drives (in that order)""" +def find_windows_image(windows_version): + """Search for a Windows source image file, returns dict. + + Searches on local drives and then the WINDOWS_SERVER share.""" image = {} imagefile = windows_version['Image File'] + imagename = windows_version['Image Name'] # Search local source - process_return = run_program('mountvol') - for tmp in re.findall(r'.*([A-Za-z]):\\', process_return.stdout.decode()): + for d in psutil.disk_partitions(): for ext in ['esd', 'wim', 'swm']: - filename = '{drive}:\\images\\{imagefile}'.format(drive=tmp[0], imagefile=imagefile) - filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) - if os.path.isfile(filename_ext): - if wim_contains_image(bin, filename_ext, windows_version['Image Name']): - image['Ext'] = ext - image['File'] = filename - image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' - image['Source'] = tmp[0] - break + path = '{}images\{}.{}'.format(d.mountpoint, imagefile, ext) + if os.path.isfile(path) and wim_contains_image(path, imagename): + image['Path'] = path + image['Source'] = letter + if ext == 'swm': + image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4]) + break - # Check for network source (if necessary) + # Check for network source if not image: mount_windows_share() if not WINDOWS_SERVER['Mounted']: return None for ext in ['esd', 'wim', 'swm']: - filename = '\\\\{IP}\\{Share}\\images\\{imagefile}'.format(imagefile=imagefile, **WINDOWS_SERVER) - filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext) - if os.path.isfile(filename_ext): - if wim_contains_image(bin, filename_ext, windows_version['Image Name']): - image['Ext'] = ext - image['File'] = filename - image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else '' - image['Source'] = None - break + path = r'\\{}\{}\images\{}.ext'.format( + WINDOWS_SERVER['IP'], WINDOWS_SERVER['Share'], imagefile, ext) + if os.path.isfile(path) and wim_contains_image(path, imagename): + image['Path'] = path + image['Source'] = None + if ext == 'swm': + image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4]) + break # Display image to be used (if any) and return - if any(image): - print_info('Using image: {File}.{Ext}'.format(**image)) + if image: + print_info('Using image: {}'.format(image['Path'])) return image else: - print_error('Failed to find Windows source image for {winver}'.format(winver=windows_version['Name'])) - abort_to_main_menu('Aborting Windows setup') + print_error('Failed to find Windows source image for {}'.format( + windows_version['Name'])) + raise GeneralAbort def format_gpt(disk=None, windows_family=None): """Format disk for use as a Windows OS drive using the GPT (UEFI) layout.""" @@ -190,17 +190,15 @@ def select_windows_version(): elif selection == 'M': abort_to_main_menu() -def setup_windows(bin=None, windows_image=None, windows_version=None): - # Bail early - if bin is None: - raise Exception('bin path not specified.') - if windows_image is None: - raise Exception('Windows image not specified.') - if windows_version is None: - raise Exception('Windows version not specified.') - - # Apply image - cmd = '{bin}\\wimlib\\wimlib-imagex apply "{File}.{Ext}" "{Image Name}" W:\\ {Glob}'.format(bin=bin, **windows_image, **windows_version) +def setup_windows(windows_image, windows_version): + cmd = [ + global_vars['Tools']['wimlib-imagex'], + 'apply', + windows_image['Path'], + windows_version['Image Name'], + 'W:\\'] + if 'Glob' in windows_image: + cmd.extend(windows_image['Glob']) run_program(cmd) def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T'): @@ -226,20 +224,15 @@ def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T') def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'): run_program('bcdboot {win}:\\Windows /s {sys}: /f {mode}'.format(win=windows_letter, sys=system_letter, mode=mode)) -def wim_contains_image(bin=None, filename=None, imagename=None): - # Bail early - if bin is None: - raise Exception('bin not specified.') - if filename is None: - raise Exception('Filename not specified.') - if imagename is None: - raise Exception('Image Name not specified.') - - cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename) +def wim_contains_image(filename, imagename): + cmd = [ + global_vars['Tools']['wimlib-imagex'], + 'info', + filename, + imagename] try: run_program(cmd) except subprocess.CalledProcessError: - print_error('Invalid image: {filename}'.format(filename=filename)) return False return True diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index 69b551fe..4b7d873b 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -211,7 +211,7 @@ def menu_setup(): windows_version = select_windows_version() # Find Windows image - windows_image = find_windows_image(bin, windows_version) + windows_image = find_windows_image(windows_version) # Scan disks try_and_print(message='Assigning letters...', function=assign_volume_letters, other_results=other_results) @@ -237,7 +237,7 @@ def menu_setup(): print(' Installing: \t{winver}'.format(winver=windows_version['Name'])) print(' Boot Method:\t{_type}'.format( _type='UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)')) - print(' Using Image:\t{File}.{Ext}'.format(**windows_image)) + print(' Using Image:\t{}'.format(windows_image['Path'])) print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format(**dest_disk)) for par in dest_disk['Partitions']: print_warning(par['Display String']) @@ -275,7 +275,7 @@ def menu_setup(): # Apply Image print(' Applying Image... \t\t', end='', flush=True) try: - setup_windows(bin, windows_image, windows_version) + setup_windows(windows_image, windows_version) print_success('Complete.') except subprocess.CalledProcessError: print_error('Failed.') From deb7c76ffbb0e12d058aa95d4a14749113c88ecc Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Fri, 1 Dec 2017 10:12:59 -0800 Subject: [PATCH 22/23] windows_setup.py done --- Scripts/functions/windows_setup.py | 81 +++++++++++++----------------- 1 file changed, 34 insertions(+), 47 deletions(-) diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py index 008c649d..1830b4eb 100644 --- a/Scripts/functions/windows_setup.py +++ b/Scripts/functions/windows_setup.py @@ -86,25 +86,17 @@ def find_windows_image(windows_version): windows_version['Name'])) raise GeneralAbort -def format_gpt(disk=None, windows_family=None): - """Format disk for use as a Windows OS drive using the GPT (UEFI) layout.""" - - # Bail early - if disk is None: - raise Exception('No disk provided.') - if windows_family is None: - raise Exception('No Windows family provided.') - - # Format drive - # print_info('Drive will use a GPT (UEFI) layout.') +def format_gpt(disk, windows_family): + """Format disk for use as a Windows OS drive using the GPT layout.""" with open(DISKPART_SCRIPT, 'w') as script: # Partition table - script.write('select disk {number}\n'.format(number=disk['Number'])) + script.write('select disk {}\n'.format(disk['Number'])) script.write('clean\n') script.write('convert gpt\n') # System partition - script.write('create partition efi size=260\n') # NOTE: Allows for Advanced Format 4K drives + # NOTE: ESP needs to be >= 260 for Advanced Format 4K drives + script.write('create partition efi size=500\n') script.write('format quick fs=fat32 label="System"\n') script.write('assign letter="S"\n') @@ -126,23 +118,14 @@ def format_gpt(disk=None, windows_family=None): script.write('gpt attributes=0x8000000000000001\n') # Run script - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + run_program(['diskpart', '/s', DISKPART_SCRIPT]) time.sleep(2) -def format_mbr(disk=None, windows_family=None): - """Format disk for use as a Windows OS drive using the MBR (legacy) layout.""" - - # Bail early - if disk is None: - raise Exception('No disk provided.') - if windows_family is None: - raise Exception('No Windows family provided.') - - # Format drive - # print_info('Drive will use a MBR (legacy) layout.') +def format_mbr(disk, windows_family): + """Format disk for use as a Windows OS drive using the MBR layout.""" with open(DISKPART_SCRIPT, 'w') as script: # Partition table - script.write('select disk {number}\n'.format(number=disk['Number'])) + script.write('select disk {}\n'.format(disk['Number'])) script.write('clean\n') # System partition @@ -165,7 +148,7 @@ def format_mbr(disk=None, windows_family=None): script.write('set id=27\n') # Run script - run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) + run_program(['diskpart', '/s', DISKPART_SCRIPT]) time.sleep(2) def mount_windows_share(): @@ -177,7 +160,9 @@ def mount_windows_share(): mount_network_share(WINDOWS_SERVER) def select_windows_version(): - actions = [{'Name': 'Main Menu', 'Letter': 'M'},] + actions = [ + {'Name': 'Main Menu', 'Letter': 'M'}, + ] # Menu loop selection = menu_select( @@ -188,7 +173,7 @@ def select_windows_version(): if selection.isnumeric(): return WINDOWS_VERSIONS[int(selection)-1] elif selection == 'M': - abort_to_main_menu() + raise GeneralAbort def setup_windows(windows_image, windows_version): cmd = [ @@ -201,28 +186,30 @@ def setup_windows(windows_image, windows_version): cmd.extend(windows_image['Glob']) run_program(cmd) -def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T'): - # Bail early - if windows_version is None: - raise Exception('Windows version not specified.') +def setup_windows_re(windows_version, windows_letter='W', tools_letter='T'): + win = r'{}:\Windows'.format(windows_letter) + winre = r'{}\System32\Recovery\WinRE.wim'.format(win) + dest = r'{}:\Recovery\WindowsRE'.format(tools_letter) - _win = '{win}:\\Windows'.format(win=windows_letter) - _winre = '{win}\\System32\\Recovery\\WinRE.wim'.format(win=_win) - _dest = '{tools}:\\Recovery\\WindowsRE'.format(tools=tools_letter) + # Copy WinRE.wim + os.makedirs(dest, exist_ok=True) + shutil.copy(winre, r'{}\WinRE.wim'.format(dest)) - if re.search(r'^(8|10)', windows_version['Family']): - # Copy WinRE.wim - os.makedirs(_dest, exist_ok=True) - shutil.copy(_winre, '{dest}\\WinRE.wim'.format(dest=_dest)) - - # Set location - run_program('{win}\\System32\\reagentc /setreimage /path {dest} /target {win}'.format(dest=_dest, win=_win)) - else: - # Only supported on Windows 8 and above - raise SetupError + # Set location + cmd = [ + r'{}\System32\ReAgentc.exe'.format(win), + '/setreimage', + '/path', dest, + '/target', win] + run_program(cmd) def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'): - run_program('bcdboot {win}:\\Windows /s {sys}: /f {mode}'.format(win=windows_letter, sys=system_letter, mode=mode)) + cmd = [ + r'{}:\Windows\System32\bcdboot.exe'.format(windows_letter), + r'{}:\Windows'.format(windows_letter), + '/s', '{}:'.format(system_letter), + '/f', mode] + run_program(cmd) def wim_contains_image(filename, imagename): cmd = [ From 90c4189942ce73cbef5366b6f9dec970518d7b59 Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Fri, 1 Dec 2017 11:05:40 -0800 Subject: [PATCH 23/23] Updated winpe_menus.py & cleanup * Removed "Windows Family" logic * Win7 will now have recovery tools setup * Added format_disk() --- Scripts/functions/disk.py | 2 +- Scripts/functions/windows_setup.py | 71 ++++++------ Scripts/functions/winpe_menus.py | 179 +++++++++++++++-------------- 3 files changed, 125 insertions(+), 127 deletions(-) diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py index 50277ca2..eb2bda4f 100644 --- a/Scripts/functions/disk.py +++ b/Scripts/functions/disk.py @@ -360,7 +360,7 @@ def select_disk(title='Which disk?', disks): if (selection.isnumeric()): return disk_options[int(selection)-1]['Disk'] elif (selection == 'M'): - abort_to_main_menu() + raise GeneralAbort if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py index 1830b4eb..cbbe1525 100644 --- a/Scripts/functions/windows_setup.py +++ b/Scripts/functions/windows_setup.py @@ -7,46 +7,38 @@ DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP']) WINDOWS_VERSIONS = [ {'Name': 'Windows 7 Home Basic', 'Image File': 'Win7', - 'Image Name': 'Windows 7 HOMEBASIC', - 'Family': '7'}, + 'Image Name': 'Windows 7 HOMEBASIC'}, {'Name': 'Windows 7 Home Premium', 'Image File': 'Win7', - 'Image Name': 'Windows 7 HOMEPREMIUM', - 'Family': '7'}, + 'Image Name': 'Windows 7 HOMEPREMIUM'}, {'Name': 'Windows 7 Professional', 'Image File': 'Win7', - 'Image Name': 'Windows 7 PROFESSIONAL', - 'Family': '7'}, + 'Image Name': 'Windows 7 PROFESSIONAL'}, {'Name': 'Windows 7 Ultimate', 'Image File': 'Win7', - 'Image Name': 'Windows 7 ULTIMATE', - 'Family': '7'}, + 'Image Name': 'Windows 7 ULTIMATE'}, {'Name': 'Windows 8.1', 'Image File': 'Win8', - 'Image Name': 'Windows 8.1', - 'Family': '8', + 'Image Name': 'Windows 8.1', 'CRLF': True}, {'Name': 'Windows 8.1 Pro', 'Image File': 'Win8', - 'Image Name': 'Windows 8.1 Pro', - 'Family': '8'}, + 'Image Name': 'Windows 8.1 Pro'}, {'Name': 'Windows 10 Home', 'Image File': 'Win10', - 'Image Name': 'Windows 10 Home', - 'Family': '10', + 'Image Name': 'Windows 10 Home', 'CRLF': True}, {'Name': 'Windows 10 Pro', 'Image File': 'Win10', - 'Image Name': 'Windows 10 Pro', - 'Family': '10'}, + 'Image Name': 'Windows 10 Pro'}, ] def find_windows_image(windows_version): """Search for a Windows source image file, returns dict. - Searches on local drives and then the WINDOWS_SERVER share.""" + Searches on local disks and then the WINDOWS_SERVER share.""" image = {} imagefile = windows_version['Image File'] imagename = windows_version['Image Name'] @@ -86,8 +78,15 @@ def find_windows_image(windows_version): windows_version['Name'])) raise GeneralAbort -def format_gpt(disk, windows_family): - """Format disk for use as a Windows OS drive using the GPT layout.""" +def format_disk(disk, use_gpt): + """Format disk for use as a Windows OS disk.""" + if use_gpt: + format_gpt(disk) + else: + format_mbr(disk) + +def format_gpt(disk): + """Format disk for use as a Windows OS disk using the GPT layout.""" with open(DISKPART_SCRIPT, 'w') as script: # Partition table script.write('select disk {}\n'.format(disk['Number'])) @@ -95,7 +94,7 @@ def format_gpt(disk, windows_family): script.write('convert gpt\n') # System partition - # NOTE: ESP needs to be >= 260 for Advanced Format 4K drives + # NOTE: ESP needs to be >= 260 for Advanced Format 4K disks script.write('create partition efi size=500\n') script.write('format quick fs=fat32 label="System"\n') script.write('assign letter="S"\n') @@ -108,21 +107,20 @@ def format_gpt(disk, windows_family): script.write('format quick fs=ntfs label="Windows"\n') script.write('assign letter="W"\n') - # Recovery Tools partition (Windows 8+) - if re.search(r'^(8|10)', windows_family): - script.write('shrink minimum=500\n') - script.write('create partition primary\n') - script.write('format quick fs=ntfs label="Recovery Tools"\n') - script.write('assign letter="T"\n') - script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n') - script.write('gpt attributes=0x8000000000000001\n') + # Recovery Tools partition + script.write('shrink minimum=500\n') + script.write('create partition primary\n') + script.write('format quick fs=ntfs label="Recovery Tools"\n') + script.write('assign letter="T"\n') + script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n') + script.write('gpt attributes=0x8000000000000001\n') # Run script run_program(['diskpart', '/s', DISKPART_SCRIPT]) time.sleep(2) -def format_mbr(disk, windows_family): - """Format disk for use as a Windows OS drive using the MBR layout.""" +def format_mbr(disk): + """Format disk for use as a Windows OS disk using the MBR layout.""" with open(DISKPART_SCRIPT, 'w') as script: # Partition table script.write('select disk {}\n'.format(disk['Number'])) @@ -139,13 +137,12 @@ def format_mbr(disk, windows_family): script.write('format fs=ntfs quick label="Windows"\n') script.write('assign letter="W"\n') - # Recovery Tools partition (Windows 8+) - if re.search(r'^(8|10)', windows_family): - script.write('shrink minimum=500\n') - script.write('create partition primary\n') - script.write('format quick fs=ntfs label="Recovery"\n') - script.write('assign letter="T"\n') - script.write('set id=27\n') + # Recovery Tools partition + script.write('shrink minimum=500\n') + script.write('create partition primary\n') + script.write('format quick fs=ntfs label="Recovery"\n') + script.write('assign letter="T"\n') + script.write('set id=27\n') # Run script run_program(['diskpart', '/s', DISKPART_SCRIPT]) diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py index 4b7d873b..c6ba4e8c 100644 --- a/Scripts/functions/winpe_menus.py +++ b/Scripts/functions/winpe_menus.py @@ -5,7 +5,7 @@ from functions.disk import * from functions.windows_setup import * # STATIC VARIABLES -FAST_COPY_ARGS = [ +FAST_COPY_PE_ARGS = [ '/cmd=noexist_only', '/utf8', '/skip_empty_dir', @@ -20,7 +20,7 @@ PE_TOOLS = { }, 'FastCopy': { 'Path': r'FastCopy\FastCopy.exe', - 'Args': FAST_COPY_ARGS, + 'Args': FAST_COPY_PE_ARGS, }, 'HWiNFO': { 'Path': r'HWiNFO\HWiNFO.exe', @@ -51,7 +51,7 @@ PE_TOOLS = { } def menu_backup(): - """Take backup images of partition(s) in the WIM format and save them to a backup share""" + """Take backup images of partition(s) in the WIM format.""" errors = False other_results = { 'Error': { @@ -65,7 +65,7 @@ def menu_backup(): set_title('{}: Backup Menu'.format(KIT_NAME_FULL)) # Set ticket Number - os.system('cls') + clear_screen() ticket_number = get_ticket_number() # Mount backup shares @@ -75,8 +75,14 @@ def menu_backup(): 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) + 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, + other_results = other_results) if result['CS']: disks = result['Out'] else: @@ -84,7 +90,7 @@ def menu_backup(): raise GenericAbort # Select disk to backup - disk = select_disk('For which drive are we creating backups?', disks) + disk = select_disk('For which disk are we creating backups?', disks) if not disk: raise GenericAbort @@ -92,9 +98,8 @@ def menu_backup(): prep_disk_for_backup(destination, disk, ticket_number) # Display details for backup task - os.system('cls') + clear_screen() print_info('Create Backup - Details:\n') - # def show_info(message='~Some message~', info='~Some info~', indent=8, width=32): show_info(message='Ticket:', info=ticket_number) show_info( message = 'Source:', @@ -115,9 +120,12 @@ def menu_backup(): # Backup partition(s) print_info('\n\nStarting task.\n') for par in disk['Partitions']: - message = 'Partition {} Backup...'.format(par['Number']) - result = try_and_print(message=message, function=backup_partition, - other_results=other_results, disk=disk, partition=par) + result = try_and_print( + message = 'Partition {} Backup...'.format(par['Number']), + function = backup_partition, + other_results = other_results, + disk = disk, + partition = par) if not result['CS']: errors = True par['Error'] = result['Error'] @@ -128,9 +136,11 @@ def menu_backup(): for par in disk['Partitions']: if par['Number'] in disk['Bad Partitions']: continue # Skip verification - message = 'Partition {} Image...'.format(par['Number']) - result = try_and_print(message=message, function=verify_wim_backup, - other_results=other_results, partition=par) + result = try_and_print( + message = 'Partition {} Image...'.format(par['Number']), + function = verify_wim_backup, + other_results = other_results, + partition = par) if not result['CS']: errors = True par['Error'] = result['Error'] @@ -199,12 +209,12 @@ def menu_root(): exit_script() def menu_setup(): - """Format a drive, partition for MBR or GPT, apply a Windows image, and rebuild the boot files""" + """Format a disk (MBR/GPT), apply a Windows image, and setup boot files.""" errors = False set_title('{}: Setup Menu'.format(KIT_NAME_FULL)) # Set ticket ID - os.system('cls') + clear_screen() ticket_number = get_ticket_number() # Select the version of Windows to apply @@ -214,16 +224,22 @@ def menu_setup(): windows_image = find_windows_image(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) + 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, + 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 - dest_disk = select_disk('To which drive are we installing Windows?', disks) + # Select disk to use as the OS disk + dest_disk = select_disk('To which disk are we installing Windows?', disks) if not disk: raise GenericAbort @@ -231,92 +247,77 @@ def menu_setup(): prep_disk_for_formatting(dest_disk) # Display details for setup task - os.system('cls') - print('Setup Windows - Details:\n') - print(' Ticket: \t{ticket_number}'.format(ticket_number=ticket_number)) - print(' Installing: \t{winver}'.format(winver=windows_version['Name'])) - print(' Boot Method:\t{_type}'.format( - _type='UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)')) - print(' Using Image:\t{}'.format(windows_image['Path'])) - print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format(**dest_disk)) + clear_screen() + print_info('Setup Windows - Details:\n') + show_info(message='Ticket:', info=ticket_number) + show_info(message='Installing:', info=windows_version['Name']) + show_info( + message = 'Boot Method:', + info = 'UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)') + show_info(message='Using Image:', info=windows_version['Path']) + print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format( + **dest_disk)) for par in dest_disk['Partitions']: print_warning(par['Display String']) print_warning(dest_disk['Format Warnings']) if (not ask('Is this correct?')): - abort_to_main_menu('Aborting Windows setup') + raise GeneralAbort - # Safety check - print('\nSAFETY CHECK') - print_warning('All data will be DELETED from the drive and partition(s) listed above.') - print_warning('This is irreversible and will lead to {CLEAR}{RED}DATA LOSS.'.format(**COLORS)) + # Safety check + print_standard('\nSAFETY CHECK') + print_warning('All data will be DELETED from the ' + 'disk & partition(s) listed above.') + print_warning('This is irreversible and will lead ' + 'to {CLEAR}{RED}DATA LOSS.'.format(**COLORS)) if (not ask('Asking again to confirm, is this correct?')): - abort_to_main_menu('Aborting Windows setup') + raise GeneralAbort - # Release currently used volume letters (ensures that the drives will get S, T, & W as needed below) + # Remove volume letters so S, T, & W can be used below remove_volume_letters(keep=windows_image['Source']) new_letter = reassign_volume_letter(letter=windows_image['Source']) if new_letter: windows_image['Source'] = new_letter - # Format and partition drive - print('\n Formatting Drive... \t\t', end='', flush=True) - try: - if (dest_disk['Use GPT']): - format_gpt(dest_disk, windows_version['Family']) - else: - format_mbr(dest_disk, windows_version['Family']) - print_success('Complete.') - except: - # We need to crash as the drive is in an unknown state - print_error('Failed.') - raise + # Format and partition disk + result = try_and_print( + message = 'Formatting disk...', + function = format_disk, + other_results = other_results, + disk = dest_disk, + use_gpt = dest_disk['Use GPT']) + if not result['CS']: + # We need to crash as the disk is in an unknown state + print_error('ERROR: Failed to format disk.') + raise GenericAbort # Apply Image - print(' Applying Image... \t\t', end='', flush=True) - try: - setup_windows(windows_image, windows_version) - print_success('Complete.') - except subprocess.CalledProcessError: - print_error('Failed.') - errors = True - except: - # We need to crash as the OS is in an unknown state - print_error('Failed.') - raise + result = try_and_print( + message = 'Applying image...', + function = setup_windows, + other_results = other_results, + windows_image = windows_image, + windows_version = windows_version) + if not result['CS']: + # We need to crash as the disk is in an unknown state + print_error('ERROR: Failed to apply image.') + raise GenericAbort # Create Boot files - print(' Update Boot Partition...\t\t', end='', flush=True) - try: - update_boot_partition() - print_success('Complete.') - except subprocess.CalledProcessError: - # Don't need to crash as this is (potentially) recoverable - print_error('Failed.') - errors = True - except: - print_error('Failed.') - raise + try_and_print( + message = 'Updating boot files...', + function = update_boot_partition, + other_results = other_results) # Setup WinRE - print(' Update Recovery Tools...\t\t', end='', flush=True) - try: - setup_windows_re(windows_version) - print_success('Complete.') - except SetupError: - print('Skipped.') - except: - # Don't need to crash as this is (potentially) recoverable - print_error('Failed.') - errors = True + try_and_print( + message = 'Updating recovery tools...', + function = setup_windows_re, + other_results = other_results, + windows_version = windows_version) # Print summary - if errors: - print_warning('\nErrors were encountered during setup.') - time.sleep(30) - else: - print_success('\nNo errors were encountered during setup.') - time.sleep(5) + print_standard('\nDone.') pause('\nPress Enter to return to main menu... ') def menu_tools(): @@ -359,9 +360,9 @@ def select_minidump_path(): # Remove RAMDisk letter if 'X' in tmp: tmp.remove('X') - for drive in tmp: - if os.path.exists('{drive}:\\Windows\\MiniDump'.format(drive=drive)): - dumps.append({'Name': '{drive}:\\Windows\\MiniDump'.format(drive=drive)}) + for disk in tmp: + if os.path.exists('{}:\\Windows\\MiniDump'.format(disk)): + dumps.append({'Name': '{}:\\Windows\\MiniDump'.format(disk)}) # Check results before showing menu if len(dumps) == 0: