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