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 import FatalError
  36 from clint.textui import progress
  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         sourcedev = self.source
 103         mode = os.stat(self.source).st_mode
 104         if stat.S_ISDIR(mode):
 105             return self._losetup(self._dir_to_disk())
 106         elif stat.S_ISREG(mode):
 107             sourcedev = self._losetup(self.source)
 108         elif not stat.S_ISBLK(mode):
 109             raise ValueError("Value for self.source is invalid")
 110
 111         # Take a snapshot and return it to the user
 112         size = blockdev('--getsize', sourcedev)
 113         cowfd, cow = tempfile.mkstemp()
 114         os.close(cowfd)
 115         self._add_cleanup(os.unlink, cow)
 116         # Create 1G cow sparse file
 117         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', 'seek=%d' % (1024 * 1024))
 118         cowdev = self._losetup(cow)
 119
 120         snapshot = uuid.uuid4().hex
 121         tablefd, table = tempfile.mkstemp()
 122         try:
 123             os.write(tablefd, "0 %d snapshot %s %s n 8" % \
 124                                         (int(size), sourcedev, cowdev))
 125             dmsetup('create', snapshot, table)
 126             self._add_cleanup(dmsetup, 'remove', snapshot)
 127         finally:
 128             os.unlink(table)
 129         new_device = DiskDevice("/dev/mapper/%s" % snapshot)
 130         self._devices.append(new_device)
 131         new_device.enable()
 132         return new_device
 133
 134     def destroy_device(self, device):
 135         """Destroys a DiskDevice instance previously created by
 136         get_device method.
 137         """
 138         self._devices.remove(device)
 139         device.destroy()
 140
 141
 142 def progress_generator(label=''):
 143     position = 0
 144     for i in progress.bar(range(100), label):
 145         if i < position:
 146             continue
 147         position = yield
 148     yield  # suppress the StopIteration exception
 149
 150
 151 class DiskDevice(object):
 152     """This class represents a block device hosting an Operating System
 153     as created by the device-mapper.
 154     """
 155
 156     def __init__(self, device, bootable=True):
 157         """Create a new DiskDevice."""
 158
 159         self.device = device
 160         self.bootable = bootable
 161         self.progress_bar = None
 162
 163         self.g = guestfs.GuestFS()
 164         self.g.add_drive_opts(self.device, readonly=0)
 165
 166         #self.g.set_trace(1)
 167         #self.g.set_verbose(1)
 168
 169         self.guestfs_enabled = False
 170
 171     def enable(self):
 172         """Enable a newly created DiskDevice"""
 173
 174         self.progressbar = progress_generator("VM lauch: ")
 175         self.progressbar.next()
 176         eh = self.g.set_event_callback(self.progress_callback,
 177                                                         guestfs.EVENT_PROGRESS)
 178         self.g.launch()
 179         self.guestfs_enabled = True
 180         self.g.delete_event_callback(eh)
 181         if self.progressbar is not None:
 182             self.progressbar.send(100)
 183             self.progressbar = None
 184
 185         roots = self.g.inspect_os()
 186         if len(roots) == 0:
 187             raise FatalError("No operating system found")
 188         if len(roots) > 1:
 189             raise FatalError("Multiple operating systems found")
 190
 191         self.root = roots[0]
 192         self.ostype = self.g.inspect_get_type(self.root)
 193         self.distro = self.g.inspect_get_distro(self.root)
 194
 195     def destroy(self):
 196         """Destroy this DiskDevice instance."""
 197
 198         if self.guestfs_enabled:
 199             self.g.umount_all()
 200             self.g.sync()
 201
 202         # Close the guestfs handler if open
 203         self.g.close()
 204
 205     def progress_callback(self, ev, eh, buf, array):
 206         position = array[2]
 207         total = array[3]
 208
 209         self.progressbar.send((position * 100) // total)
 210
 211         if position == total:
 212             self.progressbar = None
 213
 214     def mount(self):
 215         """Mount all disk partitions in a correct order."""
 216         mps = self.g.inspect_get_mountpoints(self.root)
 217
 218         # Sort the keys to mount the fs in a correct order.
 219         # / should be mounted befor /boot, etc
 220         def compare(a, b):
 221             if len(a[0]) > len(b[0]):
 222                 return 1
 223             elif len(a[0]) == len(b[0]):
 224                 return 0
 225             else:
 226                 return -1
 227         mps.sort(compare)
 228         for mp, dev in mps:
 229             try:
 230                 self.g.mount(dev, mp)
 231             except RuntimeError as msg:
 232                 print "%s (ignored)" % msg
 233
 234     def umount(self):
 235         """Umount all mounted filesystems."""
 236         self.g.umount_all()
 237
 238     def shrink(self):
 239         """Shrink the disk.
 240
 241         This is accomplished by shrinking the last filesystem in the
 242         disk and then updating the partition table. The new disk size
 243         (in bytes) is returned.
 244         """
 245         dev = self.g.part_to_dev(self.root)
 246         parttype = self.g.part_get_parttype(dev)
 247         if parttype != 'msdos':
 248             raise FatalError("You have a %s partition table. "
 249                 "Only msdos partitions are supported" % parttype)
 250
 251         last_partition = self.g.part_list(dev)[-1]
 252
 253         if last_partition['part_num'] > 4:
 254             raise FatalError("This disk contains logical partitions. "
 255                 "Only primary partitions are supported.")
 256
 257         part_dev = "%s%d" % (dev, last_partition['part_num'])
 258         fs_type = self.g.vfs_type(part_dev)
 259         if not re.match("ext[234]", fs_type):
 260             print "Warning: Don't know how to resize %s partitions." % vfs_type
 261             return
 262
 263         self.g.e2fsck_f(part_dev)
 264         self.g.resize2fs_M(part_dev)
 265         output = self.g.tune2fs_l(part_dev)
 266         block_size = int(filter(lambda x: x[0] == 'Block size', output)[0][1])
 267         block_cnt = int(filter(lambda x: x[0] == 'Block count', output)[0][1])
 268
 269         sector_size = self.g.blockdev_getss(dev)
 270
 271         start = last_partition['part_start'] / sector_size
 272         end = start + (block_size * block_cnt) / sector_size - 1
 273
 274         self.g.part_del(dev, last_partition['part_num'])
 275         self.g.part_add(dev, 'p', start, end)
 276
 277         return (end + 1) * sector_size
 278
 279     def size(self):
 280         """Returns the "payload" size of the device.
 281
 282         The size returned by this method is the size of the space occupied by
 283         the partitions (including the space before the first partition).
 284         """
 285         dev = self.g.part_to_dev(self.root)
 286         last = self.g.part_list(dev)[-1]
 287
 288         return last['part_end']
 289
 290 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :