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

@@ -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:

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 =     []

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

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

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 = [

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