204 lines
9.0 KiB
Python
204 lines
9.0 KiB
Python
# To reproduce sgdisk behaviour in v1 of AIF-NG:
|
|
# https://gist.github.com/herry13/5931cac426da99820de843477e41e89e
|
|
# https://github.com/dcantrell/pyparted/blob/master/examples/query_device_capacity.py
|
|
# TODO: Remember to replicate genfstab behaviour.
|
|
|
|
import os
|
|
import re
|
|
try:
|
|
# https://stackoverflow.com/a/34812552/733214
|
|
# https://github.com/karelzak/util-linux/blob/master/libmount/python/test_mount_context.py#L6
|
|
import libmount as mount
|
|
except ImportError:
|
|
# We should never get here. util-linux is part of core (base) in Arch and uses "libmount".
|
|
import pylibmount as mount
|
|
##
|
|
import blkinfo
|
|
import parted # https://www.gnu.org/software/parted/api/index.html
|
|
import psutil
|
|
##
|
|
import aif.constants
|
|
import aif.utils
|
|
|
|
# TODO: https://serverfault.com/questions/356534/ssd-erase-block-size-lvm-pv-on-raw-device-alignment
|
|
# TODO: caveats? https://gist.github.com/leodutra/8779d468e9062058a3e90008295d3ca6
|
|
# https://unix.stackexchange.com/questions/325886/bios-gpt-do-we-need-a-boot-flag
|
|
|
|
|
|
class Partition(object):
|
|
def __init__(self, part_xml, diskobj, start_sector, partnum, tbltype, part_type = None):
|
|
if tbltype not in ('gpt', 'msdos'):
|
|
raise ValueError('{0} must be one of gpt or msdos'.format(tbltype))
|
|
if tbltype == 'msdos' and part_type not in ('primary', 'extended', 'logical'):
|
|
raise ValueError(('You must specify if this is a '
|
|
'primary, extended, or logical partition for msdos partition tables'))
|
|
self.xml = part_xml
|
|
self.id = self.xml.attrib['id']
|
|
self.flags = set()
|
|
for f in self.xml.findall('partitionFlag'):
|
|
if f.text in aif.constants.PARTED_FLAGS:
|
|
self.flags.add(f.text)
|
|
self.flags = sorted(list(self.flags))
|
|
self.partnum = partnum
|
|
if tbltype == 'msdos':
|
|
if partnum > 4:
|
|
self.part_type = parted.PARTITION_LOGICAL
|
|
else:
|
|
if part_type == 'extended':
|
|
self.part_type = parted.PARTITION_EXTENDED
|
|
elif part_type == 'logical':
|
|
self.part_type = parted.PARTITION_LOGICAL
|
|
else:
|
|
self.part_type = parted.PARTITION_NORMAL
|
|
self.fs_type = self.xml.attrib['fsType'].lower()
|
|
if self.fs_type not in aif.constants.PARTED_FSTYPES:
|
|
raise ValueError(('{0} is not a valid partition filesystem type; '
|
|
'must be one of: {1}').format(self.xml.attrib['fsType'],
|
|
', '.join(sorted(aif.constants.PARTED_FSTYPES))))
|
|
self.disk = diskobj
|
|
self.device = self.disk.device
|
|
self.devpath = '{0}{1}'.format(self.device.path, self.partnum)
|
|
self.is_hiformatted = False
|
|
sizes = {}
|
|
for s in ('start', 'stop'):
|
|
x = dict(zip(('from_bgn', 'size', 'type'),
|
|
aif.utils.convertSizeUnit(self.xml.attrib[s])))
|
|
sectors = x['size']
|
|
if x['type'] == '%':
|
|
sectors = int(self.device.getLength() * (0.01 * x['size']))
|
|
else:
|
|
sectors = int(aif.utils.size.convertStorage(x['size'],
|
|
x['type'],
|
|
target = 'B') / self.device.sectorSize)
|
|
sizes[s] = (sectors, x['from_bgn'])
|
|
if sizes['start'][1] is not None:
|
|
if sizes['start'][1]:
|
|
self.begin = sizes['start'][0] + 0
|
|
else:
|
|
self.begin = self.device.getLength() - sizes['start'][0]
|
|
else:
|
|
self.begin = sizes['start'][0] + start_sector
|
|
if sizes['stop'][1] is not None:
|
|
if sizes['stop'][1]:
|
|
self.end = sizes['stop'][0] + 0
|
|
else:
|
|
# This *technically* should be - 34, at least for gpt, but the alignment optimizer fixes it for us.
|
|
self.end = (self.device.getLength() - 1) - sizes['stop'][0]
|
|
else:
|
|
self.end = self.begin + sizes['stop'][0]
|
|
# TECHNICALLY we could craft the Geometry object with "length = ...", but it doesn't let us be explicit
|
|
# in configs. So we manually crunch the numbers and do it all at the end.
|
|
self.geometry = parted.Geometry(device = self.device,
|
|
start = self.begin,
|
|
end = self.end)
|
|
self.filesystem = parted.FileSystem(type = self.fs_type,
|
|
geometry = self.geometry)
|
|
self.partition = parted.Partition(disk = diskobj,
|
|
type = self.part_type,
|
|
geometry = self.geometry,
|
|
fs = self.filesystem)
|
|
for f in self.flags[:]:
|
|
flag_id = aif.constants.PARTED_FLAG_IDX[f]
|
|
if self.partition.isFlagAvailable(flag_id):
|
|
self.partition.setFlag(flag_id)
|
|
else:
|
|
self.flags.remove(f)
|
|
if tbltype == 'gpt' and self.xml.attrib.get('name'):
|
|
# The name attribute setting is b0rk3n, so we operate on the underlying PedPartition object.
|
|
# https://github.com/dcantrell/pyparted/issues/49#issuecomment-540096687
|
|
# https://github.com/dcantrell/pyparted/issues/65
|
|
# self.partition.name = self.xml.attrib.get('name')
|
|
_pedpart = self.partition.getPedPartition()
|
|
_pedpart.set_name(self.xml.attrib['name'])
|
|
#
|
|
# def detect(self):
|
|
# pass # TODO; blkinfo?
|
|
|
|
|
|
class Disk(object):
|
|
def __init__(self, disk_xml):
|
|
self.xml = disk_xml
|
|
self.id = self.xml.attrib['id']
|
|
self.devpath = os.path.realpath(self.xml.attrib['device'])
|
|
self.is_lowformatted = None
|
|
self.is_hiformatted = None
|
|
self.is_partitioned = None
|
|
self.partitions = None
|
|
self._initDisk()
|
|
|
|
def _initDisk(self):
|
|
if self.devpath == 'auto':
|
|
self.devpath = '/dev/{0}'.format(blkinfo.BlkDiskInfo().get_disks()[0]['kname'])
|
|
if not os.path.isfile(self.devpath):
|
|
raise ValueError('{0} does not exist; please specify an explicit device path'.format(self.devpath))
|
|
self.table_type = self.xml.attrib.get('diskFormat', 'gpt').lower()
|
|
if self.table_type in ('bios', 'mbr', 'dos'):
|
|
self.table_type = 'msdos'
|
|
validlabels = parted.getLabels()
|
|
if self.table_type not in validlabels:
|
|
raise ValueError(('Disk format {0} is not valid for this architecture;'
|
|
'must be one of: {1}'.format(self.table_type, ', '.join(list(validlabels)))))
|
|
self.device = parted.getDevice(self.devpath)
|
|
self.disk = parted.freshDisk(self.device, self.table_type)
|
|
self.is_lowformatted = False
|
|
self.is_hiformatted = False
|
|
self.is_partitioned = False
|
|
self.partitions = []
|
|
return()
|
|
|
|
def diskFormat(self):
|
|
if self.is_lowformatted:
|
|
return()
|
|
# This is a safeguard. We do *not* want to low-format a disk that is mounted.
|
|
aif.utils.checkMounted(self.devpath)
|
|
self.disk.deleteAllPartitions()
|
|
self.disk.commit()
|
|
self.is_lowformatted = True
|
|
self.is_partitioned = False
|
|
return()
|
|
|
|
def getPartitions(self):
|
|
# For GPT, this *technically* should be 34 -- or, more precisely, 2048 (see FAQ in manual), but the alignment
|
|
# optimizer fixes it for us automatically.
|
|
# But for DOS tables, it's required.
|
|
if self.table_type == 'msdos':
|
|
start_sector = 2048
|
|
else:
|
|
start_sector = 0
|
|
self.partitions = []
|
|
xml_partitions = self.xml.findall('part')
|
|
for idx, part in enumerate(xml_partitions):
|
|
partnum = idx + 1
|
|
if self.table_type == 'gpt':
|
|
p = Partition(part, self.disk, start_sector, partnum, self.table_type)
|
|
else:
|
|
parttype = 'primary'
|
|
if len(xml_partitions) > 4:
|
|
if partnum == 4:
|
|
parttype = 'extended'
|
|
elif partnum > 4:
|
|
parttype = 'logical'
|
|
p = Partition(part, self.disk, start_sector, partnum, self.table_type, part_type = parttype)
|
|
start_sector = p.end + 1
|
|
self.partitions.append(p)
|
|
return()
|
|
|
|
def partFormat(self):
|
|
if self.is_partitioned:
|
|
return()
|
|
if not self.is_lowformatted:
|
|
self.diskFormat()
|
|
# This is a safeguard. We do *not* want to partition a disk that is mounted.
|
|
aif.utils.checkMounted(self.devpath)
|
|
if not self.partitions:
|
|
self.getPartitions()
|
|
if not self.partitions:
|
|
return()
|
|
for p in self.partitions:
|
|
self.disk.addPartition(partition = p, constraint = self.device.optimalAlignedConstraint)
|
|
self.disk.commit()
|
|
p.devpath = p.partition.path
|
|
p.is_hiformatted = True
|
|
self.is_partitioned = True
|
|
return()
|