code.grnet.gr Git - snf-image-creator/blob - image_creator/disk.py

   1 # -*- coding: utf-8 -*-
   2 #
   3 # Copyright 2012 GRNET S.A. All rights reserved.
   4 #
   5 # Redistribution and use in source and binary forms, with or
   6 # without modification, are permitted provided that the following
   7 # conditions are met:
   8 #
   9 #   1. Redistributions of source code must retain the above
  10 #      copyright notice, this list of conditions and the following
  11 #      disclaimer.
  12 #
  13 #   2. Redistributions in binary form must reproduce the above
  14 #      copyright notice, this list of conditions and the following
  15 #      disclaimer in the documentation and/or other materials
  16 #      provided with the distribution.
  17 #
  18 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  19 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  20 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  21 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  22 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  25 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  26 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  27 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  28 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  29 # POSSIBILITY OF SUCH DAMAGE.
  30 #
  31 # The views and conclusions contained in the software and
  32 # documentation are those of the authors and should not be
  33 # interpreted as representing official policies, either expressed
  34 # or implied, of GRNET S.A.
  35
  36 """Module hosting the Disk class."""
  37
  38 from image_creator.util import get_command
  39 from image_creator.util import try_fail_repeat
  40 from image_creator.util import free_space
  41 from image_creator.util import FatalError
  42 from image_creator.bundle_volume import BundleVolume
  43 from image_creator.image import Image
  44
  45 import stat
  46 import os
  47 import tempfile
  48 import uuid
  49 import shutil
  50
  51 dd = get_command('dd')
  52 dmsetup = get_command('dmsetup')
  53 losetup = get_command('losetup')
  54 blockdev = get_command('blockdev')
  55
  56
  57 def get_tmp_dir(default=None):
  58     """Check tmp directory candidates and return the one with the most
  59     available space.
  60     """
  61     if default is not None:
  62         return default
  63
  64     TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
  65
  66     space = map(free_space, TMP_CANDIDATES)
  67
  68     max_idx = 0
  69     max_val = space[0]
  70     for i, val in zip(range(len(space)), space):
  71         if val > max_val:
  72             max_val = val
  73             max_idx = i
  74
  75     # Return the candidate path with more available space
  76     return TMP_CANDIDATES[max_idx]
  77
  78
  79 class Disk(object):
  80     """This class represents a hard disk hosting an Operating System
  81
  82     A Disk instance never alters the source media it is created from.
  83     Any change is done on a snapshot created by the device-mapper of
  84     the Linux kernel.
  85     """
  86
  87     def __init__(self, source, output, tmp=None):
  88         """Create a new Disk instance out of a source media. The source
  89         media can be an image file, a block device or a directory.
  90         """
  91         self._cleanup_jobs = []
  92         self._images = []
  93         self.source = source
  94         self.out = output
  95         self.meta = {}
  96         self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
  97                                     dir=get_tmp_dir(tmp))
  98
  99         self._add_cleanup(shutil.rmtree, self.tmp)
 100
 101     def _add_cleanup(self, job, *args):
 102         """Add a new job in the cleanup list"""
 103         self._cleanup_jobs.append((job, args))
 104
 105     def _losetup(self, fname):
 106         """Setup a loop device and add it to the cleanup list. The loop device
 107         will be detached when cleanup is called.
 108         """
 109         loop = losetup('-f', '--show', fname)
 110         loop = loop.strip()  # remove the new-line char
 111         self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
 112         return loop
 113
 114     def _dir_to_disk(self):
 115         """Create a disk out of a directory"""
 116         if self.source == '/':
 117             bundle = BundleVolume(self.out, self.meta)
 118             image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex)
 119
 120             def check_unlink(path):
 121                 if os.path.exists(path):
 122                     os.unlink(path)
 123
 124             self._add_cleanup(check_unlink, image)
 125             bundle.create_image(image)
 126             return self._losetup(image)
 127         raise FatalError("Using a directory as media source is supported")
 128
 129     def cleanup(self):
 130         """Cleanup internal data. This needs to be called before the
 131         program ends.
 132         """
 133         try:
 134             while len(self._images):
 135                 image = self._images.pop()
 136                 image.destroy()
 137         finally:
 138             # Make sure those are executed even if one of the device.destroy
 139             # methods throws exeptions.
 140             while len(self._cleanup_jobs):
 141                 job, args = self._cleanup_jobs.pop()
 142                 job(*args)
 143
 144     def snapshot(self):
 145         """Creates a snapshot of the original source media of the Disk
 146         instance.
 147         """
 148
 149         self.out.output("Examining source media `%s' ..." % self.source, False)
 150         sourcedev = self.source
 151         mode = os.stat(self.source).st_mode
 152         if stat.S_ISDIR(mode):
 153             self.out.success('looks like a directory')
 154             return self._dir_to_disk()
 155         elif stat.S_ISREG(mode):
 156             self.out.success('looks like an image file')
 157             sourcedev = self._losetup(self.source)
 158         elif not stat.S_ISBLK(mode):
 159             raise FatalError("Invalid media source. Only block devices, "
 160                              "regular files and directories are supported.")
 161         else:
 162             self.out.success('looks like a block device')
 163
 164         # Take a snapshot and return it to the user
 165         self.out.output("Snapshotting media source ...", False)
 166         size = blockdev('--getsz', sourcedev)
 167         cowfd, cow = tempfile.mkstemp(dir=self.tmp)
 168         os.close(cowfd)
 169         self._add_cleanup(os.unlink, cow)
 170         # Create cow sparse file
 171         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
 172         cowdev = self._losetup(cow)
 173
 174         snapshot = uuid.uuid4().hex
 175         tablefd, table = tempfile.mkstemp()
 176         try:
 177             os.write(tablefd, "0 %d snapshot %s %s n 8" %
 178                               (int(size), sourcedev, cowdev))
 179             dmsetup('create', snapshot, table)
 180             self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
 181
 182         finally:
 183             os.unlink(table)
 184         self.out.success('done')
 185         return "/dev/mapper/%s" % snapshot
 186
 187     def get_image(self, media):
 188         """Returns a newly created Image instance."""
 189
 190         image = Image(media, self.out)
 191         self._images.append(image)
 192         image.enable()
 193         return image
 194
 195     def destroy_image(self, image):
 196         """Destroys an Image instance previously created by get_image method.
 197         """
 198         self._images.remove(image)
 199         image.destroy()
 200
 201 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :