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 try_fail_repeat
  36 from image_creator.util import free_space
  37 from image_creator.util import FatalError
  38 from image_creator.bundle_volume import BundleVolume
  39 from image_creator.image import Image
  40
  41 import stat
  42 import os
  43 import tempfile
  44 import uuid
  45 import shutil
  46
  47 dd = get_command('dd')
  48 dmsetup = get_command('dmsetup')
  49 losetup = get_command('losetup')
  50 blockdev = get_command('blockdev')
  51
  52
  53 TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']
  54
  55
  56 class Disk(object):
  57     """This class represents a hard disk hosting an Operating System
  58
  59     A Disk instance never alters the source media it is created from.
  60     Any change is done on a snapshot created by the device-mapper of
  61     the Linux kernel.
  62     """
  63
  64     def __init__(self, source, output, tmp=None):
  65         """Create a new Disk instance out of a source media. The source
  66         media can be an image file, a block device or a directory.
  67         """
  68         self._cleanup_jobs = []
  69         self._images = []
  70         self.source = source
  71         self.out = output
  72         self.meta = {}
  73         self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
  74                                     dir=self._get_tmp_dir(tmp))
  75
  76         self._add_cleanup(shutil.rmtree, self.tmp)
  77
  78     def _get_tmp_dir(self, default=None):
  79         """Check tmp directory candidates and return the one with the most
  80         available space.
  81         """
  82         if default is not None:
  83             return default
  84
  85         space = map(free_space, TMP_CANDIDATES)
  86
  87         max_idx = 0
  88         max_val = space[0]
  89         for i, val in zip(range(len(space)), space):
  90             if val > max_val:
  91                 max_val = val
  92                 max_idx = i
  93
  94         # Return the candidate path with more available space
  95         return TMP_CANDIDATES[max_idx]
  96
  97     def _add_cleanup(self, job, *args):
  98         """Add a new job in the cleanup list"""
  99         self._cleanup_jobs.append((job, args))
 100
 101     def _losetup(self, fname):
 102         """Setup a loop device and add it to the cleanup list. The loop device
 103         will be detached when cleanup is called.
 104         """
 105         loop = losetup('-f', '--show', fname)
 106         loop = loop.strip()  # remove the new-line char
 107         self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
 108         return loop
 109
 110     def _dir_to_disk(self):
 111         """Create a disk out of a directory"""
 112         if self.source == '/':
 113             bundle = BundleVolume(self.out, self.meta)
 114             image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex)
 115
 116             def check_unlink(path):
 117                 if os.path.exists(path):
 118                     os.unlink(path)
 119
 120             self._add_cleanup(check_unlink, image)
 121             bundle.create_image(image)
 122             return self._losetup(image)
 123         raise FatalError("Using a directory as media source is supported")
 124
 125     def cleanup(self):
 126         """Cleanup internal data. This needs to be called before the
 127         program ends.
 128         """
 129         try:
 130             while len(self._images):
 131                 image = self._images.pop()
 132                 image.destroy()
 133         finally:
 134             # Make sure those are executed even if one of the device.destroy
 135             # methods throws exeptions.
 136             while len(self._cleanup_jobs):
 137                 job, args = self._cleanup_jobs.pop()
 138                 job(*args)
 139
 140     def snapshot(self):
 141         """Creates a snapshot of the original source media of the Disk
 142         instance.
 143         """
 144
 145         self.out.output("Examining source media `%s' ..." % self.source, False)
 146         sourcedev = self.source
 147         mode = os.stat(self.source).st_mode
 148         if stat.S_ISDIR(mode):
 149             self.out.success('looks like a directory')
 150             return self._dir_to_disk()
 151         elif stat.S_ISREG(mode):
 152             self.out.success('looks like an image file')
 153             sourcedev = self._losetup(self.source)
 154         elif not stat.S_ISBLK(mode):
 155             raise FatalError("Invalid media source. Only block devices, "
 156                              "regular files and directories are supported.")
 157         else:
 158             self.out.success('looks like a block device')
 159
 160         # Take a snapshot and return it to the user
 161         self.out.output("Snapshotting media source...", False)
 162         size = blockdev('--getsz', sourcedev)
 163         cowfd, cow = tempfile.mkstemp(dir=self.tmp)
 164         os.close(cowfd)
 165         self._add_cleanup(os.unlink, cow)
 166         # Create cow sparse file
 167         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
 168         cowdev = self._losetup(cow)
 169
 170         snapshot = uuid.uuid4().hex
 171         tablefd, table = tempfile.mkstemp()
 172         try:
 173             os.write(tablefd, "0 %d snapshot %s %s n 8" %
 174                               (int(size), sourcedev, cowdev))
 175             dmsetup('create', snapshot, table)
 176             self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
 177
 178         finally:
 179             os.unlink(table)
 180         self.out.success('done')
 181         return "/dev/mapper/%s" % snapshot
 182
 183     def get_image(self, media):
 184         """Returns a newly created Image instance."""
 185
 186         image = Image(media, self.out)
 187         self._images.append(image)
 188         image.enable()
 189         return image
 190
 191     def destroy_image(self, image):
 192         """Destroys an Image instance previously created by get_image method.
 193         """
 194         self._images.remove(image)
 195         image.destroy()
 196
 197 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :