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("ignore partition %s of type %s" % (number, fstype), file=sys.stderr)
 100                 continue
 101
 102             path = os.path.join(basedir, 'partx', 'p'+number)
 103             mount.mount(self.image, part['start'], fstype, path)
 104
 105             num2temp[number] = path
 106             uuid2temp[part['blkid']['uuid']] = path
 107         return num2temp, uuid2temp
 108
 109     @staticmethod
 110     def _move_to_root(fstab, num2temp, uuid2temp, basedir, mount):
 111         '''Move partitions to their correct mount points according to fstab
 112         '''
 113         pairs = []
 114         for mountpoint in sorted(fstab.keys()):
 115             item = fstab[mountpoint]
 116             if 'number' in item and item['number'] in num2temp:
 117                 source = num2temp[item['number']]
 118             elif 'uuid' in item and item['uuid'] in uuid2temp:
 119                 source = uuid2temp[item['uuid']]
 120             else:
 121                 print("fstab mismatch with partition table:", \
 122                     item["entry"], file=sys.stderr)
 123                 return
 124
 125             # remove heading / otherwise the path will reduce to root
 126             target = os.path.join(basedir, 'root',
 127                                   mountpoint.lstrip(os.path.sep))
 128             pairs.append((source, target))
 129
 130         for source, target in pairs:
 131             mount.move(source, target)
 132         return True
 133
 134     def unpack(self, basedir, resourcelist):
 135         '''Unpack self into the basedir and record all resource used
 136         into resourcelist
 137         '''
 138         mount = Mount(basedir, resourcelist)
 139
 140         num2temp, uuid2temp = self._mount_to_temp(basedir, mount)
 141
 142         fstab = FSTab.guess(list(num2temp.values()))
 143         if not fstab:
 144             print("Can't find fstab file from image", file=sys.stderr)
 145             return
 146         return self._move_to_root(fstab,
 147                                   num2temp, uuid2temp,
 148                                   basedir, mount)
 149
 150
 151 def parse_args():
 152     "Parse arguments"
 153     parser = argparse.ArgumentParser()
 154     parser.add_argument('image', type=os.path.abspath,
 155                         help='image file to unpack. Only raw format is '
 156                         'supported')
 157     parser.add_argument('basedir', type=os.path.abspath,
 158                         help='directory to unpack the image')
 159     parser.add_argument('resourcelist_filename', type=os.path.abspath,
 160                         help='will record each mount point when unpacking '
 161                         'the image. Make sure call cleanup script with this '
 162                         'file name to release all allocated resources.')
 163     return parser.parse_args()
 164
 165
 166 def main():
 167     "Main"
 168     args = parse_args()
 169     img = Image(args.image)
 170     resfile = ResourceList(args.resourcelist_filename)
 171     return 0 if img.unpack(args.basedir, resfile) else 1
 172
 173
 174 if __name__ == '__main__':
 175     sys.exit(main())