Added run_diskpart(script)
* Takes list, writes script, runs script, and returns result
This commit is contained in:
Alan Mason
parent
a9822ae9bd
commit
d075f17bfa
3 changed files with 106 additions and 102 deletions
|
|
@ -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)
|
||||
|
||||
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()
|
||||
|
||||
# Run script
|
||||
try:
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
# Write script
|
||||
script = []
|
||||
for vol in get_volumes():
|
||||
script.append('select volume {}'.format(vol['Number']))
|
||||
script.append('assign')
|
||||
|
||||
# Run
|
||||
run_diskpart(script)
|
||||
|
||||
def get_boot_mode():
|
||||
boot_mode = 'Legacy'
|
||||
|
|
@ -41,17 +38,17 @@ def get_boot_mode():
|
|||
|
||||
def get_disk_details(disk):
|
||||
details = {}
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('detail disk\n')
|
||||
script = [
|
||||
'select disk {}'.format(disk['Number']),
|
||||
'detail disk']
|
||||
|
||||
# Run
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_diskpart(script)
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
output = result.stdout.decode().strip()
|
||||
# Remove empty lines
|
||||
tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
|
||||
# Set disk name
|
||||
|
|
@ -65,17 +62,15 @@ def get_disk_details(disk):
|
|||
|
||||
def get_disks():
|
||||
disks = []
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list disk\n')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_diskpart(['list disk'])
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append disk numbers
|
||||
output = result.stdout.decode().strip()
|
||||
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output):
|
||||
num = tmp[0]
|
||||
size = human_readable_size(tmp[1])
|
||||
|
|
@ -85,20 +80,20 @@ def get_disks():
|
|||
|
||||
def get_partition_details(disk, partition):
|
||||
details = {}
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('select partition {}\n'.format(partition['Number']))
|
||||
script.write('detail partition\n')
|
||||
script = [
|
||||
'select disk {}'.format(disk['Number']),
|
||||
'select partition {}'.format(partition['Number']),
|
||||
'detail partition']
|
||||
|
||||
# Diskpart details
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_diskpart(script)
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Get volume letter or RAW status
|
||||
output = result.stdout.decode().strip()
|
||||
tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output)
|
||||
if tmp:
|
||||
if tmp.group(1).upper() == 'RAW':
|
||||
|
|
@ -132,11 +127,11 @@ def get_partition_details(disk, partition):
|
|||
'{}:'.format(details['Letter'])
|
||||
]
|
||||
try:
|
||||
output = run_program(cmd)
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
output = result.stdout.decode().strip()
|
||||
# Remove empty lines from output
|
||||
tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
|
||||
# Add "Feature" lines
|
||||
|
|
@ -158,18 +153,18 @@ def get_partition_details(disk, partition):
|
|||
|
||||
def get_partitions(disk):
|
||||
partitions = []
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('list partition\n')
|
||||
script = [
|
||||
'select disk {}'.format(disk['Number']),
|
||||
'list partition']
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_diskpart(script)
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append partition numbers
|
||||
output = result.stdout.decode().strip()
|
||||
regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+'
|
||||
for tmp in re.findall(regex, output, re.IGNORECASE):
|
||||
num = tmp[0]
|
||||
|
|
@ -180,40 +175,34 @@ def get_partitions(disk):
|
|||
|
||||
def get_table_type(disk):
|
||||
part_type = 'Unknown'
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('uniqueid disk\n')
|
||||
script = [
|
||||
'select disk {}'.format(disk['Number']),
|
||||
'uniqueid disk']
|
||||
|
||||
try:
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_diskpart(script)
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
output = result.stdout.decode().strip()
|
||||
if REGEX_DISK_GPT.search(output):
|
||||
part_type = 'GPT'
|
||||
elif REGEX_DISK_MBR.search(output):
|
||||
part_type = 'MBR'
|
||||
elif REGEX_DISK_RAW.search(output):
|
||||
part_type = 'RAW'
|
||||
else:
|
||||
part_type = 'Unknown'
|
||||
|
||||
return part_type
|
||||
|
||||
def get_volumes():
|
||||
vols = []
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list volume\n')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
output = output.stdout.decode().strip()
|
||||
result = run_diskpart(['list volume'])
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append volume numbers
|
||||
output = result.stdout.decode().strip()
|
||||
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output):
|
||||
vols.append({'Number': tmp[0], 'Letter': tmp[1]})
|
||||
|
||||
|
|
@ -270,13 +259,12 @@ def reassign_volume_letter(letter, new_letter='I'):
|
|||
if not letter:
|
||||
# Ignore
|
||||
return None
|
||||
script = [
|
||||
'select volume {}'.format(letter),
|
||||
'remove noerr',
|
||||
'assign letter={}'.format(new_letter)]
|
||||
try:
|
||||
# Run 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])
|
||||
run_diskpart(script)
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
|
|
@ -285,18 +273,36 @@ def reassign_volume_letter(letter, new_letter='I'):
|
|||
def remove_volume_letters(keep=None):
|
||||
if not keep:
|
||||
keep = ''
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
if vol['Letter'].upper() != keep.upper():
|
||||
script.write('select volume {}\n'.format(vol['Number']))
|
||||
script.write('remove noerr\n')
|
||||
|
||||
script = []
|
||||
for vol in get_volumes():
|
||||
if vol['Letter'].upper() != keep.upper():
|
||||
script.append('select volume {}'.format(vol['Number']))
|
||||
script.append('remove noerr')
|
||||
|
||||
# Run script
|
||||
try:
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
run_diskpart(script)
|
||||
except subprocess.CalledProcessError:
|
||||
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():
|
||||
"""Get details about the attached disks"""
|
||||
disks = get_disks()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# Wizard Kit PE: Functions - Windows Setup
|
||||
|
||||
from functions.data import *
|
||||
from functions.disk import *
|
||||
|
||||
# STATIC VARIABLES
|
||||
WINDOWS_VERSIONS = [
|
||||
|
|
@ -89,66 +90,66 @@ def format_disk(disk, use_gpt):
|
|||
|
||||
def format_gpt(disk):
|
||||
"""Format disk for use as a Windows OS disk using the GPT layout."""
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script = [
|
||||
# Partition table
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('clean\n')
|
||||
script.write('convert gpt\n')
|
||||
'select disk {}'.format(disk['Number']),
|
||||
'clean',
|
||||
'convert gpt',
|
||||
|
||||
# System partition
|
||||
# NOTE: ESP needs to be >= 260 for Advanced Format 4K disks
|
||||
script.write('create partition efi size=500\n')
|
||||
script.write('format quick fs=fat32 label="System"\n')
|
||||
script.write('assign letter="S"\n')
|
||||
'create partition efi size=500',
|
||||
'format quick fs=fat32 label="System"',
|
||||
'assign letter="S"',
|
||||
|
||||
# Microsoft Reserved (MSR) partition
|
||||
script.write('create partition msr size=128\n')
|
||||
'create partition msr size=128',
|
||||
|
||||
# Windows partition
|
||||
script.write('create partition primary\n')
|
||||
script.write('format quick fs=ntfs label="Windows"\n')
|
||||
script.write('assign letter="W"\n')
|
||||
'create partition primary',
|
||||
'format quick fs=ntfs label="Windows"',
|
||||
'assign letter="W"',
|
||||
|
||||
# Recovery Tools partition
|
||||
script.write('shrink minimum=500\n')
|
||||
script.write('create partition primary\n')
|
||||
script.write('format quick fs=ntfs label="Recovery Tools"\n')
|
||||
script.write('assign letter="T"\n')
|
||||
script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n')
|
||||
script.write('gpt attributes=0x8000000000000001\n')
|
||||
|
||||
# Run script
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
time.sleep(2)
|
||||
'shrink minimum=500',
|
||||
'create partition primary',
|
||||
'format quick fs=ntfs label="Recovery Tools"',
|
||||
'assign letter="T"',
|
||||
'set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"',
|
||||
'gpt attributes=0x8000000000000001',
|
||||
]
|
||||
|
||||
# Run
|
||||
run_diskpart(script)
|
||||
|
||||
def format_mbr(disk):
|
||||
"""Format disk for use as a Windows OS disk using the MBR layout."""
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script = [
|
||||
# Partition table
|
||||
script.write('select disk {}\n'.format(disk['Number']))
|
||||
script.write('clean\n')
|
||||
'select disk {}'.format(disk['Number']),
|
||||
'clean',
|
||||
|
||||
# System partition
|
||||
script.write('create partition primary size=100\n')
|
||||
script.write('format fs=ntfs quick label="System Reserved"\n')
|
||||
script.write('active\n')
|
||||
script.write('assign letter="S"\n')
|
||||
'create partition primary size=100',
|
||||
'format fs=ntfs quick label="System Reserved"',
|
||||
'active',
|
||||
'assign letter="S"',
|
||||
|
||||
# Windows partition
|
||||
script.write('create partition primary\n')
|
||||
script.write('format fs=ntfs quick label="Windows"\n')
|
||||
script.write('assign letter="W"\n')
|
||||
'create partition primary',
|
||||
'format fs=ntfs quick label="Windows"',
|
||||
'assign letter="W"',
|
||||
|
||||
# Recovery Tools partition
|
||||
script.write('shrink minimum=500\n')
|
||||
script.write('create partition primary\n')
|
||||
script.write('format quick fs=ntfs label="Recovery"\n')
|
||||
script.write('assign letter="T"\n')
|
||||
script.write('set id=27\n')
|
||||
|
||||
# Run script
|
||||
run_program(['diskpart', '/s', DISKPART_SCRIPT])
|
||||
time.sleep(2)
|
||||
'shrink minimum=500',
|
||||
'create partition primary',
|
||||
'format quick fs=ntfs label="Recovery"',
|
||||
'assign letter="T"',
|
||||
'set id=27',
|
||||
]
|
||||
|
||||
# Run
|
||||
run_diskpart(script)
|
||||
|
||||
def mount_windows_share():
|
||||
"""Mount the Windows images share unless labeled as already mounted."""
|
||||
|
|
|
|||
|
|
@ -11,9 +11,6 @@ init_global_vars()
|
|||
set_title('{}: Root Menu'.format(KIT_NAME_FULL))
|
||||
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__':
|
||||
try:
|
||||
menu_root()
|
||||
|
|
|
|||
Loading…
Reference in a new issue