code.grnet.gr Git - snf-image-creator/blob - image_creator/disk.py

   1 # Copyright 2012 GRNET S.A. All rights reserved.
   2 #
   3 # Redistribution and use in source and binary forms, with or
   4 # without modification, are permitted provided that the following
   5 # conditions are met:
   6 #
   7 #   1. Redistributions of source code must retain the above
   8 #      copyright notice, this list of conditions and the following
   9 #      disclaimer.
  10 #
  11 #   2. Redistributions in binary form must reproduce the above
  12 #      copyright notice, this list of conditions and the following
  13 #      disclaimer in the documentation and/or other materials
  14 #      provided with the distribution.
  15 #
  16 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  17 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  19 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  20 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  21 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  22 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  23 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  24 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  25 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  26 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  27 # POSSIBILITY OF SUCH DAMAGE.
  28 #
  29 # The views and conclusions contained in the software and
  30 # documentation are those of the authors and should not be
  31 # interpreted as representing official policies, either expressed
  32 # or implied, of GRNET S.A.
  33
  34 from image_creator.util import get_command, warn, progress_generator
  35 from image_creator import FatalError
  36 from clint.textui import indent, puts, colored
  37
  38 import stat
  39 import os
  40 import tempfile
  41 import uuid
  42 import re
  43 import sys
  44 import guestfs
  45 import time
  46
  47
  48 class DiskError(Exception):
  49     pass
  50
  51 dd = get_command('dd')
  52 dmsetup = get_command('dmsetup')
  53 losetup = get_command('losetup')
  54 blockdev = get_command('blockdev')
  55
  56
  57 class Disk(object):
  58     """This class represents a hard disk hosting an Operating System
  59
  60     A Disk instance never alters the source media it is created from.
  61     Any change is done on a snapshot created by the device-mapper of
  62     the Linux kernel.
  63     """
  64
  65     def __init__(self, source):
  66         """Create a new Disk instance out of a source media. The source
  67         media can be an image file, a block device or a directory."""
  68         self._cleanup_jobs = []
  69         self._devices = []
  70         self.source = source
  71
  72     def _add_cleanup(self, job, *args):
  73         self._cleanup_jobs.append((job, args))
  74
  75     def _losetup(self, fname):
  76         loop = losetup('-f', '--show', fname)
  77         loop = loop.strip()  # remove the new-line char
  78         self._add_cleanup(losetup, '-d', loop)
  79         return loop
  80
  81     def _dir_to_disk(self):
  82         raise NotImplementedError
  83
  84     def cleanup(self):
  85         """Cleanup internal data. This needs to be called before the
  86         program ends.
  87         """
  88         while len(self._devices):
  89             device = self._devices.pop()
  90             device.destroy()
  91
  92         while len(self._cleanup_jobs):
  93             job, args = self._cleanup_jobs.pop()
  94             job(*args)
  95
  96     def get_device(self):
  97         """Returns a newly created DiskDevice instance.
  98
  99         This instance is a snapshot of the original source media of
 100         the Disk instance.
 101         """
 102
 103         puts("Examining source media `%s'" % self.source)
 104         with indent(4):
 105             sourcedev = self.source
 106             mode = os.stat(self.source).st_mode
 107             if stat.S_ISDIR(mode):
 108                 puts(colored.green('Looks like a directory'))
 109                 return self._losetup(self._dir_to_disk())
 110             elif stat.S_ISREG(mode):
 111                 puts(colored.green('Looks like an image file'))
 112                 sourcedev = self._losetup(self.source)
 113             elif not stat.S_ISBLK(mode):
 114                 raise ValueError("Invalid media source. Only block devices, "
 115                                 "regular files and directories are supported.")
 116             else:
 117                 puts(colored.green('Looks like a block device'))
 118             #puts()
 119
 120         # Take a snapshot and return it to the user
 121         puts("Snapshotting media source")
 122         with indent(4):
 123             size = blockdev('--getsize', sourcedev)
 124             cowfd, cow = tempfile.mkstemp()
 125             os.close(cowfd)
 126             self._add_cleanup(os.unlink, cow)
 127             # Create 1G cow sparse file
 128             dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
 129                                             'seek=%d' % (1024 * 1024))
 130             cowdev = self._losetup(cow)
 131
 132             snapshot = uuid.uuid4().hex
 133             tablefd, table = tempfile.mkstemp()
 134             try:
 135                 os.write(tablefd, "0 %d snapshot %s %s n 8" % \
 136                                             (int(size), sourcedev, cowdev))
 137                 dmsetup('create', snapshot, table)
 138                 self._add_cleanup(dmsetup, 'remove', snapshot)
 139                 # Sometimes dmsetup remove fails with Device or resource busy,
 140                 # although everything is cleaned up and the snapshot is not
 141                 # used by anyone. Add a 2 seconds delay to be on the safe side.
 142                 self._add_cleanup(time.sleep, 2)
 143
 144             finally:
 145                 os.unlink(table)
 146             puts(colored.green('Done'))
 147         # puts()
 148         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
 149         self._devices.append(new_device)
 150         new_device.enable()
 151         return new_device
 152
 153     def destroy_device(self, device):
 154         """Destroys a DiskDevice instance previously created by
 155         get_device method.
 156         """
 157         self._devices.remove(device)
 158         device.destroy()
 159
 160
 161 class DiskDevice(object):
 162     """This class represents a block device hosting an Operating System
 163     as created by the device-mapper.
 164     """
 165
 166     def __init__(self, device, bootable=True):
 167         """Create a new DiskDevice."""
 168
 169         self.device = device
 170         self.bootable = bootable
 171         self.progress_bar = None
 172
 173         self.g = guestfs.GuestFS()
 174         self.g.add_drive_opts(self.device, readonly=0)
 175
 176         #self.g.set_trace(1)
 177         #self.g.set_verbose(1)
 178
 179         self.guestfs_enabled = False
 180
 181     def enable(self):
 182         """Enable a newly created DiskDevice"""
 183         self.progressbar = progress_generator("Launching helper VM: ")
 184         with indent(4):
 185             self.progressbar.next()
 186             eh = self.g.set_event_callback(self.progress_callback,
 187                                                         guestfs.EVENT_PROGRESS)
 188             self.g.launch()
 189             self.guestfs_enabled = True
 190             self.g.delete_event_callback(eh)
 191             if self.progressbar is not None:
 192                 self.progressbar.send(100)
 193                 self.progressbar = None
 194             puts(colored.green('Done'))
 195
 196         puts('Inspecting Operating System')
 197         with indent(4):
 198             roots = self.g.inspect_os()
 199             if len(roots) == 0:
 200                 raise FatalError("No operating system found")
 201             if len(roots) > 1:
 202                 raise FatalError("Multiple operating systems found."
 203                                 "We only support images with one filesystem.")
 204             self.root = roots[0]
 205             self.ostype = self.g.inspect_get_type(self.root)
 206             self.distro = self.g.inspect_get_distro(self.root)
 207             puts(colored.green('Found a %s system' % self.distro))
 208         puts()
 209
 210     def destroy(self):
 211         """Destroy this DiskDevice instance."""
 212
 213         if self.guestfs_enabled:
 214             self.g.umount_all()
 215             self.g.sync()
 216
 217         # Close the guestfs handler if open
 218         self.g.close()
 219
 220     def progress_callback(self, ev, eh, buf, array):
 221         position = array[2]
 222         total = array[3]
 223
 224         self.progressbar.send((position * 100) // total)
 225
 226         if position == total:
 227             self.progressbar = None
 228
 229     def mount(self):
 230         """Mount all disk partitions in a correct order."""
 231         mps = self.g.inspect_get_mountpoints(self.root)
 232
 233         # Sort the keys to mount the fs in a correct order.
 234         # / should be mounted befor /boot, etc
 235         def compare(a, b):
 236             if len(a[0]) > len(b[0]):
 237                 return 1
 238             elif len(a[0]) == len(b[0]):
 239                 return 0
 240             else:
 241                 return -1
 242         mps.sort(compare)
 243         for mp, dev in mps:
 244             try:
 245                 self.g.mount(dev, mp)
 246             except RuntimeError as msg:
 247                 print "%s (ignored)" % msg
 248
 249     def umount(self):
 250         """Umount all mounted filesystems."""
 251         self.g.umount_all()
 252
 253     def shrink(self):
 254         """Shrink the disk.
 255
 256         This is accomplished by shrinking the last filesystem in the
 257         disk and then updating the partition table. The new disk size
 258         (in bytes) is returned.
 259         """
 260         puts("Shrinking image (this may take a while)")
 261
 262         dev = self.g.part_to_dev(self.root)
 263         parttype = self.g.part_get_parttype(dev)
 264         if parttype != 'msdos':
 265             raise FatalError("You have a %s partition table. "
 266                 "Only msdos partitions are supported" % parttype)
 267
 268         last_partition = self.g.part_list(dev)[-1]
 269
 270         if last_partition['part_num'] > 4:
 271             raise FatalError("This disk contains logical partitions. "
 272                 "Only primary partitions are supported.")
 273
 274         part_dev = "%s%d" % (dev, last_partition['part_num'])
 275         fs_type = self.g.vfs_type(part_dev)
 276         if not re.match("ext[234]", fs_type):
 277             warn("Don't know how to resize %s partitions." % vfs_type)
 278             return
 279
 280         with indent(4):
 281             self.g.e2fsck_f(part_dev)
 282             self.g.resize2fs_M(part_dev)
 283
 284             output = self.g.tune2fs_l(part_dev)
 285             block_size = int(
 286                 filter(lambda x: x[0] == 'Block size', output)[0][1])
 287             block_cnt = int(
 288                 filter(lambda x: x[0] == 'Block count', output)[0][1])
 289
 290             sector_size = self.g.blockdev_getss(dev)
 291
 292             start = last_partition['part_start'] / sector_size
 293             end = start + (block_size * block_cnt) / sector_size - 1
 294
 295             self.g.part_del(dev, last_partition['part_num'])
 296             self.g.part_add(dev, 'p', start, end)
 297
 298             new_size = (end + 1) * sector_size
 299             puts("  New image size is %dMB\n" % (new_size // 2 ** 20))
 300         return new_size
 301
 302     def size(self):
 303         """Returns the "payload" size of the device.
 304
 305         The size returned by this method is the size of the space occupied by
 306         the partitions (including the space before the first partition).
 307         """
 308         dev = self.g.part_to_dev(self.root)
 309         last = self.g.part_list(dev)[-1]
 310
 311         return last['part_end']
 312
 313 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :