Updated find_windows_image() & setup_windows()

* Merged File and Ext dict entries
* Using psutil instead of mountvol

Scripts/functions/windows_setup.py

@@ -43,48 +43,48 @@ WINDOWS_VERSIONS = [
         'Family': '10'},
     ]
 
-def find_windows_image(bin, windows_version):
+def find_windows_image(windows_version):
-    """Search for a Windows source image file on local drives and network drives (in that order)"""
+    """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 d in psutil.disk_partitions():
-    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)
+            path = '{}images\{}.{}'.format(d.mountpoint, imagefile, ext)
-            filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext)
+            if os.path.isfile(path) and wim_contains_image(path, imagename):
-            if os.path.isfile(filename_ext):
+                image['Path'] = path
-                if wim_contains_image(bin, filename_ext, windows_version['Image Name']):
+                image['Source'] = letter
-                    image['Ext'] = ext
+                if ext == 'swm':
-                    image['File'] = filename
+                    image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4])
-                    image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else ''
+                break
-                    image['Source'] = tmp[0]
-                    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)
+            path = r'\\{}\{}\images\{}.ext'.format(
-            filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext)
+                WINDOWS_SERVER['IP'], WINDOWS_SERVER['Share'], imagefile, ext)
-            if os.path.isfile(filename_ext):
+            if os.path.isfile(path) and wim_contains_image(path, imagename):
-                if wim_contains_image(bin, filename_ext, windows_version['Image Name']):
+                image['Path'] = path
-                    image['Ext'] = ext
+                image['Source'] = None
-                    image['File'] = filename
+                if ext == 'swm':
-                    image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else ''
+                    image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4])
-                    image['Source'] = None
+                break
-                    break
     
     # Display image to be used (if any) and return
-    if any(image):
+    if image:
-        print_info('Using image: {File}.{Ext}'.format(**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']))
+        print_error('Failed to find Windows source image for {}'.format(
-        abort_to_main_menu('Aborting Windows setup')
+            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):
+def setup_windows(windows_image, windows_version):
-    # Bail early
+    cmd = [
-    if bin is None:
+        global_vars['Tools']['wimlib-imagex'],
-        raise Exception('bin path not specified.')
+        'apply',
-    if windows_image is None:
+        windows_image['Path'],
-        raise Exception('Windows image not specified.')
+        windows_version['Image Name'],
-    if windows_version is None:
+        'W:\\']
-        raise Exception('Windows version not specified.')
+    if 'Glob' in windows_image:
-    
+        cmd.extend(windows_image['Glob'])
-    # 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'):
@@ -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):
+def wim_contains_image(filename, imagename):
-    # Bail early
+    cmd = [
-    if bin is None:
+        global_vars['Tools']['wimlib-imagex'],
-        raise Exception('bin not specified.')
+        'info',
-    if filename is None:
+        filename,
-        raise Exception('Filename not specified.')
+        imagename]
-    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

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