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