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
  35 from image_creator.util import warn, progress, success, output, FatalError
  36
  37 import stat
  38 import os
  39 import tempfile
  40 import uuid
  41 import re
  42 import sys
  43 import guestfs
  44 import time
  45 from sendfile import sendfile
  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         output("Examining source media `%s'..." % self.source, False)
 104         sourcedev = self.source
 105         mode = os.stat(self.source).st_mode
 106         if stat.S_ISDIR(mode):
 107             success('looks like a directory')
 108             return self._losetup(self._dir_to_disk())
 109         elif stat.S_ISREG(mode):
 110             success('looks like an image file')
 111             sourcedev = self._losetup(self.source)
 112         elif not stat.S_ISBLK(mode):
 113             raise ValueError("Invalid media source. Only block devices, "
 114                             "regular files and directories are supported.")
 115         else:
 116             success('looks like a block device')
 117
 118         # Take a snapshot and return it to the user
 119         output("Snapshotting media source...", False)
 120         size = blockdev('--getsize', sourcedev)
 121         cowfd, cow = tempfile.mkstemp()
 122         os.close(cowfd)
 123         self._add_cleanup(os.unlink, cow)
 124         # Create 1G cow sparse file
 125         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', \
 126                                         'seek=%d' % (1024 * 1024))
 127         cowdev = self._losetup(cow)
 128
 129         snapshot = uuid.uuid4().hex
 130         tablefd, table = tempfile.mkstemp()
 131         try:
 132             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
 133                                         (int(size), sourcedev, cowdev))
 134             dmsetup('create', snapshot, table)
 135             self._add_cleanup(dmsetup, 'remove', snapshot)
 136             # Sometimes dmsetup remove fails with Device or resource busy,
 137             # although everything is cleaned up and the snapshot is not
 138             # used by anyone. Add a 2 seconds delay to be on the safe side.
 139             self._add_cleanup(time.sleep, 2)
 140
 141         finally:
 142             os.unlink(table)
 143         success('done')
 144         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
 145         self._devices.append(new_device)
 146         new_device.enable()
 147         return new_device
 148
 149     def destroy_device(self, device):
 150         """Destroys a DiskDevice instance previously created by
 151         get_device method.
 152         """
 153         self._devices.remove(device)
 154         device.destroy()
 155
 156
 157 class DiskDevice(object):
 158     """This class represents a block device hosting an Operating System
 159     as created by the device-mapper.
 160     """
 161
 162     def __init__(self, device, bootable=True):
 163         """Create a new DiskDevice."""
 164
 165         self.device = device
 166         self.bootable = bootable
 167         self.progress_bar = None
 168
 169         self.g = guestfs.GuestFS()
 170         self.g.add_drive_opts(self.device, readonly=0)
 171
 172         #self.g.set_trace(1)
 173         #self.g.set_verbose(1)
 174
 175         self.guestfs_enabled = False
 176
 177     def enable(self):
 178         """Enable a newly created DiskDevice"""
 179         self.progressbar = progress("Launching helper VM: ")
 180         self.progressbar.next()
 181         eh = self.g.set_event_callback(self.progress_callback,
 182                                                     guestfs.EVENT_PROGRESS)
 183         self.g.launch()
 184         self.guestfs_enabled = True
 185         self.g.delete_event_callback(eh)
 186         if self.progressbar is not None:
 187             self.progressbar.send(100)
 188             self.progressbar = None
 189
 190         output('Inspecting Operating System...', False)
 191         roots = self.g.inspect_os()
 192         if len(roots) == 0:
 193             raise FatalError("No operating system found")
 194         if len(roots) > 1:
 195             raise FatalError("Multiple operating systems found."
 196                             "We only support images with one filesystem.")
 197         self.root = roots[0]
 198         self.ostype = self.g.inspect_get_type(self.root)
 199         self.distro = self.g.inspect_get_distro(self.root)
 200         success('found a %s system' % self.distro)
 201
 202     def destroy(self):
 203         """Destroy this DiskDevice instance."""
 204
 205         if self.guestfs_enabled:
 206             self.g.umount_all()
 207             self.g.sync()
 208
 209         # Close the guestfs handler if open
 210         self.g.close()
 211
 212     def progress_callback(self, ev, eh, buf, array):
 213         position = array[2]
 214         total = array[3]
 215
 216         self.progressbar.send((position * 100) // total)
 217
 218         if position == total:
 219             self.progressbar = None
 220
 221     def mount(self):
 222         """Mount all disk partitions in a correct order."""
 223
 224         output("Mounting image...", False)
 225         mps = self.g.inspect_get_mountpoints(self.root)
 226
 227         # Sort the keys to mount the fs in a correct order.
 228         # / should be mounted befor /boot, etc
 229         def compare(a, b):
 230             if len(a[0]) > len(b[0]):
 231                 return 1
 232             elif len(a[0]) == len(b[0]):
 233                 return 0
 234             else:
 235                 return -1
 236         mps.sort(compare)
 237         for mp, dev in mps:
 238             try:
 239                 self.g.mount(dev, mp)
 240             except RuntimeError as msg:
 241                 warn("%s (ignored)" % msg)
 242         success("done")
 243
 244     def umount(self):
 245         """Umount all mounted filesystems."""
 246         self.g.umount_all()
 247
 248     def shrink(self):
 249         """Shrink the disk.
 250
 251         This is accomplished by shrinking the last filesystem in the
 252         disk and then updating the partition table. The new disk size
 253         (in bytes) is returned.
 254         """
 255         output("Shrinking image (this may take a while)...", False)
 256
 257         dev = self.g.part_to_dev(self.root)
 258         parttype = self.g.part_get_parttype(dev)
 259         if parttype != 'msdos':
 260             raise FatalError("You have a %s partition table. "
 261                 "Only msdos partitions are supported" % parttype)
 262
 263         last_partition = self.g.part_list(dev)[-1]
 264
 265         if last_partition['part_num'] > 4:
 266             raise FatalError("This disk contains logical partitions. "
 267                 "Only primary partitions are supported.")
 268
 269         part_dev = "%s%d" % (dev, last_partition['part_num'])
 270         fs_type = self.g.vfs_type(part_dev)
 271         if not re.match("ext[234]", fs_type):
 272             warn("Don't know how to resize %s partitions." % vfs_type)
 273             return
 274
 275         self.g.e2fsck_f(part_dev)
 276         self.g.resize2fs_M(part_dev)
 277
 278         out = self.g.tune2fs_l(part_dev)
 279         block_size = int(
 280             filter(lambda x: x[0] == 'Block size', out)[0][1])
 281         block_cnt = int(
 282             filter(lambda x: x[0] == 'Block count', out)[0][1])
 283
 284         sector_size = self.g.blockdev_getss(dev)
 285
 286         start = last_partition['part_start'] / sector_size
 287         end = start + (block_size * block_cnt) / sector_size - 1
 288
 289         self.g.part_del(dev, last_partition['part_num'])
 290         self.g.part_add(dev, 'p', start, end)
 291
 292         new_size = (end + 1) * sector_size
 293         success("new image size is %dMB" %
 294                             ((new_size + 2 ** 20 - 1) // 2 ** 20))
 295         return new_size
 296
 297     def size(self):
 298         """Returns the "payload" size of the device.
 299
 300         The size returned by this method is the size of the space occupied by
 301         the partitions (including the space before the first partition).
 302         """
 303         dev = self.g.part_to_dev(self.root)
 304         last = self.g.part_list(dev)[-1]
 305
 306         return last['part_end'] + 1
 307
 308     def dump(self, outfile):
 309         """Dumps the content of device into a file.
 310
 311         This method will only dump the actual payload, found by reading the
 312         partition table. Empty space in the end of the device will be ignored.
 313         """
 314         blocksize = 2 ** 22  # 4MB
 315         size = self.size()
 316         progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
 317         progressbar = progress("Dumping image file: ", progress_size)
 318
 319         source = open(self.device, "r")
 320         try:
 321             dest = open(outfile, "w")
 322             try:
 323                 left = size
 324                 offset = 0
 325                 progressbar.next()
 326                 while left > 0:
 327                     length = min(left, blocksize)
 328                     sent = sendfile(dest.fileno(), source.fileno(), offset,
 329                                                                         length)
 330                     offset += sent
 331                     left -= sent
 332                     for i in range((length + 2 ** 20 - 1) // 2 ** 20):
 333                         progressbar.next()
 334             finally:
 335                 dest.close()
 336         finally:
 337             source.close()
 338
 339         success('Image file %s was successfully created' % outfile)
 340
 341 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :