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

   1 # Copyright 2013 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 FatalError
  35 from image_creator.gpt import GPTPartitionTable
  36 from image_creator.os_type import os_cls
  37
  38 import re
  39 import guestfs
  40 from sendfile import sendfile
  41
  42
  43 class Image(object):
  44     """The instances of this class can create images out of block devices."""
  45
  46     def __init__(self, device, output, bootable=True, meta={}):
  47         """Create a new Image instance"""
  48
  49         self.device = device
  50         self.out = output
  51         self.bootable = bootable
  52         self.meta = meta
  53         self.progress_bar = None
  54         self.guestfs_device = None
  55         self.size = 0
  56         self.mounted = False
  57
  58         self.g = guestfs.GuestFS()
  59         self.g.add_drive_opts(self.device, readonly=0, format="raw")
  60
  61         # Before version 1.17.14 the recovery process, which is a fork of the
  62         # original process that called libguestfs, did not close its inherited
  63         # file descriptors. This can cause problems especially if the parent
  64         # process has opened pipes. Since the recovery process is an optional
  65         # feature of libguestfs, it's better to disable it.
  66         self.g.set_recovery_proc(0)
  67         version = self.g.version()
  68         if version['major'] > 1 or \
  69             (version['major'] == 1 and (version['minor'] >= 18 or
  70                                         (version['minor'] == 17 and
  71                                          version['release'] >= 14))):
  72             self.g.set_recovery_proc(1)
  73             self.out.output("Enabling recovery proc")
  74
  75         #self.g.set_trace(1)
  76         #self.g.set_verbose(1)
  77
  78         self.guestfs_enabled = False
  79
  80     def enable(self):
  81         """Enable a newly created Image instance"""
  82
  83         self.out.output('Launching helper VM (may take a while) ...', False)
  84         # self.progressbar = self.out.Progress(100, "Launching helper VM",
  85         #                                     "percent")
  86         # eh = self.g.set_event_callback(self.progress_callback,
  87         #                               guestfs.EVENT_PROGRESS)
  88         self.g.launch()
  89         self.guestfs_enabled = True
  90         # self.g.delete_event_callback(eh)
  91         # self.progressbar.success('done')
  92         # self.progressbar = None
  93         self.out.success('done')
  94
  95         self.out.output('Inspecting Operating System ...', False)
  96         roots = self.g.inspect_os()
  97         if len(roots) == 0:
  98             raise FatalError("No operating system found")
  99         if len(roots) > 1:
 100             raise FatalError("Multiple operating systems found."
 101                              "We only support images with one OS.")
 102         self.root = roots[0]
 103         self.guestfs_device = self.g.part_to_dev(self.root)
 104         self.size = self.g.blockdev_getsize64(self.guestfs_device)
 105         self.meta['PARTITION_TABLE'] = \
 106             self.g.part_get_parttype(self.guestfs_device)
 107
 108         self.ostype = self.g.inspect_get_type(self.root)
 109         self.distro = self.g.inspect_get_distro(self.root)
 110         self.out.success('found a(n) %s system' % self.distro)
 111
 112     def _get_os(self):
 113         """Return an OS class instance for this image"""
 114         if hasattr(self, "_os"):
 115             return self._os
 116
 117         if not self.guestfs_enabled:
 118             self.enable()
 119
 120         if not self.mounted:
 121             do_unmount = True
 122             self.mount(readonly=True)
 123         else:
 124             do_unmount = False
 125
 126         try:
 127             cls = os_cls(self.distro, self.ostype)
 128             self._os = cls(self.root, self.g, self.out)
 129
 130         finally:
 131             if do_unmount:
 132                 self.umount()
 133
 134         return self._os
 135
 136     os = property(_get_os)
 137
 138     def destroy(self):
 139         """Destroy this Image instance."""
 140
 141         # In new guestfs versions, there is a handy shutdown method for this
 142         try:
 143             if self.guestfs_enabled:
 144                 self.g.umount_all()
 145                 self.g.sync()
 146         finally:
 147             # Close the guestfs handler if open
 148             self.g.close()
 149
 150 #    def progress_callback(self, ev, eh, buf, array):
 151 #        position = array[2]
 152 #        total = array[3]
 153 #
 154 #        self.progressbar.goto((position * 100) // total)
 155
 156     def mount(self, readonly=False):
 157         """Mount all disk partitions in a correct order."""
 158
 159         mount = self.g.mount_ro if readonly else self.g.mount
 160         msg = " read-only" if readonly else ""
 161         self.out.output("Mounting the media%s ..." % msg, False)
 162         mps = self.g.inspect_get_mountpoints(self.root)
 163
 164         # Sort the keys to mount the fs in a correct order.
 165         # / should be mounted befor /boot, etc
 166         def compare(a, b):
 167             if len(a[0]) > len(b[0]):
 168                 return 1
 169             elif len(a[0]) == len(b[0]):
 170                 return 0
 171             else:
 172                 return -1
 173         mps.sort(compare)
 174         for mp, dev in mps:
 175             try:
 176                 mount(dev, mp)
 177             except RuntimeError as msg:
 178                 self.out.warn("%s (ignored)" % msg)
 179
 180         self.mounted = True
 181         self.out.success("done")
 182
 183     def umount(self):
 184         """Umount all mounted filesystems."""
 185         self.g.umount_all()
 186         self.mounted = False
 187
 188     def _last_partition(self):
 189         """Return the last partition of the image disk"""
 190         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 191             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 192                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 193             raise FatalError(msg)
 194
 195         is_extended = lambda p: \
 196             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 197             in (0x5, 0xf)
 198         is_logical = lambda p: \
 199             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 200
 201         partitions = self.g.part_list(self.guestfs_device)
 202         last_partition = partitions[-1]
 203
 204         if is_logical(last_partition):
 205             # The disk contains extended and logical partitions....
 206             extended = filter(is_extended, partitions)[0]
 207             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 208
 209             # check if extended is the last primary partition
 210             if last_primary['part_num'] > extended['part_num']:
 211                 last_partition = last_primary
 212
 213         return last_partition
 214
 215     def shrink(self):
 216         """Shrink the image.
 217
 218         This is accomplished by shrinking the last file system of the
 219         image and then updating the partition table. The new disk size
 220         (in bytes) is returned.
 221
 222         ATTENTION: make sure unmount is called before shrink
 223         """
 224         get_fstype = lambda p: \
 225             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 226         is_logical = lambda p: \
 227             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 228         is_extended = lambda p: \
 229             self.meta['PARTITION_TABLE'] == 'msdos' and \
 230             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 231             in (0x5, 0xf)
 232
 233         part_add = lambda ptype, start, stop: \
 234             self.g.part_add(self.guestfs_device, ptype, start, stop)
 235         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 236         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 237         part_set_id = lambda p, id: \
 238             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 239         part_get_bootable = lambda p: \
 240             self.g.part_get_bootable(self.guestfs_device, p)
 241         part_set_bootable = lambda p, bootable: \
 242             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 243
 244         MB = 2 ** 20
 245
 246         self.out.output("Shrinking image (this may take a while) ...", False)
 247
 248         sector_size = self.g.blockdev_getss(self.guestfs_device)
 249
 250         last_part = None
 251         fstype = None
 252         while True:
 253             last_part = self._last_partition()
 254             fstype = get_fstype(last_part)
 255
 256             if fstype == 'swap':
 257                 self.meta['SWAP'] = "%d:%s" % \
 258                     (last_part['part_num'],
 259                      (last_part['part_size'] + MB - 1) // MB)
 260                 part_del(last_part['part_num'])
 261                 continue
 262             elif is_extended(last_part):
 263                 part_del(last_part['part_num'])
 264                 continue
 265
 266             # Most disk manipulation programs leave 2048 sectors after the last
 267             # partition
 268             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 269             self.size = min(self.size, new_size)
 270             break
 271
 272         if not re.match("ext[234]", fstype):
 273             self.out.warn("Don't know how to resize %s partitions." % fstype)
 274             return self.size
 275
 276         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 277         self.g.e2fsck_f(part_dev)
 278         self.g.resize2fs_M(part_dev)
 279
 280         out = self.g.tune2fs_l(part_dev)
 281         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
 282         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
 283
 284         start = last_part['part_start'] / sector_size
 285         end = start + (block_size * block_cnt) / sector_size - 1
 286
 287         if is_logical(last_part):
 288             partitions = self.g.part_list(self.guestfs_device)
 289
 290             logical = []  # logical partitions
 291             for partition in partitions:
 292                 if partition['part_num'] < 4:
 293                     continue
 294                 logical.append({
 295                     'num': partition['part_num'],
 296                     'start': partition['part_start'] / sector_size,
 297                     'end': partition['part_end'] / sector_size,
 298                     'id': part_get_id(partition['part_num']),
 299                     'bootable': part_get_bootable(partition['part_num'])
 300                 })
 301
 302             logical[-1]['end'] = end  # new end after resize
 303
 304             # Recreate the extended partition
 305             extended = filter(is_extended, partitions)[0]
 306             part_del(extended['part_num'])
 307             part_add('e', extended['part_start'] / sector_size, end)
 308
 309             # Create all the logical partitions back
 310             for l in logical:
 311                 part_add('l', l['start'], l['end'])
 312                 part_set_id(l['num'], l['id'])
 313                 part_set_bootable(l['num'], l['bootable'])
 314         else:
 315             # Recreate the last partition
 316             if self.meta['PARTITION_TABLE'] == 'msdos':
 317                 last_part['id'] = part_get_id(last_part['part_num'])
 318
 319             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 320             part_del(last_part['part_num'])
 321             part_add('p', start, end)
 322             part_set_bootable(last_part['part_num'], last_part['bootable'])
 323
 324             if self.meta['PARTITION_TABLE'] == 'msdos':
 325                 part_set_id(last_part['part_num'], last_part['id'])
 326
 327         new_size = (end + 1) * sector_size
 328
 329         assert (new_size <= self.size)
 330
 331         if self.meta['PARTITION_TABLE'] == 'gpt':
 332             ptable = GPTPartitionTable(self.device)
 333             self.size = ptable.shrink(new_size, self.size)
 334         else:
 335             self.size = min(new_size + 2048 * sector_size, self.size)
 336
 337         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 338
 339         return self.size
 340
 341     def dump(self, outfile):
 342         """Dumps the content of the image into a file.
 343
 344         This method will only dump the actual payload, found by reading the
 345         partition table. Empty space in the end of the device will be ignored.
 346         """
 347         MB = 2 ** 20
 348         blocksize = 4 * MB  # 4MB
 349         size = self.size
 350         progr_size = (size + MB - 1) // MB  # in MB
 351         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 352
 353         with open(self.device, 'r') as src:
 354             with open(outfile, "w") as dst:
 355                 left = size
 356                 offset = 0
 357                 progressbar.next()
 358                 while left > 0:
 359                     length = min(left, blocksize)
 360                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 361
 362                     # Workaround for python-sendfile API change. In
 363                     # python-sendfile 1.2.x (py-sendfile) the returning value
 364                     # of sendfile is a tuple, where in version 2.x (pysendfile)
 365                     # it is just a sigle integer.
 366                     if isinstance(sent, tuple):
 367                         sent = sent[1]
 368
 369                     offset += sent
 370                     left -= sent
 371                     progressbar.goto((size - left) // MB)
 372         progressbar.success('image file %s was successfully created' % outfile)
 373
 374 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :