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