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, False)
 104         sourcedev = self.source
 105         mode = os.stat(self.source).st_mode
 106         if stat.S_ISDIR(mode):
 107             puts(colored.green('looks like a directory'))
 108             return self._losetup(self._dir_to_disk())
 109         elif stat.S_ISREG(mode):
 110             puts(colored.green('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             puts(colored.green('looks like a block device'))
 117
 118         # Take a snapshot and return it to the user
 119         puts("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         puts(colored.green('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_generator("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         puts('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         puts(colored.green('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         mps = self.g.inspect_get_mountpoints(self.root)
 224
 225         # Sort the keys to mount the fs in a correct order.
 226         # / should be mounted befor /boot, etc
 227         def compare(a, b):
 228             if len(a[0]) > len(b[0]):
 229                 return 1
 230             elif len(a[0]) == len(b[0]):
 231                 return 0
 232             else:
 233                 return -1
 234         mps.sort(compare)
 235         for mp, dev in mps:
 236             try:
 237                 self.g.mount(dev, mp)
 238             except RuntimeError as msg:
 239                 print "%s (ignored)" % msg
 240
 241     def umount(self):
 242         """Umount all mounted filesystems."""
 243         self.g.umount_all()
 244
 245     def shrink(self):
 246         """Shrink the disk.
 247
 248         This is accomplished by shrinking the last filesystem in the
 249         disk and then updating the partition table. The new disk size
 250         (in bytes) is returned.
 251         """
 252         puts("Shrinking image (this may take a while)...", False)
 253
 254         dev = self.g.part_to_dev(self.root)
 255         parttype = self.g.part_get_parttype(dev)
 256         if parttype != 'msdos':
 257             raise FatalError("You have a %s partition table. "
 258                 "Only msdos partitions are supported" % parttype)
 259
 260         last_partition = self.g.part_list(dev)[-1]
 261
 262         if last_partition['part_num'] > 4:
 263             raise FatalError("This disk contains logical partitions. "
 264                 "Only primary partitions are supported.")
 265
 266         part_dev = "%s%d" % (dev, last_partition['part_num'])
 267         fs_type = self.g.vfs_type(part_dev)
 268         if not re.match("ext[234]", fs_type):
 269             warn("Don't know how to resize %s partitions." % vfs_type)
 270             return
 271
 272         self.g.e2fsck_f(part_dev)
 273         self.g.resize2fs_M(part_dev)
 274
 275         output = self.g.tune2fs_l(part_dev)
 276         block_size = int(
 277             filter(lambda x: x[0] == 'Block size', output)[0][1])
 278         block_cnt = int(
 279             filter(lambda x: x[0] == 'Block count', output)[0][1])
 280
 281         sector_size = self.g.blockdev_getss(dev)
 282
 283         start = last_partition['part_start'] / sector_size
 284         end = start + (block_size * block_cnt) / sector_size - 1
 285
 286         self.g.part_del(dev, last_partition['part_num'])
 287         self.g.part_add(dev, 'p', start, end)
 288
 289         new_size = (end + 1) * sector_size
 290         puts(colored.green("new image size is %dMB\n" % (new_size // 2 ** 20)))
 291
 292         return new_size
 293
 294     def size(self):
 295         """Returns the "payload" size of the device.
 296
 297         The size returned by this method is the size of the space occupied by
 298         the partitions (including the space before the first partition).
 299         """
 300         dev = self.g.part_to_dev(self.root)
 301         last = self.g.part_list(dev)[-1]
 302
 303         return last['part_end']
 304
 305 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :