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))