imgdiff/unpack.py

   1 #!/usr/bin/env python
   2 '''This script unpack a whole image into a directory
   3 '''
   4 import os
   5 import sys
   6 import errno
   7 import argparse
   8 from subprocess import check_call
   9
  10 from imgdiff.info import get_partition_info, FSTab
  11
  12
  13 def mkdir_p(path):
  14     '''Same as mkdir -p'''
  15     try:
  16         os.makedirs(path)
  17     except OSError as err:
  18         if err.errno != errno.EEXIST:
  19             raise
  20
  21
  22 class ResourceList(object):
  23     '''
  24     Record all resource allocated into a file
  25     '''
  26     def __init__(self, filename):
  27         self.filename = filename
  28
  29     def umount(self, path):
  30         '''record a mount point'''
  31         line = 'mountpoint:%s%s' % (os.path.abspath(path), os.linesep)
  32         with open(self.filename, 'a') as writer:
  33             writer.write(line)
  34
  35
  36 class Mount(object):
  37     '''
  38     Mount image partions
  39     '''
  40     def __init__(self, limited_to_dir, resourcelist):
  41         self.limited_to_dir = limited_to_dir
  42         self.resourcelist = resourcelist
  43
  44     def _check_path(self, path):
  45         '''Check whether path is ok to mount'''
  46         if not path.startswith(self.limited_to_dir):
  47             raise ValueError("Try to mount outside of jar: " + path)
  48         if os.path.ismount(path):
  49             raise Exception("Not allowed to override an exists "
  50                             "mountpoint: " + path)
  51
  52         self.resourcelist.umount(path)
  53         mkdir_p(path)
  54
  55     def mount(self, image, offset, fstype, path):
  56         '''Mount a partition starting from perticular
  57         position of a image to a direcotry
  58         '''
  59         self._check_path(path)
  60         cmd = ['sudo', 'mount',
  61                '-o', 'ro,offset=%d' % offset,
  62                '-t', fstype,
  63                image, path]
  64         print 'Mounting', '%d@%s' % (offset, image), '->', path, '...'
  65         check_call(cmd)
  66
  67     def move(self, source, target):
  68         '''Remove mount point to another path'''
  69         self._check_path(target)
  70         cmd = ['sudo', 'mount', '--make-runbindable', '/']
  71         print 'Make runbindable ...', ' '.join(cmd)
  72         check_call(cmd)
  73         cmd = ['sudo', 'mount', '-M', source, target]
  74         print 'Moving mount point from', source, 'to', target, '...'
  75         check_call(cmd)
  76
  77
  78 class Image(object):
  79     '''A raw type image'''
  80     def __init__(self, image):
  81         self.image = image
  82         self.partab = get_partition_info(self.image)
  83
  84     @staticmethod
  85     def _is_fs_supported(fstype):
  86         '''Only support ext? and *fat*.
  87         Ignore others such as swap, tmpfs etc.
  88         '''
  89         return fstype.startswith('ext') or 'fat' in fstype
  90
  91     def _mount_to_temp(self, basedir, mount):
  92         '''Mount all partitions into temp dirs like partx/p?
  93         '''
  94         num2temp, uuid2temp = {}, {}
  95         for part in self.partab:
  96             number = str(part['number'])
  97             fstype = part['blkid']['type']
  98             if not self._is_fs_supported(fstype):
  99                 print >> sys.stderr, \
 100                     "ignore partition %s of type %s" % (number, fstype)
 101                 continue
 102
 103             path = os.path.join(basedir, 'partx', 'p'+number)
 104             mount.mount(self.image, part['start'], fstype, path)
 105
 106             num2temp[number] = path
 107             uuid2temp[part['blkid']['uuid']] = path
 108         return num2temp, uuid2temp
 109
 110     @staticmethod
 111     def _move_to_root(fstab, num2temp, uuid2temp, basedir, mount):
 112         '''Move partitions to their correct mount points according to fstab
 113         '''
 114         pairs = []
 115         for mountpoint in sorted(fstab.keys()):
 116             item = fstab[mountpoint]
 117             if 'number' in item and item['number'] in num2temp:
 118                 source = num2temp[item['number']]
 119             elif 'uuid' in item and item['uuid'] in uuid2temp:
 120                 source = uuid2temp[item['uuid']]
 121             else:
 122                 print >> sys.stderr, "fstab mismatch with partition table:", \
 123                     item["entry"]
 124                 return
 125
 126             # remove heading / otherwise the path will reduce to root
 127             target = os.path.join(basedir, 'root',
 128                                   mountpoint.lstrip(os.path.sep))
 129             pairs.append((source, target))
 130
 131         for source, target in pairs:
 132             mount.move(source, target)
 133         return True
 134
 135     def unpack(self, basedir, resourcelist):
 136         '''Unpack self into the basedir and record all resource used
 137         into resourcelist
 138         '''
 139         mount = Mount(basedir, resourcelist)
 140
 141         num2temp, uuid2temp = self._mount_to_temp(basedir, mount)
 142
 143         fstab = FSTab.guess(num2temp.values())
 144         if not fstab:
 145             print >> sys.stderr, "Can't find fstab file from image"
 146             return
 147         return self._move_to_root(fstab,
 148                                   num2temp, uuid2temp,
 149                                   basedir, mount)
 150
 151
 152 def parse_args():
 153     "Parse arguments"
 154     parser = argparse.ArgumentParser()
 155     parser.add_argument('image', type=os.path.abspath,
 156                         help='image file to unpack. Only raw format is '
 157                         'supported')
 158     parser.add_argument('basedir', type=os.path.abspath,
 159                         help='directory to unpack the image')
 160     parser.add_argument('resourcelist_filename', type=os.path.abspath,
 161                         help='will record each mount point when unpacking '
 162                         'the image. Make sure call cleanup script with this '
 163                         'file name to release all allocated resources.')
 164     return parser.parse_args()
 165
 166
 167 def main():
 168     "Main"
 169     args = parse_args()
 170     img = Image(args.image)
 171     resfile = ResourceList(args.resourcelist_filename)
 172     return 0 if img.unpack(args.basedir, resfile) else 1
 173
 174
 175 if __name__ == '__main__':
 176     sys.exit(main())