2Shirt/WizardKit
2Shirt/WizardKit
1
0
Fork 0
Code Issues 2 Pull requests Projects Releases 26 Packages Wiki Activity Actions

Added run_diskpart(script)

Browse source 
* Takes list, writes script, runs script, and returns result
This commit is contained in:
Alan Mason 2017-12-01 15:27:10 -08:00
parent a9822ae9bd
commit d075f17bfa
3 changed files with 106 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

124
.bin/Scripts/functions/disk.py
View file

@ -12,19 +12,16 @@ REGEX_DISK_MBR = re.compile(r'Disk ID: [A-Z0-9]+', re.IGNORECASE)
REGEX_DISK_RAW = re.compile(r'Disk ID: 00000000', re.IGNORECASE) REGEX_DISK_RAW = re.compile(r'Disk ID: 00000000', re.IGNORECASE)
def assign_volume_letters(): def assign_volume_letters():
with open(DISKPART_SCRIPT, 'w') as script:
for vol in get_volumes():
script.write('select volume {}\n'.format(vol['Number']))
script.write('assign\n')
# Remove current letters
remove_volume_letters() remove_volume_letters()
# Run script # Write script
try: script = []
run_program(['diskpart', '/s', DISKPART_SCRIPT]) for vol in get_volumes():
except subprocess.CalledProcessError: script.append('select volume {}'.format(vol['Number']))
pass script.append('assign')
# Run
run_diskpart(script)
def get_boot_mode(): def get_boot_mode():
boot_mode = 'Legacy' boot_mode = 'Legacy'
@ -41,17 +38,17 @@ def get_boot_mode():
def get_disk_details(disk): def get_disk_details(disk):
details = {} details = {}
with open(DISKPART_SCRIPT, 'w') as script: script = [
script.write('select disk {}\n'.format(disk['Number'])) 'select disk {}'.format(disk['Number']),
script.write('detail disk\n') 'detail disk']
# Run
try: try:
# Run script result = run_diskpart(script)
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
output = result.stdout.decode().strip()
# Remove empty lines # Remove empty lines
tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
# Set disk name # Set disk name
@ -65,17 +62,15 @@ def get_disk_details(disk):
def get_disks(): def get_disks():
disks = [] disks = []
with open(DISKPART_SCRIPT, 'w') as script:
script.write('list disk\n')
try: try:
# Run script # Run script
output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) result = run_diskpart(['list disk'])
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Append disk numbers # Append disk numbers
output = result.stdout.decode().strip()
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output): for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output):
num = tmp[0] num = tmp[0]
size = human_readable_size(tmp[1]) size = human_readable_size(tmp[1])
@ -85,20 +80,20 @@ def get_disks():
def get_partition_details(disk, partition): def get_partition_details(disk, partition):
details = {} details = {}
with open(DISKPART_SCRIPT, 'w') as script: script = [
script.write('select disk {}\n'.format(disk['Number'])) 'select disk {}'.format(disk['Number']),
script.write('select partition {}\n'.format(partition['Number'])) 'select partition {}'.format(partition['Number']),
script.write('detail partition\n') 'detail partition']
# Diskpart details # Diskpart details
try: try:
# Run script # Run script
output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) result = run_diskpart(script)
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Get volume letter or RAW status # Get volume letter or RAW status
output = result.stdout.decode().strip()
tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output) tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output)
if tmp: if tmp:
if tmp.group(1).upper() == 'RAW': if tmp.group(1).upper() == 'RAW':
@ -132,11 +127,11 @@ def get_partition_details(disk, partition):
'{}:'.format(details['Letter']) '{}:'.format(details['Letter'])
] ]
try: try:
output = run_program(cmd) result = run_program(cmd)
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
output = result.stdout.decode().strip()
# Remove empty lines from output # Remove empty lines from output
tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
# Add "Feature" lines # Add "Feature" lines
@ -158,18 +153,18 @@ def get_partition_details(disk, partition):
def get_partitions(disk): def get_partitions(disk):
partitions = [] partitions = []
with open(DISKPART_SCRIPT, 'w') as script: script = [
script.write('select disk {}\n'.format(disk['Number'])) 'select disk {}'.format(disk['Number']),
script.write('list partition\n') 'list partition']
try: try:
# Run script # Run script
output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) result = run_diskpart(script)
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Append partition numbers # Append partition numbers
output = result.stdout.decode().strip()
regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+' regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+'
for tmp in re.findall(regex, output, re.IGNORECASE): for tmp in re.findall(regex, output, re.IGNORECASE):
num = tmp[0] num = tmp[0]
@ -180,40 +175,34 @@ def get_partitions(disk):
def get_table_type(disk): def get_table_type(disk):
part_type = 'Unknown' part_type = 'Unknown'
with open(DISKPART_SCRIPT, 'w') as script: script = [
script.write('select disk {}\n'.format(disk['Number'])) 'select disk {}'.format(disk['Number']),
script.write('uniqueid disk\n') 'uniqueid disk']
try: try:
output = run_program(['diskpart', '/s', DISKPART_SCRIPT]) result = run_diskpart(script)
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
output = result.stdout.decode().strip()
if REGEX_DISK_GPT.search(output): if REGEX_DISK_GPT.search(output):
part_type = 'GPT' part_type = 'GPT'
elif REGEX_DISK_MBR.search(output): elif REGEX_DISK_MBR.search(output):
part_type = 'MBR' part_type = 'MBR'
elif REGEX_DISK_RAW.search(output): elif REGEX_DISK_RAW.search(output):
part_type = 'RAW' part_type = 'RAW'
else:
part_type = 'Unknown'
return part_type return part_type
def get_volumes(): def get_volumes():
vols = [] vols = []
with open(DISKPART_SCRIPT, 'w') as script:
script.write('list volume\n')
try: try:
# Run script result = run_diskpart(['list volume'])
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Append volume numbers # Append volume numbers
output = result.stdout.decode().strip()
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output): for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output):
vols.append({'Number': tmp[0], 'Letter': tmp[1]}) vols.append({'Number': tmp[0], 'Letter': tmp[1]})
@ -270,13 +259,12 @@ def reassign_volume_letter(letter, new_letter='I'):
if not letter: if not letter:
# Ignore # Ignore
return None return None
script = [
'select volume {}'.format(letter),
'remove noerr',
'assign letter={}'.format(new_letter)]
try: try:
# Run script run_diskpart(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', DISKPART_SCRIPT])
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
@ -285,18 +273,36 @@ def reassign_volume_letter(letter, new_letter='I'):
def remove_volume_letters(keep=None): def remove_volume_letters(keep=None):
if not keep: if not keep:
keep = '' keep = ''
with open(DISKPART_SCRIPT, 'w') as script:
for vol in get_volumes(): script = []
if vol['Letter'].upper() != keep.upper(): for vol in get_volumes():
script.write('select volume {}\n'.format(vol['Number'])) if vol['Letter'].upper() != keep.upper():
script.write('remove noerr\n') script.append('select volume {}'.format(vol['Number']))
script.append('remove noerr')
# Run script # Run script
try: try:
run_program(['diskpart', '/s', DISKPART_SCRIPT]) run_diskpart(script)
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
def run_diskpart(script):
tempfile = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
# Write script
with open(tempfile, 'w') as f:
for line in script:
f.write('{}\n'.format(line))
# Run script
cmd = [
r'{}\Windows\System32\diskpart.exe'.format(
global_vars['Env']['SYSTEMDRIVE']),
'/s', tempfile]
result = run_program(cmd)
time.sleep(2)
return result
def scan_disks(): def scan_disks():
"""Get details about the attached disks""" """Get details about the attached disks"""
disks = get_disks() disks = get_disks()

81
.bin/Scripts/functions/windows_setup.py
View file

@ -1,6 +1,7 @@
# Wizard Kit PE: Functions - Windows Setup # Wizard Kit PE: Functions - Windows Setup
from functions.data import * from functions.data import *
from functions.disk import *
# STATIC VARIABLES # STATIC VARIABLES
WINDOWS_VERSIONS = [ WINDOWS_VERSIONS = [
@ -89,66 +90,66 @@ def format_disk(disk, use_gpt):
def format_gpt(disk): def format_gpt(disk):
"""Format disk for use as a Windows OS disk using the GPT layout.""" """Format disk for use as a Windows OS disk using the GPT layout."""
with open(DISKPART_SCRIPT, 'w') as script: script = [
# Partition table # Partition table
script.write('select disk {}\n'.format(disk['Number'])) 'select disk {}'.format(disk['Number']),
script.write('clean\n') 'clean',
script.write('convert gpt\n') 'convert gpt',
# System partition # System partition
# NOTE: ESP needs to be >= 260 for Advanced Format 4K disks # NOTE: ESP needs to be >= 260 for Advanced Format 4K disks
script.write('create partition efi size=500\n') 'create partition efi size=500',
script.write('format quick fs=fat32 label="System"\n') 'format quick fs=fat32 label="System"',
script.write('assign letter="S"\n') 'assign letter="S"',
# Microsoft Reserved (MSR) partition # Microsoft Reserved (MSR) partition
script.write('create partition msr size=128\n') 'create partition msr size=128',
# Windows partition # Windows partition
script.write('create partition primary\n') 'create partition primary',
script.write('format quick fs=ntfs label="Windows"\n') 'format quick fs=ntfs label="Windows"',
script.write('assign letter="W"\n') 'assign letter="W"',
# Recovery Tools partition # Recovery Tools partition
script.write('shrink minimum=500\n') 'shrink minimum=500',
script.write('create partition primary\n') 'create partition primary',
script.write('format quick fs=ntfs label="Recovery Tools"\n') 'format quick fs=ntfs label="Recovery Tools"',
script.write('assign letter="T"\n') 'assign letter="T"',
script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n') 'set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"',
script.write('gpt attributes=0x8000000000000001\n') 'gpt attributes=0x8000000000000001',
]
# Run script
run_program(['diskpart', '/s', DISKPART_SCRIPT]) # Run
time.sleep(2) run_diskpart(script)
def format_mbr(disk): def format_mbr(disk):
"""Format disk for use as a Windows OS disk using the MBR layout.""" """Format disk for use as a Windows OS disk using the MBR layout."""
with open(DISKPART_SCRIPT, 'w') as script: script = [
# Partition table # Partition table
script.write('select disk {}\n'.format(disk['Number'])) 'select disk {}'.format(disk['Number']),
script.write('clean\n') 'clean',
# System partition # System partition
script.write('create partition primary size=100\n') 'create partition primary size=100',
script.write('format fs=ntfs quick label="System Reserved"\n') 'format fs=ntfs quick label="System Reserved"',
script.write('active\n') 'active',
script.write('assign letter="S"\n') 'assign letter="S"',
# Windows partition # Windows partition
script.write('create partition primary\n') 'create partition primary',
script.write('format fs=ntfs quick label="Windows"\n') 'format fs=ntfs quick label="Windows"',
script.write('assign letter="W"\n') 'assign letter="W"',
# Recovery Tools partition # Recovery Tools partition
script.write('shrink minimum=500\n') 'shrink minimum=500',
script.write('create partition primary\n') 'create partition primary',
script.write('format quick fs=ntfs label="Recovery"\n') 'format quick fs=ntfs label="Recovery"',
script.write('assign letter="T"\n') 'assign letter="T"',
script.write('set id=27\n') 'set id=27',
]
# Run script
run_program(['diskpart', '/s', DISKPART_SCRIPT]) # Run
time.sleep(2) run_diskpart(script)
def mount_windows_share(): def mount_windows_share():
"""Mount the Windows images share unless labeled as already mounted.""" """Mount the Windows images share unless labeled as already mounted."""

3
.bin/Scripts/winpe_root_menu.py
View file

@ -11,9 +11,6 @@ init_global_vars()
set_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) global_vars['LogFile'] = r'{LogDir}\WinPE.log'.format(**global_vars)
# STATIC VARIABLES
DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
if __name__ == '__main__': if __name__ == '__main__':
try: try:
menu_root() menu_root()