code.grnet.gr Git - snf-image-creator/blob - image_creator/image.py

   1 # -*- coding: utf-8 -*-
   2 #
   3 # Copyright 2013 GRNET S.A. All rights reserved.
   4 #
   5 # Redistribution and use in source and binary forms, with or
   6 # without modification, are permitted provided that the following
   7 # conditions are met:
   8 #
   9 #   1. Redistributions of source code must retain the above
  10 #      copyright notice, this list of conditions and the following
  11 #      disclaimer.
  12 #
  13 #   2. Redistributions in binary form must reproduce the above
  14 #      copyright notice, this list of conditions and the following
  15 #      disclaimer in the documentation and/or other materials
  16 #      provided with the distribution.
  17 #
  18 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  19 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  20 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  21 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  22 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  23 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  24 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  25 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  26 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  27 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  28 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  29 # POSSIBILITY OF SUCH DAMAGE.
  30 #
  31 # The views and conclusions contained in the software and
  32 # documentation are those of the authors and should not be
  33 # interpreted as representing official policies, either expressed
  34 # or implied, of GRNET S.A.
  35
  36 from image_creator.util import FatalError
  37 from image_creator.gpt import GPTPartitionTable
  38 from image_creator.os_type import os_cls
  39
  40 import re
  41 import guestfs
  42 from sendfile import sendfile
  43
  44
  45 class Image(object):
  46     """The instances of this class can create images out of block devices."""
  47
  48     def __init__(self, device, output, **kargs):
  49         """Create a new Image instance"""
  50
  51         self.device = device
  52         self.out = output
  53
  54         self.meta = kargs['meta'] if 'meta' in kargs else {}
  55         self.sysprep_params = \
  56             kargs['sysprep_params'] if 'sysprep_params' in kargs else {}
  57
  58         self.progress_bar = None
  59         self.guestfs_device = None
  60         self.size = 0
  61
  62         self.g = guestfs.GuestFS()
  63         self.guestfs_enabled = False
  64         self.guestfs_version = self.g.version()
  65
  66     def check_guestfs_version(self, major, minor, release):
  67         """Checks if the version of the used libguestfs is smaller, equal or
  68         greater than the one specified by the major, minor and release triplet
  69
  70         Returns:
  71             < 0 if the installed version is smaller than the specified one
  72             = 0 if they are equal
  73             > 0 if the installed one is greater than the specified one
  74         """
  75
  76         for (a, b) in (self.guestfs_version['major'], major), \
  77                 (self.guestfs_version['minor'], minor), \
  78                 (self.guestfs_version['release'], release):
  79             if a != b:
  80                 return a - b
  81
  82         return 0
  83
  84     def enable(self):
  85         """Enable a newly created Image instance"""
  86
  87         self.enable_guestfs()
  88
  89         self.out.output('Inspecting Operating System ...', False)
  90         roots = self.g.inspect_os()
  91         if len(roots) == 0:
  92             raise FatalError("No operating system found")
  93         if len(roots) > 1:
  94             raise FatalError("Multiple operating systems found."
  95                              "We only support images with one OS.")
  96         self.root = roots[0]
  97         self.guestfs_device = self.g.part_to_dev(self.root)
  98         self.size = self.g.blockdev_getsize64(self.guestfs_device)
  99         self.meta['PARTITION_TABLE'] = \
 100             self.g.part_get_parttype(self.guestfs_device)
 101
 102         self.ostype = self.g.inspect_get_type(self.root)
 103         self.distro = self.g.inspect_get_distro(self.root)
 104         self.out.success(
 105             'found a(n) %s system' %
 106             self.ostype if self.distro == "unknown" else self.distro)
 107
 108     def enable_guestfs(self):
 109         """Enable the guestfs handler"""
 110
 111         if self.guestfs_enabled:
 112             self.out.warn("Guestfs is already enabled")
 113             return
 114
 115         # Before version 1.18.4 the behaviour of kill_subprocess was different
 116         # and you need to reset the guestfs handler to relaunch a previously
 117         # shut down qemu backend
 118         if self.check_guestfs_version(1, 18, 4) < 0:
 119             self.g = guestfs.GuestFS()
 120
 121         self.g.add_drive_opts(self.device, readonly=0, format="raw")
 122
 123         # Before version 1.17.14 the recovery process, which is a fork of the
 124         # original process that called libguestfs, did not close its inherited
 125         # file descriptors. This can cause problems especially if the parent
 126         # process has opened pipes. Since the recovery process is an optional
 127         # feature of libguestfs, it's better to disable it.
 128         if self.check_guestfs_version(1, 17, 14) >= 0:
 129             self.out.output("Enabling recovery proc")
 130             self.g.set_recovery_proc(1)
 131         else:
 132             self.g.set_recovery_proc(0)
 133
 134         #self.g.set_trace(1)
 135         #self.g.set_verbose(1)
 136
 137         self.out.output('Launching helper VM (may take a while) ...', False)
 138         # self.progressbar = self.out.Progress(100, "Launching helper VM",
 139         #                                     "percent")
 140         # eh = self.g.set_event_callback(self.progress_callback,
 141         #                               guestfs.EVENT_PROGRESS)
 142         self.g.launch()
 143         self.guestfs_enabled = True
 144         # self.g.delete_event_callback(eh)
 145         # self.progressbar.success('done')
 146         # self.progressbar = None
 147
 148         if self.check_guestfs_version(1, 18, 4) < 0:
 149             self.g.inspect_os()  # some calls need this
 150
 151         self.out.success('done')
 152
 153     def disable_guestfs(self):
 154         """Disable the guestfs handler"""
 155
 156         if not self.guestfs_enabled:
 157             self.out.warn("Guestfs is already disabled")
 158             return
 159
 160         self.out.output("Shutting down helper VM ...", False)
 161         self.g.sync()
 162         # guestfs_shutdown which is the prefered way to shutdown the backend
 163         # process was introduced in version 1.19.16
 164         if self.check_guestfs_version(1, 19, 16) >= 0:
 165             self.g.shutdown()
 166         else:
 167             self.g.kill_subprocess()
 168
 169         self.guestfs_enabled = False
 170         self.out.success('done')
 171
 172     def _get_os(self):
 173         """Return an OS class instance for this image"""
 174         if hasattr(self, "_os"):
 175             return self._os
 176
 177         if not self.guestfs_enabled:
 178             self.enable()
 179
 180         cls = os_cls(self.distro, self.ostype)
 181         self._os = cls(self, sysprep_params=self.sysprep_params)
 182
 183         self._os.collect_metadata()
 184
 185         return self._os
 186
 187     os = property(_get_os)
 188
 189     def destroy(self):
 190         """Destroy this Image instance."""
 191
 192         # In new guestfs versions, there is a handy shutdown method for this
 193         try:
 194             if self.guestfs_enabled:
 195                 self.g.umount_all()
 196                 self.g.sync()
 197         finally:
 198             # Close the guestfs handler if open
 199             self.g.close()
 200
 201 #    def progress_callback(self, ev, eh, buf, array):
 202 #        position = array[2]
 203 #        total = array[3]
 204 #
 205 #        self.progressbar.goto((position * 100) // total)
 206
 207     def _last_partition(self):
 208         """Return the last partition of the image disk"""
 209         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 210             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 211                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 212             raise FatalError(msg)
 213
 214         is_extended = lambda p: \
 215             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 216             in (0x5, 0xf)
 217         is_logical = lambda p: \
 218             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 219
 220         partitions = self.g.part_list(self.guestfs_device)
 221         last_partition = partitions[-1]
 222
 223         if is_logical(last_partition):
 224             # The disk contains extended and logical partitions....
 225             extended = filter(is_extended, partitions)[0]
 226             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 227
 228             # check if extended is the last primary partition
 229             if last_primary['part_num'] > extended['part_num']:
 230                 last_partition = last_primary
 231
 232         return last_partition
 233
 234     def shrink(self):
 235         """Shrink the image.
 236
 237         This is accomplished by shrinking the last file system of the
 238         image and then updating the partition table. The new disk size
 239         (in bytes) is returned.
 240
 241         ATTENTION: make sure unmount is called before shrink
 242         """
 243         get_fstype = lambda p: \
 244             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 245         is_logical = lambda p: \
 246             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 247         is_extended = lambda p: \
 248             self.meta['PARTITION_TABLE'] == 'msdos' and \
 249             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 250             in (0x5, 0xf)
 251
 252         part_add = lambda ptype, start, stop: \
 253             self.g.part_add(self.guestfs_device, ptype, start, stop)
 254         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 255         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 256         part_set_id = lambda p, id: \
 257             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 258         part_get_bootable = lambda p: \
 259             self.g.part_get_bootable(self.guestfs_device, p)
 260         part_set_bootable = lambda p, bootable: \
 261             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 262
 263         MB = 2 ** 20
 264
 265         self.out.output("Shrinking image (this may take a while) ...", False)
 266
 267         sector_size = self.g.blockdev_getss(self.guestfs_device)
 268
 269         last_part = None
 270         fstype = None
 271         while True:
 272             last_part = self._last_partition()
 273             fstype = get_fstype(last_part)
 274
 275             if fstype == 'swap':
 276                 self.meta['SWAP'] = "%d:%s" % \
 277                     (last_part['part_num'],
 278                      (last_part['part_size'] + MB - 1) // MB)
 279                 part_del(last_part['part_num'])
 280                 continue
 281             elif is_extended(last_part):
 282                 part_del(last_part['part_num'])
 283                 continue
 284
 285             # Most disk manipulation programs leave 2048 sectors after the last
 286             # partition
 287             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 288             self.size = min(self.size, new_size)
 289             break
 290
 291         if not re.match("ext[234]", fstype):
 292             self.out.warn("Don't know how to shrink %s partitions." % fstype)
 293             return self.size
 294
 295         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 296         self.g.e2fsck_f(part_dev)
 297         self.g.resize2fs_M(part_dev)
 298
 299         out = self.g.tune2fs_l(part_dev)
 300         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
 301         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
 302
 303         start = last_part['part_start'] / sector_size
 304         end = start + (block_size * block_cnt) / sector_size - 1
 305
 306         if is_logical(last_part):
 307             partitions = self.g.part_list(self.guestfs_device)
 308
 309             logical = []  # logical partitions
 310             for partition in partitions:
 311                 if partition['part_num'] < 4:
 312                     continue
 313                 logical.append({
 314                     'num': partition['part_num'],
 315                     'start': partition['part_start'] / sector_size,
 316                     'end': partition['part_end'] / sector_size,
 317                     'id': part_get_id(partition['part_num']),
 318                     'bootable': part_get_bootable(partition['part_num'])
 319                 })
 320
 321             logical[-1]['end'] = end  # new end after resize
 322
 323             # Recreate the extended partition
 324             extended = filter(is_extended, partitions)[0]
 325             part_del(extended['part_num'])
 326             part_add('e', extended['part_start'] / sector_size, end)
 327
 328             # Create all the logical partitions back
 329             for l in logical:
 330                 part_add('l', l['start'], l['end'])
 331                 part_set_id(l['num'], l['id'])
 332                 part_set_bootable(l['num'], l['bootable'])
 333         else:
 334             # Recreate the last partition
 335             if self.meta['PARTITION_TABLE'] == 'msdos':
 336                 last_part['id'] = part_get_id(last_part['part_num'])
 337
 338             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 339             part_del(last_part['part_num'])
 340             part_add('p', start, end)
 341             part_set_bootable(last_part['part_num'], last_part['bootable'])
 342
 343             if self.meta['PARTITION_TABLE'] == 'msdos':
 344                 part_set_id(last_part['part_num'], last_part['id'])
 345
 346         new_size = (end + 1) * sector_size
 347
 348         assert (new_size <= self.size)
 349
 350         if self.meta['PARTITION_TABLE'] == 'gpt':
 351             ptable = GPTPartitionTable(self.device)
 352             self.size = ptable.shrink(new_size, self.size)
 353         else:
 354             self.size = min(new_size + 2048 * sector_size, self.size)
 355
 356         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 357
 358         return self.size
 359
 360     def dump(self, outfile):
 361         """Dumps the content of the image into a file.
 362
 363         This method will only dump the actual payload, found by reading the
 364         partition table. Empty space in the end of the device will be ignored.
 365         """
 366         MB = 2 ** 20
 367         blocksize = 4 * MB  # 4MB
 368         size = self.size
 369         progr_size = (size + MB - 1) // MB  # in MB
 370         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 371
 372         with open(self.device, 'r') as src:
 373             with open(outfile, "w") as dst:
 374                 left = size
 375                 offset = 0
 376                 progressbar.next()
 377                 while left > 0:
 378                     length = min(left, blocksize)
 379                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 380
 381                     # Workaround for python-sendfile API change. In
 382                     # python-sendfile 1.2.x (py-sendfile) the returning value
 383                     # of sendfile is a tuple, where in version 2.x (pysendfile)
 384                     # it is just a sigle integer.
 385                     if isinstance(sent, tuple):
 386                         sent = sent[1]
 387
 388                     offset += sent
 389                     left -= sent
 390                     progressbar.goto((size - left) // MB)
 391         progressbar.success('image file %s was successfully created' % outfile)
 392
 393 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :