X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/567891a691b1e753ea5dbee6c9c43ee05da40f66..29fd973e0da4b1bf61f76a7809b67b18ef9db195:/image_creator/disk.py diff --git a/image_creator/disk.py b/image_creator/disk.py index 8a06f78..b4311f8 100644 --- a/image_creator/disk.py +++ b/image_creator/disk.py @@ -32,20 +32,17 @@ # or implied, of GRNET S.A. from image_creator.util import get_command +from image_creator.util import try_fail_repeat +from image_creator.util import free_space from image_creator.util import FatalError -from image_creator.gpt import GPTPartitionTable from image_creator.bundle_volume import BundleVolume +from image_creator.image import Image import stat import os import tempfile import uuid -import re -import sys -import guestfs -import time -from sendfile import sendfile - +import shutil dd = get_command('dd') dmsetup = get_command('dmsetup') @@ -53,6 +50,9 @@ losetup = get_command('losetup') blockdev = get_command('blockdev') +TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt'] + + class Disk(object): """This class represents a hard disk hosting an Operating System @@ -61,28 +61,57 @@ class Disk(object): the Linux kernel. """ - def __init__(self, source, output): + def __init__(self, source, output, tmp=None): """Create a new Disk instance out of a source media. The source - media can be an image file, a block device or a directory.""" + media can be an image file, a block device or a directory. + """ self._cleanup_jobs = [] - self._devices = [] + self._images = [] self.source = source self.out = output self.meta = {} + self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.', + dir=self._get_tmp_dir(tmp)) + + self._add_cleanup(shutil.rmtree, self.tmp) + + def _get_tmp_dir(self, default=None): + """Check tmp directory candidates and return the one with the most + available space. + """ + if default is not None: + return default + + space = map(free_space, TMP_CANDIDATES) + + max_idx = 0 + max_val = space[0] + for i, val in zip(range(len(space)), space): + if val > max_val: + max_val = val + max_idx = i + + # Return the candidate path with more available space + return TMP_CANDIDATES[max_idx] def _add_cleanup(self, job, *args): + """Add a new job in the cleanup list""" self._cleanup_jobs.append((job, args)) def _losetup(self, fname): + """Setup a loop device and add it to the cleanup list. The loop device + will be detached when cleanup is called. + """ loop = losetup('-f', '--show', fname) loop = loop.strip() # remove the new-line char - self._add_cleanup(losetup, '-d', loop) + self._add_cleanup(try_fail_repeat, losetup, '-d', loop) return loop def _dir_to_disk(self): + """Create a disk out of a directory""" if self.source == '/': bundle = BundleVolume(self.out, self.meta) - image = '/var/tmp/%s.diskdump' % uuid.uuid4().hex + image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex) def check_unlink(path): if os.path.exists(path): @@ -98,9 +127,9 @@ class Disk(object): program ends. """ try: - while len(self._devices): - device = self._devices.pop() - device.destroy() + while len(self._images): + image = self._images.pop() + image.destroy() finally: # Make sure those are executed even if one of the device.destroy # methods throws exeptions. @@ -113,7 +142,7 @@ class Disk(object): instance. """ - self.out.output("Examining source media `%s'..." % self.source, False) + self.out.output("Examining source media `%s' ..." % self.source, False) sourcedev = self.source mode = os.stat(self.source).st_mode if stat.S_ISDIR(mode): @@ -123,7 +152,7 @@ class Disk(object): self.out.success('looks like an image file') sourcedev = self._losetup(self.source) elif not stat.S_ISBLK(mode): - raise ValueError("Invalid media source. Only block devices, " + raise FatalError("Invalid media source. Only block devices, " "regular files and directories are supported.") else: self.out.success('looks like a block device') @@ -131,7 +160,7 @@ class Disk(object): # Take a snapshot and return it to the user self.out.output("Snapshotting media source...", False) size = blockdev('--getsz', sourcedev) - cowfd, cow = tempfile.mkstemp() + cowfd, cow = tempfile.mkstemp(dir=self.tmp) os.close(cowfd) self._add_cleanup(os.unlink, cow) # Create cow sparse file @@ -144,326 +173,25 @@ class Disk(object): os.write(tablefd, "0 %d snapshot %s %s n 8" % (int(size), sourcedev, cowdev)) dmsetup('create', snapshot, table) - self._add_cleanup(dmsetup, 'remove', snapshot) - # Sometimes dmsetup remove fails with Device or resource busy, - # although everything is cleaned up and the snapshot is not - # used by anyone. Add a 2 seconds delay to be on the safe side. - self._add_cleanup(time.sleep, 2) + self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot) finally: os.unlink(table) self.out.success('done') return "/dev/mapper/%s" % snapshot - def get_device(self, media): - """Returns a newly created DiskDevice instance.""" + def get_image(self, media): + """Returns a newly created Image instance.""" - new_device = DiskDevice(media, self.out) - self._devices.append(new_device) - new_device.enable() - return new_device + image = Image(media, self.out) + self._images.append(image) + image.enable() + return image - def destroy_device(self, device): - """Destroys a DiskDevice instance previously created by - get_device method. + def destroy_image(self, image): + """Destroys an Image instance previously created by get_image method. """ - self._devices.remove(device) - device.destroy() - - -class DiskDevice(object): - """This class represents a block device hosting an Operating System - as created by the device-mapper. - """ - - def __init__(self, device, output, bootable=True, meta={}): - """Create a new DiskDevice.""" - - self.real_device = device - self.out = output - self.bootable = bootable - self.meta = meta - self.progress_bar = None - self.guestfs_device = None - self.size = 0 - - self.g = guestfs.GuestFS() - self.g.add_drive_opts(self.real_device, readonly=0) - - # Before version 1.17.14 the recovery process, which is a fork of the - # original process that called libguestfs, did not close its inherited - # file descriptors. This can cause problems especially if the parent - # process has opened pipes. Since the recovery process is an optional - # feature of libguestfs, it's better to disable it. - self.g.set_recovery_proc(0) - version = self.g.version() - if version['major'] > 1 or \ - (version['major'] == 1 and (version['minor'] >= 18 or - (version['minor'] == 17 and - version['release'] >= 14))): - self.g.set_recovery_proc(1) - self.out.output("Enabling recovery proc") - - #self.g.set_trace(1) - #self.g.set_verbose(1) - - self.guestfs_enabled = False - - def enable(self): - """Enable a newly created DiskDevice""" - - self.out.output('Launching helper VM (may take a while) ...', False) - # self.progressbar = self.out.Progress(100, "Launching helper VM", - # "percent") - # eh = self.g.set_event_callback(self.progress_callback, - # guestfs.EVENT_PROGRESS) - self.g.launch() - self.guestfs_enabled = True - # self.g.delete_event_callback(eh) - # self.progressbar.success('done') - # self.progressbar = None - self.out.success('done') - - self.out.output('Inspecting Operating System ...', False) - roots = self.g.inspect_os() - if len(roots) == 0: - raise FatalError("No operating system found") - if len(roots) > 1: - raise FatalError("Multiple operating systems found." - "We only support images with one OS.") - self.root = roots[0] - self.guestfs_device = self.g.part_to_dev(self.root) - self.size = self.g.blockdev_getsize64(self.guestfs_device) - self.meta['PARTITION_TABLE'] = \ - self.g.part_get_parttype(self.guestfs_device) - - self.ostype = self.g.inspect_get_type(self.root) - self.distro = self.g.inspect_get_distro(self.root) - self.out.success('found a(n) %s system' % self.distro) - - def destroy(self): - """Destroy this DiskDevice instance.""" - - # In new guestfs versions, there is a handy shutdown method for this - try: - if self.guestfs_enabled: - self.g.umount_all() - self.g.sync() - finally: - # Close the guestfs handler if open - self.g.close() - -# def progress_callback(self, ev, eh, buf, array): -# position = array[2] -# total = array[3] -# -# self.progressbar.goto((position * 100) // total) - - def mount(self, readonly=False): - """Mount all disk partitions in a correct order.""" - - mount = self.g.mount_ro if readonly else self.g.mount - msg = " read-only" if readonly else "" - self.out.output("Mounting the media%s..." % msg, False) - mps = self.g.inspect_get_mountpoints(self.root) - - # Sort the keys to mount the fs in a correct order. - # / should be mounted befor /boot, etc - def compare(a, b): - if len(a[0]) > len(b[0]): - return 1 - elif len(a[0]) == len(b[0]): - return 0 - else: - return -1 - mps.sort(compare) - for mp, dev in mps: - try: - mount(dev, mp) - except RuntimeError as msg: - self.out.warn("%s (ignored)" % msg) - self.out.success("done") - - def umount(self): - """Umount all mounted filesystems.""" - self.g.umount_all() - - def _last_partition(self): - if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt': - msg = "Unsupported partition table: %s. Only msdos and gpt " \ - "partition tables are supported" % self.meta['PARTITION_TABLE'] - raise FatalError(msg) - - is_extended = lambda p: \ - self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5 - is_logical = lambda p: \ - self.meta['PARTITION_TABLE'] != 'msdos' and p['part_num'] > 4 - - partitions = self.g.part_list(self.guestfs_device) - last_partition = partitions[-1] - - if is_logical(last_partition): - # The disk contains extended and logical partitions.... - extended = [p for p in partitions if is_extended(p)][0] - last_primary = [p for p in partitions if p['part_num'] <= 4][-1] - - # check if extended is the last primary partition - if last_primary['part_num'] > extended['part_num']: - last_partition = last_primary - - return last_partition - - def shrink(self): - """Shrink the disk. - - This is accomplished by shrinking the last filesystem in the - disk and then updating the partition table. The new disk size - (in bytes) is returned. - - ATTENTION: make sure unmount is called before shrink - """ - get_fstype = lambda p: \ - self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num'])) - is_logical = lambda p: \ - self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4 - is_extended = lambda p: \ - self.meta['PARTITION_TABLE'] == 'msdos' and \ - self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5 - - part_add = lambda ptype, start, stop: \ - self.g.part_add(self.guestfs_device, ptype, start, stop) - part_del = lambda p: self.g.part_del(self.guestfs_device, p) - part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p) - part_set_id = lambda p, id: \ - self.g.part_set_mbr_id(self.guestfs_device, p, id) - part_get_bootable = lambda p: \ - self.g.part_get_bootable(self.guestfs_device, p) - part_set_bootable = lambda p, bootable: \ - self.g.part_set_bootable(self.guestfs_device, p, bootable) - - MB = 2 ** 20 - - self.out.output("Shrinking image (this may take a while)...", False) - - sector_size = self.g.blockdev_getss(self.guestfs_device) - - last_part = None - fstype = None - while True: - last_part = self._last_partition() - fstype = get_fstype(last_part) - - if fstype == 'swap': - self.meta['SWAP'] = "%d:%s" % \ - (last_part['part_num'], - (last_part['part_size'] + MB - 1) // MB) - part_del(last_part['part_num']) - continue - elif is_extended(last_part): - part_del(last_part['part_num']) - continue - - # Most disk manipulation programs leave 2048 sectors after the last - # partition - new_size = last_part['part_end'] + 1 + 2048 * sector_size - self.size = min(self.size, new_size) - break - - if not re.match("ext[234]", fstype): - self.out.warn("Don't know how to resize %s partitions." % fstype) - return self.size - - part_dev = "%s%d" % (self.guestfs_device, last_part['part_num']) - self.g.e2fsck_f(part_dev) - self.g.resize2fs_M(part_dev) - - out = self.g.tune2fs_l(part_dev) - block_size = int( - filter(lambda x: x[0] == 'Block size', out)[0][1]) - block_cnt = int( - filter(lambda x: x[0] == 'Block count', out)[0][1]) - - start = last_part['part_start'] / sector_size - end = start + (block_size * block_cnt) / sector_size - 1 - - if is_logical(last_part): - partitions = self.g.part_list(self.guestfs_device) - - logical = [] # logical partitions - for partition in partitions: - if partition['part_num'] < 4: - continue - logical.append({ - 'num': partition['part_num'], - 'start': partition['part_start'] / sector_size, - 'end': partition['part_end'] / sector_size, - 'id': part_get_(partition['part_num']), - 'bootable': part_get_bootable(partition['part_num']) - }) - - logical[-1]['end'] = end # new end after resize - - # Recreate the extended partition - extended = [p for p in partitions if self._is_extended(p)][0] - part_del(extended['part_num']) - part_add('e', extended['part_start'], end) - - # Create all the logical partitions back - for l in logical: - part_add('l', l['start'], l['end']) - part_set_id(l['num'], l['id']) - part_set_bootable(l['num'], l['bootable']) - else: - # Recreate the last partition - if self.meta['PARTITION_TABLE'] == 'msdos': - last_part['id'] = part_get_id(last_part['part_num']) - - last_part['bootable'] = part_get_bootable(last_part['part_num']) - part_del(last_part['part_num']) - part_add('p', start, end) - part_set_bootable(last_part['part_num'], last_part['bootable']) - - if self.meta['PARTITION_TABLE'] == 'msdos': - part_set_id(last_part['part_num'], last_part['id']) - - new_size = (end + 1) * sector_size - - assert (new_size <= self.size) - - if self.meta['PARTITION_TABLE'] == 'gpt': - ptable = GPTPartitionTable(self.real_device) - self.size = ptable.shrink(new_size, self.size) - else: - self.size = min(new_size + 2048 * sector_size, self.size) - - self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB)) - - return self.size - - def dump(self, outfile): - """Dumps the content of device into a file. - - This method will only dump the actual payload, found by reading the - partition table. Empty space in the end of the device will be ignored. - """ - MB = 2 ** 20 - blocksize = 4 * MB # 4MB - size = self.size - progr_size = (size + MB - 1) // MB # in MB - progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb') - - with open(self.real_device, 'r') as src: - with open(outfile, "w") as dst: - left = size - offset = 0 - progressbar.next() - while left > 0: - length = min(left, blocksize) - _, sent = sendfile(dst.fileno(), src.fileno(), offset, - length) - offset += sent - left -= sent - progressbar.goto((size - left) // MB) - progressbar.success('image file %s was successfully created' % outfile) + self._images.remove(image) + image.destroy() # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :