summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorEdward Cragg <edward.cragg@codethink.co.uk>2015-08-04 12:17:22 +0100
committerPedro Alvarez <pedro.alvarez@codethink.co.uk>2015-09-22 14:43:54 +0000
commit746001dae609d752c1b32f511cbc632e8810f231 (patch)
tree9cdc7d8c8de99249fdfd7e6815b78fc02bc8e7ee
parentfd47533d721238f5b6dc7e4b557048ec149a1b6d (diff)
downloaddefinitions-746001dae609d752c1b32f511cbc632e8810f231.tar.gz
Rawdisk partitioning v2: Add pyfdisk.py module
This module is a GPL2 licenced python wrapper for fdisk, used to provide partitioning capabilities Change-Id: I5a28310f081045dd65d37bd5da4356ff3e6410f9
-rw-r--r--extensions/pyfdisk.py623
1 files changed, 623 insertions, 0 deletions
diff --git a/extensions/pyfdisk.py b/extensions/pyfdisk.py
new file mode 100644
index 00000000..3fbfe91c
--- /dev/null
+++ b/extensions/pyfdisk.py
@@ -0,0 +1,623 @@
+#!/usr/bin/python
+# Copyright (C) 2012-2015 Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program. If not, see <http://www.gnu.org/licenses/>.
+
+"""
+A simple Python wrapper for command line fdisk.
+
+It is intended to work on Linux, though may work on other operating
+systems using fdisk from util-linux.
+
+"""
+
+import contextlib
+from copy import deepcopy
+import re
+import subprocess
+import time
+import yaml
+
+
+class Extent(object):
+ """
+ A class to hold start and end points for other objects
+
+ Start and end points are measured in sectors. This class transparently
+ handles the inclusive nature of the start and end sectors of blocks of
+ storage. It also allows extents to be aligned within other extents.
+ """
+
+ def __init__(self, start=0, length=0, end=0):
+ if length and not start:
+ raise PartitioningError('Extent requires a non-zero start '
+ 'point and length')
+ if start and length:
+ self.start = int(start)
+ self.end = int(start) + int(length) - 1
+ else:
+ self.start = int(start)
+ self.end = int(end)
+
+ self.filled_sectors = 0
+
+ def __max__(self):
+ return self.end
+
+ def __min__(self):
+ return self.start
+
+ def __len__(self):
+ return self.end - self.start + 1
+
+ def __add__(self, other):
+ """Return the sum of two extents"""
+ return Extent(start=self.start,
+ length=(len(self) + len(other)))
+
+ def __iadd__(self, other):
+ """+="""
+ self.end += len(other)
+ return self
+
+ def __gt__(self, other):
+ return len(self) > len(other)
+
+ def __lt__(self, other):
+ return not self > other
+
+ def __str__(self):
+ return ('<Extent: Start=%d, End=%d, Length=%d>' %
+ (self.start, self.end, len(self)))
+
+ def pack(self, other):
+ """
+ Return a new Extent aligned to self's first unused sector
+
+ This is done by length, to quantify fitting an area of disk space
+ inside the other. The filled space in self is calculated and updated.
+
+ Returns:
+ A new Extent, starting at the first available sector in `self`,
+ with the same length as `other`.
+ """
+ length_other = len(other)
+ first_free_sector = self.start + self.filled_sectors
+ if length_other + self.filled_sectors > len(self):
+ raise PartitioningError('Not enough free space to pack Extent')
+ self.filled_sectors += length_other
+ return Extent(start=first_free_sector, length=length_other)
+
+ def free_sectors(self):
+ return len(self) - self.filled_sectors
+
+
+class PartitionList(object):
+ """
+ An iterable object to contain and process a list of Partition objects
+
+ This class eases the calculation of partition sizes and numbering, since
+ the properties of a given partition depend on each of the other partitions
+ in the list.
+
+ Attributes:
+ device: A Device class containing the partition list
+ """
+
+ def __init__(self, device):
+ """
+ Initialisation function
+
+ Args:
+ device: A Device object
+ """
+ self.device = device
+ self.extent = device.extent
+
+ self.__cached_list_hash = 0
+
+ self.__partition_list = []
+ self.__iter_index = 0
+
+ def append(self, partition):
+ """Append a new Partition object to the list"""
+ if isinstance(partition, Partition):
+ for part in self.__partition_list:
+ dup_attrib = part.compare(partition)
+ if dup_attrib:
+ raise PartitioningError('Duplicated partition attribute '
+ '\'%s\'' % dup_attrib)
+ self.__partition_list.append(partition)
+ else:
+ raise PartitioningError('PartitionList can only '
+ 'contain Partition objects')
+
+ def __iter__(self):
+ """Return self as an iterable object"""
+ self.__iter_index = 0
+ return self
+
+ def __next__(self):
+ """Return the next item in an iteration"""
+ if self.__iter_index == len(self.__partition_list):
+ raise StopIteration
+ else:
+ partition = self[self.__iter_index]
+ self.__iter_index += 1
+ return partition
+
+ def next(self):
+ """Provide a next() method for Python 2 compatibility"""
+ return self.__next__()
+
+ def __getitem__(self, i):
+ """Return an partition from the list, sorted by partition number"""
+ part_list = sorted(self.__update_partition_list(),
+ key=lambda part: part.number)
+ return part_list[i]
+
+ def free_sectors(self):
+ """Calculate the amount of unused space in the list"""
+ part_list = self.__update_partition_list()
+ self.extent.filled_sectors = 0
+ for part in part_list:
+ self.extent.pack(part.extent)
+ return self.extent.free_sectors()
+
+ def __update_partition_list(self):
+ """
+ Allocate extent and numbering for each Partition object in the list
+
+ A copy of the partition list is made so that any Partition object
+ returned from this list is a copy of a stored Partition object, thus
+ any partitions stored in the partition list remain intact even if a
+ copy is modified after is is returned. Hashing is used to avoid
+ updating the list when the partition list has not changed.
+ """
+ current_list_hash = hash(str(self.__partition_list))
+ if current_list_hash == self.__cached_list_hash:
+ return self.__cached_list
+
+ self.__cached_list_hash = current_list_hash
+
+ part_list = deepcopy(self.__partition_list)
+ used_numbers = set()
+ fill_partitions = set(partition for partition in part_list
+ if partition.size == 'fill')
+ requested_numbers = set(partition.number for partition in part_list
+ if hasattr(partition, 'number'))
+
+ # Get free space and the size of 'fill' partitions
+ self.extent.filled_sectors = 0
+ for part in part_list:
+ if part.size != 'fill':
+ extent = Extent(start=1,
+ length=self.get_length_sectors(part.size))
+ part.extent = extent
+ self.extent.pack(extent)
+
+ # Allocate aligned Extents and process partition numbers
+ if len(fill_partitions):
+ fill_size = self.extent.free_sectors() / len(fill_partitions)
+ # Set size of fill partitions
+ for part in fill_partitions:
+ part.size = fill_size * self.device.sector_size
+ part.extent = Extent(start=1, length=fill_size)
+
+ self.extent.filled_sectors = 0
+ for part in part_list:
+ part.extent = self.extent.pack(part.extent)
+
+ # Find the next unused partition number if not assigned
+ if hasattr(part, 'number'):
+ num = part.number
+ else:
+ for n in range(1, self.device.max_allowed_partitions + 1):
+ if n not in used_numbers and n not in requested_numbers:
+ num = n
+ break
+
+ part.number = num
+ used_numbers.add(num)
+
+ self.__cached_list = part_list
+ return part_list
+
+ def get_length_sectors(self, size_bytes):
+ """Get a length in sectors, aligned to 4096 byte boundaries"""
+ return (int(size_bytes) / self.device.sector_size +
+ ((int(size_bytes) % 4096) != 0) *
+ (4096 / self.device.sector_size))
+
+ def __str__(self):
+ string = ''
+ for part in self:
+ string = '%s\n%s\n' % (part, string)
+ return string.rstrip()
+
+ def __len__(self):
+ return len(self.__partition_list)
+
+ def __setitem__(self, i, value):
+ """ Update the ith item in the list """
+ self.append(partition)
+
+
+class Partition(object):
+ """
+ A class to describe a partition in a disk or image
+
+ The required attributes are loaded via kwargs.
+
+ Required attributes:
+ size: String describing the size of the partition in bytes
+ This may also be 'fill' to indicate that this partition should
+ be expanded to fill all unused space. Where there is more than
+ one fill partition, unused space is divided equally between the
+ fill partitions.
+ fdisk_type: A number describing the hexadecimal code used by fdisk
+ to describe the partition type. Any partitions with
+ fdisk_type='none' create an area of unused space.
+
+ Optional attributes:
+ **kwargs: A mapping of any keyword arguments
+ filesystem: A string describing the filesystem format for the
+ partition, or 'none' to skip filesystem creation.
+ description: A string describing the partition, for documentation
+ boot: Boolean string describing whether to set the bootable flag
+ mountpoint: String describing the mountpoint for the partition
+ number: Number used to override partition numbering for the
+ partition (Possible only when using an MBR partition table)
+ """
+ def __init__(self, size=0, fdisk_type=0x81, filesystem='none', **kwargs):
+ if not size and 'size' not in kwargs:
+ raise PartitioningError('Partition must have a non-zero size')
+ self.filesystem = filesystem
+ self.fdisk_type = fdisk_type
+ self.size = decode_human_size(size)
+ self.__dict__.update(**kwargs)
+
+ def compare(self, other):
+ """Check for mutually exclusive attributes"""
+ non_duplicable = ('number', 'mountpoint')
+ for attrib in non_duplicable:
+ if hasattr(self, attrib) and hasattr(other, attrib):
+ if getattr(self, attrib) == getattr(other, attrib):
+ return attrib
+ return False
+
+ def __str__(self):
+ string = ('Partition\n'
+ ' size: %s\n'
+ ' fdisk type: %s\n'
+ ' filesystem: %s'
+ % (self.size, hex(self.fdisk_type), self.filesystem))
+ if hasattr(self, 'extent'):
+ string += (
+ '\n start: %s'
+ '\n end: %s'
+ % (self.extent.start, self.extent.end))
+ if hasattr(self, 'number'):
+ string += '\n number: %s' % self.number
+ if hasattr(self, 'mountpoint'):
+ string += '\n mountpoint: %s' % self.mountpoint
+ if hasattr(self, 'boot'):
+ string += '\n bootable: %s' % self.boot
+
+ return string
+
+
+class Device(object):
+ """
+ A class to describe a disk or image, and its partition layout
+
+ Attributes are loaded from **kwargs, containing key-value pairs describing
+ the required attributes. This can be loaded from a YAML file, using the
+ module function load_yaml().
+
+ Required attributes:
+ location: The location of the device or disk image
+ size: A size in bytes describing the total amount of space the
+ partition table on the device will occupy, or 'fill' to
+ automatically fill the available space.
+
+ Optional attributes:
+ **kwargs: A mapping of any keyword arguments
+ start_offset: The first 512 byte sector of the first partition
+ (default: 2048)
+ partition_table_format: A string describing the type of partition
+ table used on the device (default: 'gpt')
+ partitions: A list of mappings for the attributes for each Partition
+ object. update_partitions() populates the partition list
+ based on the contents of this attribute.
+ """
+ min_start_bytes = 1024**2
+
+ def __init__(self, location, size, **kwargs):
+
+ if 'partition_table_format' not in kwargs:
+ self.partition_table_format = 'gpt'
+ if 'start_offset' not in kwargs:
+ self.start_offset = 2048
+
+ target_size = get_disk_size(location)
+ if str(size).lower() == 'fill':
+ self.size = target_size
+ else:
+ self.size = decode_human_size(size)
+
+ if self.size > target_size:
+ raise PartitioningError('Not enough space available on target')
+
+ if self.size <= self.min_start_bytes:
+ raise PartitioningError('Device size must be greater than %d '
+ 'bytes' % self.min_start_bytes)
+
+ # Get sector size
+ self.sector_size = get_sector_size(location)
+ self.location = location
+
+ # Populate Device attributes from keyword args
+ self.__dict__.update(**kwargs)
+
+ if self.partition_table_format.lower() == 'gpt':
+ self.max_allowed_partitions = 128
+ else:
+ self.max_allowed_partitions = 4
+
+ # Process Device size
+ start = (self.start_offset * 512) / self.sector_size
+ # Sector quantities in the specification are assumed to be 512 bytes
+ # This converts to the real sector size
+ if (start * self.sector_size) < self.min_start_bytes:
+ raise PartitioningError('Start offset should be greater than '
+ '%d, for %d byte sectors' %
+ (min_start_bytes / self.sector_size,
+ self.sector_size))
+ # Check the disk's first partition starts on a 4096 byte boundary
+ # this ensures alignment, and avoiding a reduction in performance
+ # on disks which use a 4096 byte physical sector size
+ if (start * self.sector_size) % 4096 != 0:
+ print('WARNING: Start sector is not aligned '
+ 'to 4096 byte sector boundaries')
+
+ # End sector is one sector less than the disk length
+ disk_end_sector = (self.size / self.sector_size) - 1
+ if self.partition_table_format == 'gpt':
+ # GPT partition table is duplicated at the end of the device.
+ # GPT header takes one sector, whatever the sector size,
+ # with a 16384 byte 'minimum' area for partition entries,
+ # supporting up to 128 partitions (128 bytes per entry).
+ # The duplicate GPT does not include the 'protective' MBR
+ gpt_size = ((16 * 1024) / self.sector_size) + 1
+ self.extent = Extent(start=start,
+ end=(disk_end_sector - gpt_size))
+ else:
+ self.extent = Extent(start=start, end=disk_end_sector)
+
+ self.update_partitions()
+
+ def update_partitions(self, partitions=None):
+ """
+ Reset list, populate with partitions from a list of attributes
+
+ Args:
+ partitions: A list of partition keyword attributes
+ """
+ self.partitionlist = PartitionList(self)
+ if partitions:
+ self.partitions = partitions
+ if hasattr(self, 'partitions'):
+ for partition_args in self.partitions:
+ self.add_partition(Partition(**partition_args))
+
+ def add_partition(self, partition):
+ """
+ Add a Partition object to the device's list of partitions
+
+ Args:
+ partition: a Partition class
+ """
+ if len(self.partitionlist) < self.max_allowed_partitions:
+ self.partitionlist.append(partition)
+ else:
+ raise PartitioningError('Exceeded maximum number of partitions '
+ 'for %s partition table (%d)' %
+ (self.partition_table_format.upper(),
+ self.max_allowed_partitions))
+
+ def commit(self):
+ """Write the partition table to the disk or image"""
+ pt_format = self.partition_table_format.lower()
+ print("Creating %s partition table on %s" %
+ (pt_format.upper(), self.location))
+
+ # Create a new partition table
+ if pt_format in ('mbr', 'dos'):
+ cmd = "o\n"
+ elif pt_format == 'gpt':
+ cmd = "g\n"
+ else:
+ raise PartitioningError('Unrecognised partition '
+ 'table type \'%s\'' % pt_format)
+
+ for partition in self.partitionlist:
+ # Create partitions
+ if str(partition.fdisk_type).lower() != 'none':
+ cmd += "n\n"
+ if pt_format in ('mbr', 'dos'):
+ cmd += "p\n"
+ cmd += (str(partition.number) + "\n"
+ "" + str(partition.extent.start) + "\n"
+ "" + str(partition.extent.end) + "\n")
+
+ # Set partition types
+ cmd += "t\n"
+ if partition.number > 1:
+ # fdisk does not ask for a partition
+ # number when setting the type of the
+ # first created partition
+ cmd += str(partition.number) + "\n"
+ cmd += str(partition.fdisk_type) + "\n"
+
+ # Set bootable flag
+ if hasattr(partition, 'boot') and pt_format == 'mbr':
+ if str(partition.boot).lower() in ('yes', 'true'):
+ cmd += "a\n"
+ if partition.number > 1:
+ cmd += str(partition.number) + "\n"
+
+ # Write changes
+ cmd += ("w\n"
+ "q\n")
+ p = subprocess.Popen(["fdisk", self.location],
+ stdin=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ stdout=subprocess.PIPE)
+ output = p.communicate(cmd)
+
+ errors = output[1].split('\n')[1:-1]
+ if errors:
+ # Note that the message saying 'disk does not contain a valid
+ # partition table' is never an error, it's a status message
+ # printed to stderr when fdisk starts with a blank device.
+ # Exception handling is done like this since fdisk will not
+ # return a failure exit code if it finds problems with the input
+ raise FdiskError('"%s"' % ' '.join(str(x) for x in errors))
+
+ def create_filesystems(self, skip=None):
+ """Create filesystems on the disk or image
+
+ Args:
+ skipping: A list of strings denoting partitions to skip filesystem
+ creation on, for example if custom settings are required
+ """
+ for part in self.partitionlist:
+ if hasattr(part, 'mountpoint') and part.mountpoint in skip:
+ continue
+ if part.filesystem.lower() != 'none':
+ with create_loopback(self.location,
+ part.extent.start * self.sector_size,
+ part.size) as loop:
+ print ('Creating %s filesystem on partition %s' %
+ (part.filesystem, part.number))
+ subprocess.check_output(['mkfs.' + part.filesystem, loop])
+
+ def __str__(self):
+ return ('<Device: location=%s, size=%s, partitions: %s>' %
+ (self.location, self.size, len(self.partitionlist)))
+
+
+class PartitioningError(Exception):
+
+ def __init__(self, msg=None):
+ self.msg = msg
+
+ def __str__(self):
+ return self.msg
+
+
+class FdiskError(Exception):
+
+ def __init__(self, msg=None):
+ self.msg = msg
+
+ def __str__(self):
+ return self.msg
+
+
+def load_yaml(location, size, yaml_file):
+ """
+ Load partition data from a yaml specification
+
+ The YAML file describes the attributes documented in the Device
+ and Partition classes.
+
+ Args:
+ yaml_file: String path to a YAML file to load
+ location: Path to the device node or image to use for partitioning
+ size: The desired device size in bytes (may be 'fill' to occupy the
+ entire device
+
+ Returns:
+ A Device object
+ """
+ with open(yaml_file, 'r') as f:
+ kwargs = yaml.safe_load(f)
+ return Device(location, size, **kwargs)
+
+
+def get_sector_size(location):
+ """Get the logical sector size of a block device or image, in bytes"""
+ return int(__filter_fdisk_list_output('Sector size.*?(\d+) bytes',
+ location)[0])
+
+def get_disk_size(location):
+ """Get the total size of a block device or image, in bytes"""
+ return int(__filter_fdisk_list_output('Disk.*?(\d+) bytes',
+ location)[0])
+
+def get_disk_offsets(location):
+ """Return an array of the partition start sectors in a device or image"""
+ return map(int, __filter_fdisk_list_output('%s\d+[\s*]+(\d+)' % location,
+ location))
+
+def __filter_fdisk_list_output(regex, location):
+ r = re.compile(regex, re.DOTALL)
+ m = re.findall(r, subprocess.check_output(['fdisk', '-l', location]))
+ if m:
+ return m
+ else:
+ raise PartitioningError('Error reading information from fdisk')
+
+@contextlib.contextmanager
+def create_loopback(mount_path, offset=0, size=0):
+ """
+ Create a loopback device for accessing partitions in block devices
+
+ Args:
+ mount_path: String path to mount
+ offset: Offset of the start of a partition in bytes (default 0)
+ size: Limits the size of the partition, in bytes (default 0)
+ Returns:
+ The path to a created loopback device node
+ """
+
+ try:
+ base_args = ['losetup', '--show', '-f', '-P', '-o', str(offset)]
+ if size and offset:
+ cmd = base_args + ['--sizelimit', str(size), mount_path]
+ else:
+ cmd = base_args + [mount_path]
+ loop_device = subprocess.check_output(cmd).rstrip()
+ # Allow the system time to see the new device
+ # On some systems, mounts created on the loopdev
+ # too soon after creating the loopback device
+ # may be unreliable, even though the -P option
+ # (--partscan) is passed to losetup
+ time.sleep(1)
+ except subprocess.CalledProcessError:
+ PartitioningError('Error creating loopback')
+ try:
+ yield loop_device
+ finally:
+ subprocess.check_call(['losetup', '-d', loop_device])
+
+def decode_human_size(size_string):
+ """Parse strings for human readable size factors"""
+
+ facts_of_1024 = ['', 'k', 'm', 'g', 't']
+ m = re.match('^(\d+)([kmgtKMGT]?)$', str(size_string))
+ if not m:
+ return size_string
+ return int(m.group(1)) * (1024 ** facts_of_1024.index(m.group(2).lower()))