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
  92         if len(roots) == 0 or len(roots) > 1:
  93             self.root = None
  94             self.ostype = "unsupported"
  95             self.distro = "unsupported"
  96             self.guestfs_device = '/dev/sda'
  97             self.size = self.g.blockdev_getsize64(self.guestfs_device)
  98
  99             if len(roots) > 1:
 100                 reason = "Multiple operating systems found on the media."
 101             else:
 102                 reason = "Unable to detect any operating system on the media."
 103
 104             self.set_unsupported(reason)
 105             return
 106
 107         self.root = roots[0]
 108         self.meta['PARTITION_TABLE'] = self.g.part_get_parttype('/dev/sda')
 109         self.guestfs_device = '/dev/sda'  # self.g.part_to_dev(self.root)
 110         self.size = self.g.blockdev_getsize64(self.guestfs_device)
 111
 112         self.ostype = self.g.inspect_get_type(self.root)
 113         self.distro = self.g.inspect_get_distro(self.root)
 114         self.out.success(
 115             'found a(n) %s system' %
 116             self.ostype if self.distro == "unknown" else self.distro)
 117
 118         # Inspect the OS
 119         self.os.inspect()
 120
 121     def set_unsupported(self, reason):
 122         """Flag this image us ansupported"""
 123
 124         self._unsupported = reason
 125         self.meta['UNSUPPORTED'] = reason
 126         self.out.warn('Media is not supported. Reason: %s' % reason)
 127
 128     def is_unsupported(self):
 129         """Returns if this image is unsupported"""
 130         return hasattr(self, '_unsupported')
 131
 132     def enable_guestfs(self):
 133         """Enable the guestfs handler"""
 134
 135         if self.guestfs_enabled:
 136             self.out.warn("Guestfs is already enabled")
 137             return
 138
 139         # Before version 1.18.4 the behaviour of kill_subprocess was different
 140         # and you need to reset the guestfs handler to relaunch a previously
 141         # shut down qemu backend
 142         if self.check_guestfs_version(1, 18, 4) < 0:
 143             self.g = guestfs.GuestFS()
 144
 145         self.g.add_drive_opts(self.device, readonly=0, format="raw")
 146
 147         # Before version 1.17.14 the recovery process, which is a fork of the
 148         # original process that called libguestfs, did not close its inherited
 149         # file descriptors. This can cause problems especially if the parent
 150         # process has opened pipes. Since the recovery process is an optional
 151         # feature of libguestfs, it's better to disable it.
 152         if self.check_guestfs_version(1, 17, 14) >= 0:
 153             self.out.output("Enabling recovery proc")
 154             self.g.set_recovery_proc(1)
 155         else:
 156             self.g.set_recovery_proc(0)
 157
 158         #self.g.set_trace(1)
 159         #self.g.set_verbose(1)
 160
 161         self.out.output('Launching helper VM (may take a while) ...', False)
 162         # self.progressbar = self.out.Progress(100, "Launching helper VM",
 163         #                                     "percent")
 164         # eh = self.g.set_event_callback(self.progress_callback,
 165         #                               guestfs.EVENT_PROGRESS)
 166         self.g.launch()
 167         self.guestfs_enabled = True
 168         # self.g.delete_event_callback(eh)
 169         # self.progressbar.success('done')
 170         # self.progressbar = None
 171
 172         if self.check_guestfs_version(1, 18, 4) < 0:
 173             self.g.inspect_os()  # some calls need this
 174
 175         self.out.success('done')
 176
 177     def disable_guestfs(self):
 178         """Disable the guestfs handler"""
 179
 180         if not self.guestfs_enabled:
 181             self.out.warn("Guestfs is already disabled")
 182             return
 183
 184         self.out.output("Shutting down helper VM ...", False)
 185         self.g.sync()
 186         # guestfs_shutdown which is the prefered way to shutdown the backend
 187         # process was introduced in version 1.19.16
 188         if self.check_guestfs_version(1, 19, 16) >= 0:
 189             self.g.shutdown()
 190         else:
 191             self.g.kill_subprocess()
 192
 193         self.guestfs_enabled = False
 194         self.out.success('done')
 195
 196     def _get_os(self):
 197         """Return an OS class instance for this image"""
 198         if hasattr(self, "_os"):
 199             return self._os
 200
 201         if not self.guestfs_enabled:
 202             self.enable()
 203
 204         cls = os_cls(self.distro, self.ostype)
 205         self._os = cls(self, sysprep_params=self.sysprep_params)
 206
 207         self._os.collect_metadata()
 208
 209         return self._os
 210
 211     os = property(_get_os)
 212
 213     def destroy(self):
 214         """Destroy this Image instance."""
 215
 216         # In new guestfs versions, there is a handy shutdown method for this
 217         try:
 218             if self.guestfs_enabled:
 219                 self.g.umount_all()
 220                 self.g.sync()
 221         finally:
 222             # Close the guestfs handler if open
 223             self.g.close()
 224
 225 #    def progress_callback(self, ev, eh, buf, array):
 226 #        position = array[2]
 227 #        total = array[3]
 228 #
 229 #        self.progressbar.goto((position * 100) // total)
 230
 231     def _last_partition(self):
 232         """Return the last partition of the image disk"""
 233         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 234             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 235                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 236             raise FatalError(msg)
 237
 238         is_extended = lambda p: \
 239             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 240             in (0x5, 0xf)
 241         is_logical = lambda p: \
 242             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 243
 244         partitions = self.g.part_list(self.guestfs_device)
 245         last_partition = partitions[-1]
 246
 247         if is_logical(last_partition):
 248             # The disk contains extended and logical partitions....
 249             extended = filter(is_extended, partitions)[0]
 250             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 251
 252             # check if extended is the last primary partition
 253             if last_primary['part_num'] > extended['part_num']:
 254                 last_partition = last_primary
 255
 256         return last_partition
 257
 258     def shrink(self):
 259         """Shrink the image.
 260
 261         This is accomplished by shrinking the last file system of the
 262         image and then updating the partition table. The new disk size
 263         (in bytes) is returned.
 264
 265         ATTENTION: make sure unmount is called before shrink
 266         """
 267         get_fstype = lambda p: \
 268             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 269         is_logical = lambda p: \
 270             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 271         is_extended = lambda p: \
 272             self.meta['PARTITION_TABLE'] == 'msdos' and \
 273             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 274             in (0x5, 0xf)
 275
 276         part_add = lambda ptype, start, stop: \
 277             self.g.part_add(self.guestfs_device, ptype, start, stop)
 278         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 279         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 280         part_set_id = lambda p, id: \
 281             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 282         part_get_bootable = lambda p: \
 283             self.g.part_get_bootable(self.guestfs_device, p)
 284         part_set_bootable = lambda p, bootable: \
 285             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 286
 287         MB = 2 ** 20
 288
 289         self.out.output("Shrinking image (this may take a while) ...", False)
 290
 291         if self.is_unsupported():
 292             self.out.warn("Shrinking is disabled for unsupported images")
 293             return self.size
 294
 295         sector_size = self.g.blockdev_getss(self.guestfs_device)
 296
 297         last_part = None
 298         fstype = None
 299         while True:
 300             last_part = self._last_partition()
 301             fstype = get_fstype(last_part)
 302
 303             if fstype == 'swap':
 304                 self.meta['SWAP'] = "%d:%s" % \
 305                     (last_part['part_num'],
 306                      (last_part['part_size'] + MB - 1) // MB)
 307                 part_del(last_part['part_num'])
 308                 continue
 309             elif is_extended(last_part):
 310                 part_del(last_part['part_num'])
 311                 continue
 312
 313             # Most disk manipulation programs leave 2048 sectors after the last
 314             # partition
 315             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 316             self.size = min(self.size, new_size)
 317             break
 318
 319         if not re.match("ext[234]", fstype):
 320             self.out.warn("Don't know how to shrink %s partitions." % fstype)
 321             return self.size
 322
 323         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 324
 325         if self.check_guestfs_version(1, 15, 17) >= 0:
 326             self.g.e2fsck(part_dev, forceall=1)
 327         else:
 328             self.g.e2fsck_f(part_dev)
 329
 330         self.g.resize2fs_M(part_dev)
 331
 332         out = self.g.tune2fs_l(part_dev)
 333         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
 334         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
 335
 336         start = last_part['part_start'] / sector_size
 337         end = start + (block_size * block_cnt) / sector_size - 1
 338
 339         if is_logical(last_part):
 340             partitions = self.g.part_list(self.guestfs_device)
 341
 342             logical = []  # logical partitions
 343             for partition in partitions:
 344                 if partition['part_num'] < 4:
 345                     continue
 346                 logical.append({
 347                     'num': partition['part_num'],
 348                     'start': partition['part_start'] / sector_size,
 349                     'end': partition['part_end'] / sector_size,
 350                     'id': part_get_id(partition['part_num']),
 351                     'bootable': part_get_bootable(partition['part_num'])
 352                 })
 353
 354             logical[-1]['end'] = end  # new end after resize
 355
 356             # Recreate the extended partition
 357             extended = filter(is_extended, partitions)[0]
 358             part_del(extended['part_num'])
 359             part_add('e', extended['part_start'] / sector_size, end)
 360
 361             # Create all the logical partitions back
 362             for l in logical:
 363                 part_add('l', l['start'], l['end'])
 364                 part_set_id(l['num'], l['id'])
 365                 part_set_bootable(l['num'], l['bootable'])
 366         else:
 367             # Recreate the last partition
 368             if self.meta['PARTITION_TABLE'] == 'msdos':
 369                 last_part['id'] = part_get_id(last_part['part_num'])
 370
 371             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 372             part_del(last_part['part_num'])
 373             part_add('p', start, end)
 374             part_set_bootable(last_part['part_num'], last_part['bootable'])
 375
 376             if self.meta['PARTITION_TABLE'] == 'msdos':
 377                 part_set_id(last_part['part_num'], last_part['id'])
 378
 379         new_size = (end + 1) * sector_size
 380
 381         assert (new_size <= self.size)
 382
 383         if self.meta['PARTITION_TABLE'] == 'gpt':
 384             ptable = GPTPartitionTable(self.device)
 385             self.size = ptable.shrink(new_size, self.size)
 386         else:
 387             self.size = min(new_size + 2048 * sector_size, self.size)
 388
 389         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 390
 391         return self.size
 392
 393     def dump(self, outfile):
 394         """Dumps the content of the image into a file.
 395
 396         This method will only dump the actual payload, found by reading the
 397         partition table. Empty space in the end of the device will be ignored.
 398         """
 399         MB = 2 ** 20
 400         blocksize = 4 * MB  # 4MB
 401         size = self.size
 402         progr_size = (size + MB - 1) // MB  # in MB
 403         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 404
 405         with open(self.device, 'r') as src:
 406             with open(outfile, "w") as dst:
 407                 left = size
 408                 offset = 0
 409                 progressbar.next()
 410                 while left > 0:
 411                     length = min(left, blocksize)
 412                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 413
 414                     # Workaround for python-sendfile API change. In
 415                     # python-sendfile 1.2.x (py-sendfile) the returning value
 416                     # of sendfile is a tuple, where in version 2.x (pysendfile)
 417                     # it is just a sigle integer.
 418                     if isinstance(sent, tuple):
 419                         sent = sent[1]
 420
 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 :