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