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