disk.py 6.76 KB
Newer Older
Nikos Skalkotos's avatar
Nikos Skalkotos committed
1 2
# -*- coding: utf-8 -*-
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
3
# Copyright (C) 2011-2014 GRNET S.A.
4
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
5 6 7 8
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
9
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
10 11 12 13
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
14
#
Nikos Skalkotos's avatar
Nikos Skalkotos committed
15 16
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
Nikos Skalkotos's avatar
Nikos Skalkotos committed
17

Nikos Skalkotos's avatar
Nikos Skalkotos committed
18 19
"""Module hosting the Disk class."""

20
from image_creator.util import get_command, try_fail_repeat, free_space, \
21
    FatalError, create_snapshot
22
from image_creator.bundle_volume import BundleVolume
23
from image_creator.image import Image
24

Nikos Skalkotos's avatar
Nikos Skalkotos committed
25 26 27 28
import stat
import os
import tempfile
import uuid
Nikos Skalkotos's avatar
Nikos Skalkotos committed
29
import shutil
Nikos Skalkotos's avatar
Nikos Skalkotos committed
30

31 32 33 34
dd = get_command('dd')
dmsetup = get_command('dmsetup')
losetup = get_command('losetup')
blockdev = get_command('blockdev')
35 36


Nikos Skalkotos's avatar
Nikos Skalkotos committed
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
def get_tmp_dir(default=None):
    """Check tmp directory candidates and return the one with the most
    available space.
    """
    if default is not None:
        return default

    TMP_CANDIDATES = ['/var/tmp', os.path.expanduser('~'), '/mnt']

    space = map(free_space, TMP_CANDIDATES)

    max_idx = 0
    max_val = space[0]
    for i, val in zip(range(len(space)), space):
        if val > max_val:
            max_val = val
            max_idx = i

    # Return the candidate path with more available space
    return TMP_CANDIDATES[max_idx]
57 58


Nikos Skalkotos's avatar
Nikos Skalkotos committed
59
class Disk(object):
60 61 62 63 64 65
    """This class represents a hard disk hosting an Operating System

    A Disk instance never alters the source media it is created from.
    Any change is done on a snapshot created by the device-mapper of
    the Linux kernel.
    """
Nikos Skalkotos's avatar
Nikos Skalkotos committed
66

67
    def __init__(self, source, output, tmp=None):
68
        """Create a new Disk instance out of a source media. The source
69 70
        media can be an image file, a block device or a directory.
        """
Nikos Skalkotos's avatar
Nikos Skalkotos committed
71
        self._cleanup_jobs = []
72
        self._images = []
73
        self._file = None
Nikos Skalkotos's avatar
Nikos Skalkotos committed
74
        self.source = source
75
        self.out = output
76
        self.meta = {}
77
        self.tmp = tempfile.mkdtemp(prefix='.snf_image_creator.',
Nikos Skalkotos's avatar
Nikos Skalkotos committed
78
                                    dir=get_tmp_dir(tmp))
79

Nikos Skalkotos's avatar
Nikos Skalkotos committed
80
        self._add_cleanup(shutil.rmtree, self.tmp)
81

Nikos Skalkotos's avatar
Nikos Skalkotos committed
82
    def _add_cleanup(self, job, *args):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
83
        """Add a new job in the cleanup list"""
Nikos Skalkotos's avatar
Nikos Skalkotos committed
84 85 86
        self._cleanup_jobs.append((job, args))

    def _losetup(self, fname):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
87 88 89
        """Setup a loop device and add it to the cleanup list. The loop device
        will be detached when cleanup is called.
        """
90
        loop = losetup('-f', '--show', fname)
91
        loop = loop.strip()  # remove the new-line char
92
        self._add_cleanup(try_fail_repeat, losetup, '-d', loop)
93
        return loop
Nikos Skalkotos's avatar
Nikos Skalkotos committed
94 95

    def _dir_to_disk(self):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
96
        """Create a disk out of a directory"""
Nikos Skalkotos's avatar
Nikos Skalkotos committed
97
        if self.source == '/':
98
            bundle = BundleVolume(self.out, self.meta)
99
            image = '%s/%s.raw' % (self.tmp, uuid.uuid4().hex)
100 101

            def check_unlink(path):
102
                """Unlinks file if exists"""
103 104 105 106 107
                if os.path.exists(path):
                    os.unlink(path)

            self._add_cleanup(check_unlink, image)
            bundle.create_image(image)
108
            return image
Nikos Skalkotos's avatar
Nikos Skalkotos committed
109
        raise FatalError("Using a directory as media source is supported")
Nikos Skalkotos's avatar
Nikos Skalkotos committed
110 111

    def cleanup(self):
112 113 114
        """Cleanup internal data. This needs to be called before the
        program ends.
        """
115
        try:
116 117 118
            while len(self._images):
                image = self._images.pop()
                image.destroy()
119 120 121 122 123 124
        finally:
            # Make sure those are executed even if one of the device.destroy
            # methods throws exeptions.
            while len(self._cleanup_jobs):
                job, args = self._cleanup_jobs.pop()
                job(*args)
Nikos Skalkotos's avatar
Nikos Skalkotos committed
125

126
    @property
127 128
    def file(self):
        """Convert the source media into a file"""
129

130 131
        if self._file is not None:
            return self._file
Nikos Skalkotos's avatar
Nikos Skalkotos committed
132

133
        self.out.output("Examining source media `%s' ..." % self.source, False)
134 135
        mode = os.stat(self.source).st_mode
        if stat.S_ISDIR(mode):
136
            self.out.success('looks like a directory')
137
            self._file = self._dir_to_disk()
138
        elif stat.S_ISREG(mode):
139
            self.out.success('looks like an image file')
140
            self._file = self.source
141
        elif not stat.S_ISBLK(mode):
142
            raise FatalError("Invalid media source. Only block devices, "
Nikos Skalkotos's avatar
Nikos Skalkotos committed
143
                             "regular files and directories are supported.")
144
        else:
145
            self.out.success('looks like a block device')
146
            self._file = self.source
Nikos Skalkotos's avatar
Nikos Skalkotos committed
147

148
        return self._file
149 150 151 152 153

    def snapshot(self):
        """Creates a snapshot of the original source media of the Disk
        instance.
        """
Nikos Skalkotos's avatar
Nikos Skalkotos committed
154 155 156 157 158

        if self.source == '/':
            self.out.warn("Snapshotting ignored for host bundling mode.")
            return self.file

159
        self.out.output("Snapshotting media source ...", False)
160 161 162 163 164 165 166 167 168 169 170

        # Create a qcow2 snapshot for image files
        if not stat.S_ISBLK(os.stat(self.file).st_mode):
            snapshot = create_snapshot(self.file, self.tmp)
            self._add_cleanup(os.unlink, snapshot)
            self.out.success('done')
            return snapshot

        # Create a device-mapper snapshot for block devices
        size = int(blockdev('--getsz', self.file))

171
        cowfd, cow = tempfile.mkstemp(dir=self.tmp)
172 173
        os.close(cowfd)
        self._add_cleanup(os.unlink, cow)
174
        # Create cow sparse file
175
        dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % size)
176 177
        cowdev = self._losetup(cow)

178
        snapshot = 'snf-image-creator-snapshot-%s' % uuid.uuid4().hex
179 180
        tablefd, table = tempfile.mkstemp()
        try:
181
            try:
182 183
                os.write(tablefd, "0 %d snapshot %s %s n 8\n" %
                         (size, self.file, cowdev))
184 185 186
            finally:
                os.close(tablefd)

187
            dmsetup('create', snapshot, table)
188
            self._add_cleanup(try_fail_repeat, dmsetup, 'remove', snapshot)
189 190
        finally:
            os.unlink(table)
191
        self.out.success('done')
192 193
        return "/dev/mapper/%s" % snapshot

194
    def get_image(self, media, **kargs):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
195
        """Returns a newly created Image instance."""
196

197
        image = Image(media, self.out, **kargs)
198 199 200
        self._images.append(image)
        image.enable()
        return image
Nikos Skalkotos's avatar
Nikos Skalkotos committed
201

202
    def destroy_image(self, image):
Nikos Skalkotos's avatar
Nikos Skalkotos committed
203
        """Destroys an Image instance previously created by get_image method.
204
        """
205 206
        self._images.remove(image)
        image.destroy()
207

Nikos Skalkotos's avatar
Nikos Skalkotos committed
208
# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :