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