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