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 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, output):
  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         self.out = output
  72
  73     def _add_cleanup(self, job, *args):
  74         self._cleanup_jobs.append((job, args))
  75
  76     def _losetup(self, fname):
  77         loop = losetup('-f', '--show', fname)
  78         loop = loop.strip()  # remove the new-line char
  79         self._add_cleanup(losetup, '-d', loop)
  80         return loop
  81
  82     def _dir_to_disk(self):
  83         raise NotImplementedError
  84
  85     def cleanup(self):
  86         """Cleanup internal data. This needs to be called before the
  87         program ends.
  88         """
  89         while len(self._devices):
  90             device = self._devices.pop()
  91             device.destroy()
  92
  93         while len(self._cleanup_jobs):
  94             job, args = self._cleanup_jobs.pop()
  95             job(*args)
  96
  97     def snapshot(self):
  98         """Creates a snapshot of the original source media of the Disk
  99         instance.
 100         """
 101
 102         self.out.output("Examining source media `%s'..." % self.source, False)
 103         sourcedev = self.source
 104         mode = os.stat(self.source).st_mode
 105         if stat.S_ISDIR(mode):
 106             success('looks like a directory')
 107             return self._losetup(self._dir_to_disk())
 108         elif stat.S_ISREG(mode):
 109             success('looks like an image file')
 110             sourcedev = self._losetup(self.source)
 111         elif not stat.S_ISBLK(mode):
 112             raise ValueError("Invalid media source. Only block devices, "
 113                             "regular files and directories are supported.")
 114         else:
 115             self.out.success('looks like a block device')
 116
 117         # Take a snapshot and return it to the user
 118         self.out.output("Snapshotting media source...", False)
 119         size = blockdev('--getsize', sourcedev)
 120         cowfd, cow = tempfile.mkstemp()
 121         os.close(cowfd)
 122         self._add_cleanup(os.unlink, cow)
 123         # Create 1G cow sparse file
 124         dd('if=/dev/null', 'of=%s' % cow, 'bs=1k', '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         self.out.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, self.out)
 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, output, bootable=True):
 166         """Create a new DiskDevice."""
 167
 168         self.real_device = device
 169         self.out = output
 170         self.bootable = bootable
 171         self.progress_bar = None
 172         self.guestfs_device = None
 173         self.meta = {}
 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 = self.out.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         self.progressbar.success('done')
 194         self.progressbar = None
 195
 196         self.out.output('Inspecting Operating System...', False)
 197         roots = self.g.inspect_os()
 198         if len(roots) == 0:
 199             raise FatalError("No operating system found")
 200         if len(roots) > 1:
 201             raise FatalError("Multiple operating systems found."
 202                             "We only support images with one filesystem.")
 203         self.root = roots[0]
 204         self.guestfs_device = self.g.part_to_dev(self.root)
 205         self.meta['SIZE'] = self.g.blockdev_getsize64(self.guestfs_device)
 206         self.meta['PARTITION_TABLE'] = \
 207                                 self.g.part_get_parttype(self.guestfs_device)
 208
 209         self.ostype = self.g.inspect_get_type(self.root)
 210         self.distro = self.g.inspect_get_distro(self.root)
 211         self.out.success('found a(n) %s system' % self.distro)
 212
 213     def destroy(self):
 214         """Destroy this DiskDevice instance."""
 215
 216         if self.guestfs_enabled:
 217             self.g.umount_all()
 218             self.g.sync()
 219
 220         # Close the guestfs handler if open
 221         self.g.close()
 222
 223     def progress_callback(self, ev, eh, buf, array):
 224         position = array[2]
 225         total = array[3]
 226
 227         self.progressbar.goto((position * 100) // total)
 228
 229     def mount(self):
 230         """Mount all disk partitions in a correct order."""
 231
 232         self.out.output("Mounting image...", False)
 233         mps = self.g.inspect_get_mountpoints(self.root)
 234
 235         # Sort the keys to mount the fs in a correct order.
 236         # / should be mounted befor /boot, etc
 237         def compare(a, b):
 238             if len(a[0]) > len(b[0]):
 239                 return 1
 240             elif len(a[0]) == len(b[0]):
 241                 return 0
 242             else:
 243                 return -1
 244         mps.sort(compare)
 245         for mp, dev in mps:
 246             try:
 247                 self.g.mount(dev, mp)
 248             except RuntimeError as msg:
 249                 self.out.warn("%s (ignored)" % msg)
 250         self.out.success("done")
 251
 252     def umount(self):
 253         """Umount all mounted filesystems."""
 254         self.g.umount_all()
 255
 256     def _last_partition(self):
 257         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 258             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 259             "partition tables are supported" % self.meta['PARTITION_TABLE']
 260             raise FatalError(msg)
 261
 262         is_extended = lambda p: self.g.part_get_mbr_id(
 263                                     self.guestfs_device, p['part_num']) == 5
 264         is_logical = lambda p: self.meta['PARTITION_TABLE'] != 'msdos' and \
 265                                                             p['part_num'] > 4
 266
 267         partitions = self.g.part_list(self.guestfs_device)
 268         last_partition = partitions[-1]
 269
 270         if is_logical(last_partition):
 271             # The disk contains extended and logical partitions....
 272             extended = [p for p in partitions if is_extended(p)][0]
 273             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 274
 275             # check if extended is the last primary partition
 276             if last_primary['part_num'] > extended['part_num']:
 277                 last_partition = last_primary
 278
 279         return last_partition
 280
 281     def shrink(self):
 282         """Shrink the disk.
 283
 284         This is accomplished by shrinking the last filesystem in the
 285         disk and then updating the partition table. The new disk size
 286         (in bytes) is returned.
 287
 288         ATTENTION: make sure unmount is called before shrink
 289         """
 290         get_fstype = lambda p: self.g.vfs_type("%s%d" % \
 291                                         (self.guestfs_device, p['part_num']))
 292         is_logical = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
 293                                                             p['part_num'] > 4
 294         is_extended = lambda p: self.meta['PARTITION_TABLE'] == 'msdos' and \
 295                 self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
 296
 297         part_add = lambda ptype, start, stop: \
 298                     self.g.part_add(self.guestfs_device, ptype, start, stop)
 299         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 300         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 301         part_set_id = lambda p, id: self.g.part_set_mbr_id(
 302                                                     self.guestfs_device, p, id)
 303         part_get_bootable = lambda p: self.g.part_get_bootable(
 304                                                         self.guestfs_device, p)
 305         part_set_bootable = lambda p, bootable: self.g.part_set_bootable(
 306                                             self.guestfs_device, p, bootable)
 307
 308         MB = 2 ** 20
 309
 310         self.out.output("Shrinking image (this may take a while)...", False)
 311
 312         last_part = None
 313         fstype = None
 314         while True:
 315             last_part = self._last_partition()
 316             fstype = get_fstype(last_part)
 317
 318             if fstype == 'swap':
 319                 self.meta['SWAP'] = "%d:%s" % \
 320                         (last_part['part_num'],
 321                         (last_part['part_size'] + MB - 1) // MB)
 322                 part_del(last_part['part_num'])
 323                 continue
 324             elif is_extended(last_part):
 325                 part_del(last_part['part_num'])
 326                 continue
 327
 328             self.meta['SIZE'] = last_part['part_end'] + 1
 329             break
 330
 331         if not re.match("ext[234]", fstype):
 332             self.out.warn("Don't know how to resize %s partitions." % fstype)
 333             return self.meta['SIZE']
 334
 335         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 336         self.g.e2fsck_f(part_dev)
 337         self.g.resize2fs_M(part_dev)
 338
 339         out = self.g.tune2fs_l(part_dev)
 340         block_size = int(
 341             filter(lambda x: x[0] == 'Block size', out)[0][1])
 342         block_cnt = int(
 343             filter(lambda x: x[0] == 'Block count', out)[0][1])
 344
 345         sector_size = self.g.blockdev_getss(self.guestfs_device)
 346         start = last_part['part_start'] / sector_size
 347         end = start + (block_size * block_cnt) / sector_size - 1
 348
 349         if is_logical(last_part):
 350             partitions = self.g.part_list(self.guestfs_device)
 351
 352             logical = []  # logical partitions
 353             for partition in partitions:
 354                 if partition['part_num'] < 4:
 355                     continue
 356                 logical.append({
 357                     'num': partition['part_num'],
 358                     'start': partition['part_start'] / sector_size,
 359                     'end': partition['part_end'] / sector_size,
 360                     'id': part_get_(partition['part_num']),
 361                     'bootable': part_get_bootable(partition['part_num'])
 362                 })
 363
 364             logical[-1]['end'] = end  # new end after resize
 365
 366             # Recreate the extended partition
 367             extended = [p for p in partitions if self._is_extended(p)][0]
 368             part_del(extended['part_num'])
 369             part_add('e', extended['part_start'], end)
 370
 371             # Create all the logical partitions back
 372             for l in logical:
 373                 part_add('l', l['start'], l['end'])
 374                 part_set_id(l['num'], l['id'])
 375                 part_set_bootable(l['num'], l['bootable'])
 376         else:
 377             # Recreate the last partition
 378             if self.meta['PARTITION_TABLE'] == 'msdos':
 379                 last_part['id'] = part_get_id(last_part['part_num'])
 380
 381             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 382             part_del(last_part['part_num'])
 383             part_add('p', start, end)
 384             part_set_bootable(last_part['part_num'], last_part['bootable'])
 385
 386             if self.meta['PARTITION_TABLE'] == 'msdos':
 387                 part_set_id(last_part['part_num'], last_part['id'])
 388
 389         new_size = (end + 1) * sector_size
 390         self.out.success("new size is %dMB" % ((new_size + MB - 1) // MB))
 391
 392         if self.meta['PARTITION_TABLE'] == 'gpt':
 393             ptable = GPTPartitionTable(self.real_device)
 394             self.meta['SIZE'] = ptable.shrink(new_size)
 395         else:
 396             self.meta['SIZE'] = new_size
 397
 398         return self.meta['SIZE']
 399
 400     def dump(self, outfile):
 401         """Dumps the content of device into a file.
 402
 403         This method will only dump the actual payload, found by reading the
 404         partition table. Empty space in the end of the device will be ignored.
 405         """
 406         MB = 2 ** 20
 407         blocksize = 4 * MB  # 4MB
 408         size = self.meta['SIZE']
 409         progress_size = (size + MB - 1) // MB  # in MB
 410         progressbar = self.out.Progress("Dumping image file", 'mb')
 411         progressbar.max = progress_size
 412
 413         with open(self.real_device, 'r') as src:
 414             with open(outfile, "w") as dst:
 415                 left = size
 416                 offset = 0
 417                 progressbar.next()
 418                 while left > 0:
 419                     length = min(left, blocksize)
 420                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 421                     offset += sent
 422                     left -= sent
 423                     progressbar.goto((size - left) // MB)
 424         progressbar.success('image file %s was successfully created' % outfile)
 425
 426 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :