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] 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.')