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         self.mounted_ro = False
  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         if not self.mounted:
 124             do_unmount = True
 125             self.mount(readonly=True)
 126         else:
 127             do_unmount = False
 128
 129         try:
 130             cls = os_cls(self.distro, self.ostype)
 131             self._os = cls(self.root, self.g, self.out)
 132
 133         finally:
 134             if do_unmount:
 135                 self.umount()
 136
 137         return self._os
 138
 139     os = property(_get_os)
 140
 141     def destroy(self):
 142         """Destroy this Image instance."""
 143
 144         # In new guestfs versions, there is a handy shutdown method for this
 145         try:
 146             if self.guestfs_enabled:
 147                 self.g.umount_all()
 148                 self.g.sync()
 149         finally:
 150             # Close the guestfs handler if open
 151             self.g.close()
 152
 153 #    def progress_callback(self, ev, eh, buf, array):
 154 #        position = array[2]
 155 #        total = array[3]
 156 #
 157 #        self.progressbar.goto((position * 100) // total)
 158
 159     def mount(self, readonly=False):
 160         """Mount all disk partitions in a correct order."""
 161
 162         msg = "Mounting the media%s ..." % (" read-only" if readonly else "")
 163         self.out.output(msg, False)
 164
 165         #If something goes wrong when mounting rw, remount the filesystem ro
 166         remount_ro = False
 167         rw_mpoints = ('/', '/etc', '/root', '/home', '/var')
 168
 169         # Sort the keys to mount the fs in a correct order.
 170         # / should be mounted befor /boot, etc
 171         def compare(a, b):
 172             if len(a[0]) > len(b[0]):
 173                 return 1
 174             elif len(a[0]) == len(b[0]):
 175                 return 0
 176             else:
 177                 return -1
 178         mps = self.g.inspect_get_mountpoints(self.root)
 179         mps.sort(compare)
 180
 181         mopts = 'ro' if readonly else 'rw'
 182         for mp, dev in mps:
 183             if self.ostype == 'freebsd':
 184                 # libguestfs can't handle correct freebsd partitions on GUID
 185                 # Partition Table. We have to do the translation to linux
 186                 # device names ourselves
 187                 m = re.match('^/dev/((?:ada)|(?:vtbd))(\d+)p(\d+)$', dev)
 188                 if m:
 189                     m2 = int(m.group(2))
 190                     m3 = int(m.group(3))
 191                     dev = '/dev/sd%c%d' % (chr(ord('a') + m2), m3)
 192             try:
 193                 self.g.mount_options(mopts, dev, mp)
 194             except RuntimeError as msg:
 195                 if self.ostype == 'freebsd':
 196                     freebsd_mopts = "ufstype=ufs2,%s" % mopts
 197                     try:
 198                         self.g.mount_vfs(freebsd_mopts, 'ufs', dev, mp)
 199                     except RuntimeError as msg:
 200                         if readonly is False and mp in rw_mpoints:
 201                             remount_ro = True
 202                             break
 203                 elif readonly is False and mp in rw_mpoints:
 204                     remount_ro = True
 205                     break
 206                 else:
 207                     self.out.warn("%s (ignored)" % msg)
 208         if remount_ro:
 209             self.out.warn("Unable to mount %s read-write. "
 210                           "Remounting everything read-only..." % mp)
 211             self.umount()
 212             self.mount(True)
 213         else:
 214             self.mounted = True
 215             self.mounted_ro = readonly
 216             self.out.success("done")
 217
 218     def umount(self):
 219         """Umount all mounted filesystems."""
 220         self.g.umount_all()
 221         self.mounted = False
 222
 223     def _last_partition(self):
 224         """Return the last partition of the image disk"""
 225         if self.meta['PARTITION_TABLE'] not in 'msdos' 'gpt':
 226             msg = "Unsupported partition table: %s. Only msdos and gpt " \
 227                 "partition tables are supported" % self.meta['PARTITION_TABLE']
 228             raise FatalError(msg)
 229
 230         is_extended = lambda p: \
 231             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 232             in (0x5, 0xf)
 233         is_logical = lambda p: \
 234             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 235
 236         partitions = self.g.part_list(self.guestfs_device)
 237         last_partition = partitions[-1]
 238
 239         if is_logical(last_partition):
 240             # The disk contains extended and logical partitions....
 241             extended = filter(is_extended, partitions)[0]
 242             last_primary = [p for p in partitions if p['part_num'] <= 4][-1]
 243
 244             # check if extended is the last primary partition
 245             if last_primary['part_num'] > extended['part_num']:
 246                 last_partition = last_primary
 247
 248         return last_partition
 249
 250     def shrink(self):
 251         """Shrink the image.
 252
 253         This is accomplished by shrinking the last file system of the
 254         image and then updating the partition table. The new disk size
 255         (in bytes) is returned.
 256
 257         ATTENTION: make sure unmount is called before shrink
 258         """
 259         get_fstype = lambda p: \
 260             self.g.vfs_type("%s%d" % (self.guestfs_device, p['part_num']))
 261         is_logical = lambda p: \
 262             self.meta['PARTITION_TABLE'] == 'msdos' and p['part_num'] > 4
 263         is_extended = lambda p: \
 264             self.meta['PARTITION_TABLE'] == 'msdos' and \
 265             self.g.part_get_mbr_id(self.guestfs_device, p['part_num']) \
 266             in (0x5, 0xf)
 267
 268         part_add = lambda ptype, start, stop: \
 269             self.g.part_add(self.guestfs_device, ptype, start, stop)
 270         part_del = lambda p: self.g.part_del(self.guestfs_device, p)
 271         part_get_id = lambda p: self.g.part_get_mbr_id(self.guestfs_device, p)
 272         part_set_id = lambda p, id: \
 273             self.g.part_set_mbr_id(self.guestfs_device, p, id)
 274         part_get_bootable = lambda p: \
 275             self.g.part_get_bootable(self.guestfs_device, p)
 276         part_set_bootable = lambda p, bootable: \
 277             self.g.part_set_bootable(self.guestfs_device, p, bootable)
 278
 279         MB = 2 ** 20
 280
 281         self.out.output("Shrinking image (this may take a while) ...", False)
 282
 283         sector_size = self.g.blockdev_getss(self.guestfs_device)
 284
 285         last_part = None
 286         fstype = None
 287         while True:
 288             last_part = self._last_partition()
 289             fstype = get_fstype(last_part)
 290
 291             if fstype == 'swap':
 292                 self.meta['SWAP'] = "%d:%s" % \
 293                     (last_part['part_num'],
 294                      (last_part['part_size'] + MB - 1) // MB)
 295                 part_del(last_part['part_num'])
 296                 continue
 297             elif is_extended(last_part):
 298                 part_del(last_part['part_num'])
 299                 continue
 300
 301             # Most disk manipulation programs leave 2048 sectors after the last
 302             # partition
 303             new_size = last_part['part_end'] + 1 + 2048 * sector_size
 304             self.size = min(self.size, new_size)
 305             break
 306
 307         if not re.match("ext[234]", fstype):
 308             self.out.warn("Don't know how to resize %s partitions." % fstype)
 309             return self.size
 310
 311         part_dev = "%s%d" % (self.guestfs_device, last_part['part_num'])
 312         self.g.e2fsck_f(part_dev)
 313         self.g.resize2fs_M(part_dev)
 314
 315         out = self.g.tune2fs_l(part_dev)
 316         block_size = int(filter(lambda x: x[0] == 'Block size', out)[0][1])
 317         block_cnt = int(filter(lambda x: x[0] == 'Block count', out)[0][1])
 318
 319         start = last_part['part_start'] / sector_size
 320         end = start + (block_size * block_cnt) / sector_size - 1
 321
 322         if is_logical(last_part):
 323             partitions = self.g.part_list(self.guestfs_device)
 324
 325             logical = []  # logical partitions
 326             for partition in partitions:
 327                 if partition['part_num'] < 4:
 328                     continue
 329                 logical.append({
 330                     'num': partition['part_num'],
 331                     'start': partition['part_start'] / sector_size,
 332                     'end': partition['part_end'] / sector_size,
 333                     'id': part_get_id(partition['part_num']),
 334                     'bootable': part_get_bootable(partition['part_num'])
 335                 })
 336
 337             logical[-1]['end'] = end  # new end after resize
 338
 339             # Recreate the extended partition
 340             extended = filter(is_extended, partitions)[0]
 341             part_del(extended['part_num'])
 342             part_add('e', extended['part_start'] / sector_size, end)
 343
 344             # Create all the logical partitions back
 345             for l in logical:
 346                 part_add('l', l['start'], l['end'])
 347                 part_set_id(l['num'], l['id'])
 348                 part_set_bootable(l['num'], l['bootable'])
 349         else:
 350             # Recreate the last partition
 351             if self.meta['PARTITION_TABLE'] == 'msdos':
 352                 last_part['id'] = part_get_id(last_part['part_num'])
 353
 354             last_part['bootable'] = part_get_bootable(last_part['part_num'])
 355             part_del(last_part['part_num'])
 356             part_add('p', start, end)
 357             part_set_bootable(last_part['part_num'], last_part['bootable'])
 358
 359             if self.meta['PARTITION_TABLE'] == 'msdos':
 360                 part_set_id(last_part['part_num'], last_part['id'])
 361
 362         new_size = (end + 1) * sector_size
 363
 364         assert (new_size <= self.size)
 365
 366         if self.meta['PARTITION_TABLE'] == 'gpt':
 367             ptable = GPTPartitionTable(self.device)
 368             self.size = ptable.shrink(new_size, self.size)
 369         else:
 370             self.size = min(new_size + 2048 * sector_size, self.size)
 371
 372         self.out.success("new size is %dMB" % ((self.size + MB - 1) // MB))
 373
 374         return self.size
 375
 376     def dump(self, outfile):
 377         """Dumps the content of the image into a file.
 378
 379         This method will only dump the actual payload, found by reading the
 380         partition table. Empty space in the end of the device will be ignored.
 381         """
 382         MB = 2 ** 20
 383         blocksize = 4 * MB  # 4MB
 384         size = self.size
 385         progr_size = (size + MB - 1) // MB  # in MB
 386         progressbar = self.out.Progress(progr_size, "Dumping image file", 'mb')
 387
 388         with open(self.device, 'r') as src:
 389             with open(outfile, "w") as dst:
 390                 left = size
 391                 offset = 0
 392                 progressbar.next()
 393                 while left > 0:
 394                     length = min(left, blocksize)
 395                     sent = sendfile(dst.fileno(), src.fileno(), offset, length)
 396
 397                     # Workaround for python-sendfile API change. In
 398                     # python-sendfile 1.2.x (py-sendfile) the returning value
 399                     # of sendfile is a tuple, where in version 2.x (pysendfile)
 400                     # it is just a sigle integer.
 401                     if isinstance(sent, tuple):
 402                         sent = sent[1]
 403
 404                     offset += sent
 405                     left -= sent
 406                     progressbar.goto((size - left) // MB)
 407         progressbar.success('image file %s was successfully created' % outfile)
 408
 409 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :