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 ImageCreator."""
  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 ImageCreator"""
  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         if hasattr(self, "_os"):
 114             return self._os
 115
 116         if not self.guestfs_enabled:
 117             self.enable()
 118
 119         if not self.mounted:
 120             do_unmount = True
 121             self.mount(readonly=True)
 122         else:
 123             do_unmount = False
 124
 125         try:
 126             cls = os_cls(self.distro, self.ostype)
 127             self._os = cls(self.root, self.g, self.out)
 128
 129         finally:
 130             if do_unmount:
 131                 self.umount()
 132
 133         return self._os
 134
 135     os = property(_get_os)
 136
 137     def destroy(self):
 138         """Destroy this ImageCreator instance."""
 139
 140         # In new guestfs versions, there is a handy shutdown method for this
 141         try:
 142             if self.guestfs_enabled:
 143                 self.g.umount_all()
 144                 self.g.sync()
 145         finally:
 146             # Close the guestfs handler if open
 147             self.g.close()
 148
 149 #    def progress_callback(self, ev, eh, buf, array):
 150 #        position = array[2]
 151 #        total = array[3]
 152 #
 153 #        self.progressbar.goto((position * 100) // total)
 154
 155     def mount(self, readonly=False):
 156         """Mount all disk partitions in a correct order."""
 157
 158         mount = self.g.mount_ro if readonly else self.g.mount
 159         msg = " read-only" if readonly else ""
 160         self.out.output("Mounting the media%s ..." % msg, False)
 161         mps = self.g.inspect_get_mountpoints(self.root)
 162
 163         # Sort the keys to mount the fs in a correct order.
 164         # / should be mounted befor /boot, etc
 165         def compare(a, b):
 166             if len(a[0]) > len(b[0]):
 167                 return 1
 168             elif len(a[0]) == len(b[0]):
 169                 return 0
 170             else:
 171                 return -1
 172         mps.sort(compare)
 173         for mp, dev in mps:
 174             try:
 175                 mount(dev, mp)
 176             except RuntimeError as msg:
 177                 self.out.warn("%s (ignored)" % msg)
 178
 179         self.mounted = True
 180         self.out.success("done")
 181
 182     def umount(self):
 183         """Umount all mounted filesystems."""
 184         self.g.umount_all()
 185         self.mounted = False
 186
 187     def _last_partition(self):
 188         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 189             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 190                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 191             raise FatalError(msg)
 192
 193         is_extended = lambda p: \
 194             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 195             in (0x5, 0xf)
 196         is_logical = lambda p: \
 197             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 198
 199         partitions = self.g.part_list(self.guestfs_device)
 200         last_partition = partitions[-1]
 201
 202         if is_logical(last_partition):
 203             # The disk contains extended and logical partitions....
 204             extended = filter(is_extended, partitions)[0]
 205             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 206
 207             # check if extended is the last primary partition
 208             if last_primary['part_num'] > extended['part_num']:
 209                 last_partition = last_primary
 210
 211         return last_partition
 212
 213     def shrink(self):
 214         """Shrink the disk.
 215
 216         This is accomplished by shrinking the last filesystem in the
 217         disk and then updating the partition table. The new disk size
 218         (in bytes) is returned.
 219
 220         ATTENTION: make sure unmount is called before shrink
 221         """
 222         get_fstype = lambda p: \
 223             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 224         is_logical = lambda p: \
 225             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 226         is_extended = lambda p: \
 227             self.meta['PARTITION_TABLE'] == 'msdos' and \
 228             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 229             in (0x5, 0xf)
 230
 231         part_add = lambda ptype, start, stop: \
 232             self.g.part_add(self.guestfs_device, ptype, start, stop)
 233         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 234         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 235         part_set_id = lambda p, id: \
 236             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 237         part_get_bootable = lambda p: \
 238             self.g.part_get_bootable(self.guestfs_device, p)
 239         part_set_bootable = lambda p, bootable: \
 240             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 241
 242         MB = 2 ** 20
 243
 244         self.out.output("Shrinking image (this may take a while) ...", False)
 245
 246         sector_size = self.g.blockdev_getss(self.guestfs_device)
 247
 248         last_part = None
 249         fstype = None
 250         while True:
 251             last_part = self._last_partition()
 252             fstype = get_fstype(last_part)
 253
 254             if fstype == 'swap':
 255                 self.meta['SWAP'] = "%d:%s" % \
 256                     (last_part['part_num'],
 257                      (last_part['part_size'] + MB - 1) // MB)
 258                 part_del(last_part['part_num'])
 259                 continue
 260             elif is_extended(last_part):
 261                 part_del(last_part['part_num'])
 262                 continue
 263
 264             # Most disk manipulation programs leave 2048 sectors after the last
 265             # partition
 266             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 267             self.size = min(self.size, new_size)
 268             break
 269
 270         if not re.match("ext[234]", fstype):
 271             self.out.warn("Don't know how to resize %s partitions." % fstype)
 272             return self.size
 273
 274         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 275         self.g.e2fsck_f(part_dev)
 276         self.g.resize2fs_M(part_dev)
 277
 278         out = self.g.tune2fs_l(part_dev)
 279         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
 280         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
 281
 282         start = last_part['part_start'] / sector_size
 283         end = start + (block_size * block_cnt) / sector_size - 1
 284
 285         if is_logical(last_part):
 286             partitions = self.g.part_list(self.guestfs_device)
 287
 288             logical = []  # logical partitions
 289             for partition in partitions:
 290                 if partition['part_num'] < 4:
 291                     continue
 292                 logical.append({
 293                     'num': partition['part_num'],
 294                     'start': partition['part_start'] / sector_size,
 295                     'end': partition['part_end'] / sector_size,
 296                     'id': part_get_id(partition['part_num']),
 297                     'bootable': part_get_bootable(partition['part_num'])
 298                 })
 299
 300             logical[-1]['end'] = end  # new end after resize
 301
 302             # Recreate the extended partition
 303             extended = filter(is_extended, partitions)[0]
 304             part_del(extended['part_num'])
 305             part_add('e', extended['part_start'] / sector_size, end)
 306
 307             # Create all the logical partitions back
 308             for l in logical:
 309                 part_add('l', l['start'], l['end'])
 310                 part_set_id(l['num'], l['id'])
 311                 part_set_bootable(l['num'], l['bootable'])
 312         else:
 313             # Recreate the last partition
 314             if self.meta['PARTITION_TABLE'] == 'msdos':
 315                 last_part['id'] = part_get_id(last_part['part_num'])
 316
 317             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 318             part_del(last_part['part_num'])
 319             part_add('p', start, end)
 320             part_set_bootable(last_part['part_num'], last_part['bootable'])
 321
 322             if self.meta['PARTITION_TABLE'] == 'msdos':
 323                 part_set_id(last_part['part_num'], last_part['id'])
 324
 325         new_size = (end + 1) * sector_size
 326
 327         assert (new_size <= self.size)
 328
 329         if self.meta['PARTITION_TABLE'] == 'gpt':
 330             ptable = GPTPartitionTable(self.device)
 331             self.size = ptable.shrink(new_size, self.size)
 332         else:
 333             self.size = min(new_size + 2048 * sector_size, self.size)
 334
 335         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 336
 337         return self.size
 338
 339     def dump(self, outfile):
 340         """Dumps the content of device into a file.
 341
 342         This method will only dump the actual payload, found by reading the
 343         partition table. Empty space in the end of the device will be ignored.
 344         """
 345         MB = 2 ** 20
 346         blocksize = 4 * MB  # 4MB
 347         size = self.size
 348         progr_size = (size + MB - 1) // MB  # in MB
 349         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 350
 351         with open(self.device, 'r') as src:
 352             with open(outfile, "w") as dst:
 353                 left = size
 354                 offset = 0
 355                 progressbar.next()
 356                 while left > 0:
 357                     length = min(left, blocksize)
 358                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 359
 360                     # Workaround for python-sendfile API change. In
 361                     # python-sendfile 1.2.x (py-sendfile) the returning value
 362                     # of sendfile is a tuple, where in version 2.x (pysendfile)
 363                     # it is just a sigle integer.
 364                     if isinstance(sent, tuple):
 365                         sent = sent[1]
 366
 367                     offset += sent
 368                     left -= sent
 369                     progressbar.goto((size - left) // MB)
 370         progressbar.success('image file %s was successfully created' % outfile)
 371
 372 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :