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