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 from image_creator.bundle_volume import BundleVolume
  38
  39 import stat
  40 import os
  41 import tempfile
  42 import uuid
  43 import re
  44 import sys
  45 import guestfs
  46 import time
  47 from sendfile import sendfile
  48
  49
  50 dd = get_command('dd')
  51 dmsetup = get_command('dmsetup')
  52 losetup = get_command('losetup')
  53 blockdev = get_command('blockdev')
  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):
  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         self._cleanup_jobs = []
  68         self._devices = []
  69         self.source = source
  70         self.out = output
  71         self.meta = {}
  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         if self.source == '/':
  84             bundle = BundleVolume(self.out, self.meta)
  85             image = bundle.create_image()
  86             self._add_cleanup(os.unlink, image)
  87             return self._losetup(image)
  88         raise FatalError("Using a directory as media source is supported")
  89
  90     def cleanup(self):
  91         """Cleanup internal data. This needs to be called before the
  92         program ends.
  93         """
  94         try:
  95             while len(self._devices):
  96                 device = self._devices.pop()
  97                 device.destroy()
  98         finally:
  99             # Make sure those are executed even if one of the device.destroy
 100             # methods throws exeptions.
 101             while len(self._cleanup_jobs):
 102                 job, args = self._cleanup_jobs.pop()
 103                 job(*args)
 104
 105     def snapshot(self):
 106         """Creates a snapshot of the original source media of the Disk
 107         instance.
 108         """
 109
 110         self.out.output("Examining source media `%s'..." % self.source, False)
 111         sourcedev = self.source
 112         mode = os.stat(self.source).st_mode
 113         if stat.S_ISDIR(mode):
 114             self.out.success('looks like a directory')
 115             return self._dir_to_disk()
 116         elif stat.S_ISREG(mode):
 117             self.out.success('looks like an image file')
 118             sourcedev = self._losetup(self.source)
 119         elif not stat.S_ISBLK(mode):
 120             raise ValueError("Invalid media source. Only block devices, "
 121                              "regular files and directories are supported.")
 122         else:
 123             self.out.success('looks like a block device')
 124
 125         # Take a snapshot and return it to the user
 126         self.out.output("Snapshotting media source...", False)
 127         size = blockdev('--getsz', sourcedev)
 128         cowfd, cow = tempfile.mkstemp()
 129         os.close(cowfd)
 130         self._add_cleanup(os.unlink, cow)
 131         # Create cow sparse file
 132         dd('if=/dev/null', 'of=%s' % cow, 'bs=512', 'seek=%d' % int(size))
 133         cowdev = self._losetup(cow)
 134
 135         snapshot = uuid.uuid4().hex
 136         tablefd, table = tempfile.mkstemp()
 137         try:
 138             os.write(tablefd, "0 %d snapshot %s %s n 8" %
 139                               (int(size), sourcedev, cowdev))
 140             dmsetup('create', snapshot, table)
 141             self._add_cleanup(dmsetup, 'remove', snapshot)
 142             # Sometimes dmsetup remove fails with Device or resource busy,
 143             # although everything is cleaned up and the snapshot is not
 144             # used by anyone. Add a 2 seconds delay to be on the safe side.
 145             self._add_cleanup(time.sleep, 2)
 146
 147         finally:
 148             os.unlink(table)
 149         self.out.success('done')
 150         return "/dev/mapper/%s" % snapshot
 151
 152     def get_device(self, media):
 153         """Returns a newly created DiskDevice instance."""
 154
 155         new_device = DiskDevice(media, self.out)
 156         self._devices.append(new_device)
 157         new_device.enable()
 158         return new_device
 159
 160     def destroy_device(self, device):
 161         """Destroys a DiskDevice instance previously created by
 162         get_device method.
 163         """
 164         self._devices.remove(device)
 165         device.destroy()
 166
 167
 168 class DiskDevice(object):
 169     """This class represents a block device hosting an Operating System
 170     as created by the device-mapper.
 171     """
 172
 173     def __init__(self, device, output, bootable=True, meta={}):
 174         """Create a new DiskDevice."""
 175
 176         self.real_device = device
 177         self.out = output
 178         self.bootable = bootable
 179         self.meta = meta
 180         self.progress_bar = None
 181         self.guestfs_device = None
 182         self.size = 0
 183
 184         self.g = guestfs.GuestFS()
 185         self.g.add_drive_opts(self.real_device, readonly=0)
 186
 187         # Before version 1.17.14 the recovery process, which is a fork of the
 188         # original process that called libguestfs, did not close its inherited
 189         # file descriptors. This can cause problems especially if the parent
 190         # process has opened pipes. Since the recovery process is an optional
 191         # feature of libguestfs, it's better to disable it.
 192         self.g.set_recovery_proc(0)
 193         version = self.g.version()
 194         if version['major'] > 1 or \
 195             (version['major'] == 1 and (version['minor'] >= 18 or
 196                                         (version['minor'] == 17 and
 197                                          version['release'] >= 14))):
 198             self.g.set_recovery_proc(1)
 199             self.out.output("Enabling recovery proc")
 200
 201         #self.g.set_trace(1)
 202         #self.g.set_verbose(1)
 203
 204         self.guestfs_enabled = False
 205
 206     def enable(self):
 207         """Enable a newly created DiskDevice"""
 208
 209         self.out.output('Launching helper VM (may take a while) ...', False)
 210         # self.progressbar = self.out.Progress(100, "Launching helper VM",
 211         #                                     "percent")
 212         # eh = self.g.set_event_callback(self.progress_callback,
 213         #                               guestfs.EVENT_PROGRESS)
 214         self.g.launch()
 215         self.guestfs_enabled = True
 216         # self.g.delete_event_callback(eh)
 217         # self.progressbar.success('done')
 218         # self.progressbar = None
 219         self.out.success('done')
 220
 221         self.out.output('Inspecting Operating System ...', False)
 222         roots = self.g.inspect_os()
 223         if len(roots) == 0:
 224             raise FatalError("No operating system found")
 225         if len(roots) > 1:
 226             raise FatalError("Multiple operating systems found."
 227                              "We only support images with one OS.")
 228         self.root = roots[0]
 229         self.guestfs_device = self.g.part_to_dev(self.root)
 230         self.size = self.g.blockdev_getsize64(self.guestfs_device)
 231         self.meta['PARTITION_TABLE'] = \
 232             self.g.part_get_parttype(self.guestfs_device)
 233
 234         self.ostype = self.g.inspect_get_type(self.root)
 235         self.distro = self.g.inspect_get_distro(self.root)
 236         self.out.success('found a(n) %s system' % self.distro)
 237
 238     def destroy(self):
 239         """Destroy this DiskDevice instance."""
 240
 241         # In new guestfs versions, there is a handy shutdown method for this
 242         try:
 243             if self.guestfs_enabled:
 244                 self.g.umount_all()
 245                 self.g.sync()
 246         finally:
 247             # Close the guestfs handler if open
 248             self.g.close()
 249
 250 #    def progress_callback(self, ev, eh, buf, array):
 251 #        position = array[2]
 252 #        total = array[3]
 253 #
 254 #        self.progressbar.goto((position * 100) // total)
 255
 256     def mount(self, readonly=False):
 257         """Mount all disk partitions in a correct order."""
 258
 259         mount = self.g.mount_ro if readonly else self.g.mount
 260         msg = " read-only" if readonly else ""
 261         self.out.output("Mounting the media%s..." % msg, False)
 262         mps = self.g.inspect_get_mountpoints(self.root)
 263
 264         # Sort the keys to mount the fs in a correct order.
 265         # / should be mounted befor /boot, etc
 266         def compare(a, b):
 267             if len(a[0]) > len(b[0]):
 268                 return 1
 269             elif len(a[0]) == len(b[0]):
 270                 return 0
 271             else:
 272                 return -1
 273         mps.sort(compare)
 274         for mp, dev in mps:
 275             try:
 276                 mount(dev, mp)
 277             except RuntimeError as msg:
 278                 self.out.warn("%s (ignored)" % msg)
 279         self.out.success("done")
 280
 281     def umount(self):
 282         """Umount all mounted filesystems."""
 283         self.g.umount_all()
 284
 285     def _last_partition(self):
 286         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 287             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 288                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 289             raise FatalError(msg)
 290
 291         is_extended = lambda p: \
 292             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
 293         is_logical = lambda p: \
 294             self.meta['PARTITION_TABLE'] != 'msdos' and p['part_num'] > 4
 295
 296         partitions = self.g.part_list(self.guestfs_device)
 297         last_partition = partitions[-1]
 298
 299         if is_logical(last_partition):
 300             # The disk contains extended and logical partitions....
 301             extended = [p for p in partitions if is_extended(p)][0]
 302             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 303
 304             # check if extended is the last primary partition
 305             if last_primary['part_num'] > extended['part_num']:
 306                 last_partition = last_primary
 307
 308         return last_partition
 309
 310     def shrink(self):
 311         """Shrink the disk.
 312
 313         This is accomplished by shrinking the last filesystem in the
 314         disk and then updating the partition table. The new disk size
 315         (in bytes) is returned.
 316
 317         ATTENTION: make sure unmount is called before shrink
 318         """
 319         get_fstype = lambda p: \
 320             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 321         is_logical = lambda p: \
 322             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 323         is_extended = lambda p: \
 324             self.meta['PARTITION_TABLE'] == 'msdos' and \
 325             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) == 5
 326
 327         part_add = lambda ptype, start, stop: \
 328             self.g.part_add(self.guestfs_device, ptype, start, stop)
 329         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 330         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 331         part_set_id = lambda p, id: \
 332             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 333         part_get_bootable = lambda p: \
 334             self.g.part_get_bootable(self.guestfs_device, p)
 335         part_set_bootable = lambda p, bootable: \
 336             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 337
 338         MB = 2 ** 20
 339
 340         self.out.output("Shrinking image (this may take a while)...", False)
 341
 342         sector_size = self.g.blockdev_getss(self.guestfs_device)
 343
 344         last_part = None
 345         fstype = None
 346         while True:
 347             last_part = self._last_partition()
 348             fstype = get_fstype(last_part)
 349
 350             if fstype == 'swap':
 351                 self.meta['SWAP'] = "%d:%s" % \
 352                     (last_part['part_num'],
 353                      (last_part['part_size'] + MB - 1) // MB)
 354                 part_del(last_part['part_num'])
 355                 continue
 356             elif is_extended(last_part):
 357                 part_del(last_part['part_num'])
 358                 continue
 359
 360             # Most disk manipulation programs leave 2048 sectors after the last
 361             # partition
 362             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 363             self.size = min(self.size, new_size)
 364             break
 365
 366         if not re.match("ext[234]", fstype):
 367             self.out.warn("Don't know how to resize %s partitions." % fstype)
 368             return self.size
 369
 370         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 371         self.g.e2fsck_f(part_dev)
 372         self.g.resize2fs_M(part_dev)
 373
 374         out = self.g.tune2fs_l(part_dev)
 375         block_size = int(
 376             filter(lambda x: x[0] == 'Block size', out)[0][1])
 377         block_cnt = int(
 378             filter(lambda x: x[0] == 'Block count', out)[0][1])
 379
 380         start = last_part['part_start'] / sector_size
 381         end = start + (block_size * block_cnt) / sector_size - 1
 382
 383         if is_logical(last_part):
 384             partitions = self.g.part_list(self.guestfs_device)
 385
 386             logical = []  # logical partitions
 387             for partition in partitions:
 388                 if partition['part_num'] < 4:
 389                     continue
 390                 logical.append({
 391                     'num': partition['part_num'],
 392                     'start': partition['part_start'] / sector_size,
 393                     'end': partition['part_end'] / sector_size,
 394                     'id': part_get_(partition['part_num']),
 395                     'bootable': part_get_bootable(partition['part_num'])
 396                 })
 397
 398             logical[-1]['end'] = end  # new end after resize
 399
 400             # Recreate the extended partition
 401             extended = [p for p in partitions if self._is_extended(p)][0]
 402             part_del(extended['part_num'])
 403             part_add('e', extended['part_start'], end)
 404
 405             # Create all the logical partitions back
 406             for l in logical:
 407                 part_add('l', l['start'], l['end'])
 408                 part_set_id(l['num'], l['id'])
 409                 part_set_bootable(l['num'], l['bootable'])
 410         else:
 411             # Recreate the last partition
 412             if self.meta['PARTITION_TABLE'] == 'msdos':
 413                 last_part['id'] = part_get_id(last_part['part_num'])
 414
 415             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 416             part_del(last_part['part_num'])
 417             part_add('p', start, end)
 418             part_set_bootable(last_part['part_num'], last_part['bootable'])
 419
 420             if self.meta['PARTITION_TABLE'] == 'msdos':
 421                 part_set_id(last_part['part_num'], last_part['id'])
 422
 423         new_size = (end + 1) * sector_size
 424
 425         assert (new_size <= self.size)
 426
 427         if self.meta['PARTITION_TABLE'] == 'gpt':
 428             ptable = GPTPartitionTable(self.real_device)
 429             self.size = ptable.shrink(new_size, self.size)
 430         else:
 431             self.size = min(new_size + 2048 * sector_size, self.size)
 432
 433         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 434
 435         return self.size
 436
 437     def dump(self, outfile):
 438         """Dumps the content of device into a file.
 439
 440         This method will only dump the actual payload, found by reading the
 441         partition table. Empty space in the end of the device will be ignored.
 442         """
 443         MB = 2 ** 20
 444         blocksize = 4 * MB  # 4MB
 445         size = self.size
 446         progr_size = (size + MB - 1) // MB  # in MB
 447         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 448
 449         with open(self.real_device, 'r') as src:
 450             with open(outfile, "w") as dst:
 451                 left = size
 452                 offset = 0
 453                 progressbar.next()
 454                 while left > 0:
 455                     length = min(left, blocksize)
 456                     _, sent = sendfile(dst.fileno(), src.fileno(), offset,
 457                         length)
 458                     offset += sent
 459                     left -= sent
 460                     progressbar.goto((size - left) // MB)
 461         progressbar.success('image file %s was successfully created' % outfile)
 462
 463 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :