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