Rename DiskDevice class to Image
[snf-image-creator] / 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         if default is not None:
80             return default
81
82         space = map(free_space, TMP_CANDIDATES)
83
84         max_idx = 0
85         max_val = space[0]
86         for i, val in zip(range(len(space)), space):
87             if val > max_val:
88                 max_val = val
89                 max_idx = i
90
91         # Return the candidate path with more available space
92         return TMP_CANDIDATES[max_idx]
93
94     def _add_cleanup(self, job, *args):
95         self._cleanup_jobs.append((job, args))
96
97     def _losetup(self, fname):
98         loop = losetup('-f', '--show', fname)
99         loop = loop.strip()  # remove the new-line char
100         self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
101         return loop
102
103     def _dir_to_disk(self):
104         if self.source == '/':
105             bundle = BundleVolume(self.out, self.meta)
106             image = '%s/%s.diskdump' % (self.tmp, uuid.uuid4().hex)
107
108             def check_unlink(path):
109                 if os.path.exists(path):
110                     os.unlink(path)
111
112             self._add_cleanup(check_unlink, image)
113             bundle.create_image(image)
114             return self._losetup(image)
115         raise FatalError("Using a directory as media source is supported")
116
117     def cleanup(self):
118         """Cleanup internal data. This needs to be called before the
119         program ends.
120         """
121         try:
122             while len(self._images):
123                 image = self._images.pop()
124                 image.destroy()
125         finally:
126             # Make sure those are executed even if one of the device.destroy
127             # methods throws exeptions.
128             while len(self._cleanup_jobs):
129                 job, args = self._cleanup_jobs.pop()
130                 job(*args)
131
132     def snapshot(self):
133         """Creates a snapshot of the original source media of the Disk
134         instance.
135         """
136
137         self.out.output("Examining source media `%s' ..." % self.source, False)
138         sourcedev = self.source
139         mode = os.stat(self.source).st_mode
140         if stat.S_ISDIR(mode):
141             self.out.success('looks like a directory')
142             return self._dir_to_disk()
143         elif stat.S_ISREG(mode):
144             self.out.success('looks like an image file')
145             sourcedev = self._losetup(self.source)
146         elif not stat.S_ISBLK(mode):
147             raise ValueError("Invalid media source. Only block devices, "
148                              "regular files and directories are supported.")
149         else:
150             self.out.success('looks like a block device')
151
152         # Take a snapshot and return it to the user
153         self.out.output("Snapshotting media source...", False)
154         size = blockdev('--getsz', sourcedev)
155         cowfd, cow = tempfile.mkstemp(dir=self.tmp)
156         os.close(cowfd)
157         self._add_cleanup(os.unlink, cow)
158         # Create cow sparse file
159         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
160         cowdev = self._losetup(cow)
161
162         snapshot = uuid.uuid4().hex
163         tablefd, table = tempfile.mkstemp()
164         try:
165             os.write(tablefd, "0 %d snapshot %s %s n 8" %
166                               (int(size), sourcedev, cowdev))
167             dmsetup('create', snapshot, table)
168             self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
169
170         finally:
171             os.unlink(table)
172         self.out.success('done')
173         return "/dev/mapper/%s" % snapshot
174
175     def get_image(self, media):
176         """Returns a newly created ImageCreator instance."""
177
178         image = Image(media, self.out)
179         self._images.append(image)
180         image.enable()
181         return image
182
183     def destroy_image(self, image):
184         """Destroys an ImageCreator instance previously created by
185         get_image_creator method.
186         """
187         self._images.remove(image)
188         image.destroy()
189
190 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :