1#!/usr/bin/python2
2"""Diff 2 chromiumos images by comparing each elf file.
3
4   The script diffs every *ELF* files by dissembling every *executable*
5   section, which means it is not a FULL elf differ.
6
7   A simple usage example -
8     chromiumos_image_diff.py --image1 image-path-1 --image2 image-path-2
9
10   Note that image path should be inside the chroot, if not (ie, image is
11   downloaded from web), please specify a chromiumos checkout via
12   "--chromeos_root".
13
14   And this script should be executed outside chroot.
15"""
16
17from __future__ import print_function
18
19__author__ = 'shenhan@google.com (Han Shen)'
20
21import argparse
22import os
23import re
24import sys
25import tempfile
26
27import image_chromeos
28from cros_utils import command_executer
29from cros_utils import logger
30from cros_utils import misc
31
32
33class CrosImage(object):
34  """A cros image object."""
35
36  def __init__(self, image, chromeos_root, no_unmount):
37    self.image = image
38    self.chromeos_root = chromeos_root
39    self.mounted = False
40    self._ce = command_executer.GetCommandExecuter()
41    self.logger = logger.GetLogger()
42    self.elf_files = []
43    self.no_unmount = no_unmount
44    self.unmount_script = ''
45    self.stateful = ''
46    self.rootfs = ''
47
48  def MountImage(self, mount_basename):
49    """Mount/unpack the image."""
50
51    if mount_basename:
52      self.rootfs = '/tmp/{0}.rootfs'.format(mount_basename)
53      self.stateful = '/tmp/{0}.stateful'.format(mount_basename)
54      self.unmount_script = '/tmp/{0}.unmount.sh'.format(mount_basename)
55    else:
56      self.rootfs = tempfile.mkdtemp(suffix='.rootfs',
57                                     prefix='chromiumos_image_diff')
58      ## rootfs is like /tmp/tmpxyz012.rootfs.
59      match = re.match(r'^(.*)\.rootfs$', self.rootfs)
60      basename = match.group(1)
61      self.stateful = basename + '.stateful'
62      os.mkdir(self.stateful)
63      self.unmount_script = '{0}.unmount.sh'.format(basename)
64
65    self.logger.LogOutput('Mounting "{0}" onto "{1}" and "{2}"'.format(
66        self.image, self.rootfs, self.stateful))
67    ## First of all creating an unmount image
68    self.CreateUnmountScript()
69    command = image_chromeos.GetImageMountCommand(
70        self.chromeos_root, self.image, self.rootfs, self.stateful)
71    rv = self._ce.RunCommand(command, print_to_console=True)
72    self.mounted = (rv == 0)
73    if not self.mounted:
74      self.logger.LogError('Failed to mount "{0}" onto "{1}" and "{2}".'.format(
75          self.image, self.rootfs, self.stateful))
76    return self.mounted
77
78  def CreateUnmountScript(self):
79    command = ('sudo umount {r}/usr/local {r}/usr/share/oem '
80               '{r}/var {r}/mnt/stateful_partition {r}; sudo umount {s} ; '
81               'rmdir {r} ; rmdir {s}\n').format(r=self.rootfs, s=self.stateful)
82    f = open(self.unmount_script, 'w')
83    f.write(command)
84    f.close()
85    self._ce.RunCommand('chmod +x {}'.format(self.unmount_script),
86                        print_to_console=False)
87    self.logger.LogOutput('Created an unmount script - "{0}"'.format(
88        self.unmount_script))
89
90  def UnmountImage(self):
91    """Unmount the image and delete mount point."""
92
93    self.logger.LogOutput('Unmounting image "{0}" from "{1}" and "{2}"'.format(
94        self.image, self.rootfs, self.stateful))
95    if self.mounted:
96      command = 'bash "{0}"'.format(self.unmount_script)
97      if self.no_unmount:
98        self.logger.LogOutput(('Please unmount manually - \n'
99                               '\t bash "{0}"'.format(self.unmount_script)))
100      else:
101        if self._ce.RunCommand(command, print_to_console=True) == 0:
102          self._ce.RunCommand('rm {0}'.format(self.unmount_script))
103          self.mounted = False
104          self.rootfs = None
105          self.stateful = None
106          self.unmount_script = None
107
108    return not self.mounted
109
110  def FindElfFiles(self):
111    """Find all elf files for the image.
112
113    Returns:
114      Always true
115    """
116
117    self.logger.LogOutput('Finding all elf files in "{0}" ...'.format(
118        self.rootfs))
119    # Note '\;' must be prefixed by 'r'.
120    command = ('find "{0}" -type f -exec '
121               'bash -c \'file -b "{{}}" | grep -q "ELF"\'' r' \; '
122               r'-exec echo "{{}}" \;').format(self.rootfs)
123    self.logger.LogCmd(command)
124    _, out, _ = self._ce.RunCommandWOutput(command, print_to_console=False)
125    self.elf_files = out.splitlines()
126    self.logger.LogOutput(
127        'Total {0} elf files found.'.format(len(self.elf_files)))
128    return True
129
130
131class ImageComparator(object):
132  """A class that wraps comparsion actions."""
133
134  def __init__(self, images, diff_file):
135    self.images = images
136    self.logger = logger.GetLogger()
137    self.diff_file = diff_file
138    self.tempf1 = None
139    self.tempf2 = None
140
141  def Cleanup(self):
142    if self.tempf1 and self.tempf2:
143      command_executer.GetCommandExecuter().RunCommand(
144          'rm {0} {1}'.format(self.tempf1, self.tempf2))
145      logger.GetLogger('Removed "{0}" and "{1}".'.format(
146          self.tempf1, self.tempf2))
147
148  def CheckElfFileSetEquality(self):
149    """Checking whether images have exactly number of elf files."""
150
151    self.logger.LogOutput('Checking elf file equality ...')
152    i1 = self.images[0]
153    i2 = self.images[1]
154    t1 = i1.rootfs + '/'
155    elfset1 = set([e.replace(t1, '') for e in i1.elf_files])
156    t2 = i2.rootfs + '/'
157    elfset2 = set([e.replace(t2, '') for e in i2.elf_files])
158    dif1 = elfset1.difference(elfset2)
159    msg = None
160    if dif1:
161      msg = 'The following files are not in "{image}" - "{rootfs}":\n'.format(
162          image=i2.image, rootfs=i2.rootfs)
163      for d in dif1:
164        msg += '\t' + d + '\n'
165    dif2 = elfset2.difference(elfset1)
166    if dif2:
167      msg = 'The following files are not in "{image}" - "{rootfs}":\n'.format(
168          image=i1.image, rootfs=i1.rootfs)
169      for d in dif2:
170        msg += '\t' + d + '\n'
171    if msg:
172      self.logger.LogError(msg)
173      return False
174    return True
175
176  def CompareImages(self):
177    """Do the comparsion work."""
178
179    if not self.CheckElfFileSetEquality():
180      return False
181
182    mismatch_list = []
183    match_count = 0
184    i1 = self.images[0]
185    i2 = self.images[1]
186    self.logger.LogOutput('Start comparing {0} elf file by file ...'.format(
187        len(i1.elf_files)))
188    ## Note - i1.elf_files and i2.elf_files have exactly the same entries here.
189
190    ## Create 2 temp files to be used for all disassembed files.
191    handle, self.tempf1 = tempfile.mkstemp()
192    os.close(handle)  # We do not need the handle
193    handle, self.tempf2 = tempfile.mkstemp()
194    os.close(handle)
195
196    cmde = command_executer.GetCommandExecuter()
197    for elf1 in i1.elf_files:
198      tmp_rootfs = i1.rootfs + '/'
199      f1 = elf1.replace(tmp_rootfs, '')
200      full_path1 = elf1
201      full_path2 = elf1.replace(i1.rootfs, i2.rootfs)
202
203      if full_path1 == full_path2:
204        self.logger.LogError(
205            'Error:  We\'re comparing the SAME file - {0}'.format(f1))
206        continue
207
208      command = ('objdump -d "{f1}" > {tempf1} ; '
209                 'objdump -d "{f2}" > {tempf2} ; '
210                 # Remove path string inside the dissemble
211                 'sed -i \'s!{rootfs1}!!g\' {tempf1} ; '
212                 'sed -i \'s!{rootfs2}!!g\' {tempf2} ; '
213                 'diff {tempf1} {tempf2} 1>/dev/null 2>&1').format(
214                     f1=full_path1, f2=full_path2,
215                     rootfs1=i1.rootfs, rootfs2=i2.rootfs,
216                     tempf1=self.tempf1, tempf2=self.tempf2)
217      ret = cmde.RunCommand(command, print_to_console=False)
218      if ret != 0:
219        self.logger.LogOutput('*** Not match - "{0}" "{1}"'.format(
220            full_path1, full_path2))
221        mismatch_list.append(f1)
222        if self.diff_file:
223          command = (
224              'echo "Diffs of disassemble of \"{f1}\" and \"{f2}\"" '
225              '>> {diff_file} ; diff {tempf1} {tempf2} '
226              '>> {diff_file}').format(
227                  f1=full_path1, f2=full_path2, diff_file=self.diff_file,
228                  tempf1=self.tempf1, tempf2=self.tempf2)
229          cmde.RunCommand(command, print_to_console=False)
230      else:
231        match_count += 1
232    ## End of comparing every elf files.
233
234    if not mismatch_list:
235      self.logger.LogOutput('** COOL, ALL {0} BINARIES MATCHED!! **'.format(
236          match_count))
237      return True
238
239    mismatch_str = 'Found {0} mismatch:\n'.format(len(mismatch_list))
240    for b in mismatch_list:
241      mismatch_str += '\t' + b + '\n'
242
243    self.logger.LogOutput(mismatch_str)
244    return False
245
246
247def Main(argv):
248  """The main function."""
249
250  command_executer.InitCommandExecuter()
251  images = []
252
253  parser = argparse.ArgumentParser()
254  parser.add_argument(
255      '--no_unmount', action='store_true', dest='no_unmount', default=False,
256      help='Do not unmount after finish, this is useful for debugging.')
257  parser.add_argument(
258      '--chromeos_root', dest='chromeos_root', default=None, action='store',
259      help=('[Optional] Specify a chromeos tree instead of '
260            'deducing it from image path so that we can compare '
261            '2 images that are downloaded.'))
262  parser.add_argument(
263      '--mount_basename', dest='mount_basename', default=None, action='store',
264      help=('Specify a meaningful name for the mount point. With this being '
265            'set, the mount points would be "/tmp/mount_basename.x.rootfs" '
266            ' and "/tmp/mount_basename.x.stateful". (x is 1 or 2).'))
267  parser.add_argument('--diff_file', dest='diff_file', default=None,
268                      help='Dumping all the diffs (if any) to the diff file')
269  parser.add_argument('--image1', dest='image1', default=None,
270                      required=True, help=('Image 1 file name.'))
271  parser.add_argument('--image2', dest='image2', default=None,
272                      required=True, help=('Image 2 file name.'))
273  options = parser.parse_args(argv[1:])
274
275  if options.mount_basename and options.mount_basename.find('/') >= 0:
276    logger.GetLogger().LogError(
277        '"--mount_basename" must be a name, not a path.')
278    parser.print_help()
279    return 1
280
281  result = False
282  image_comparator = None
283  try:
284    for i, image_path in enumerate([options.image1, options.image2], start=1):
285      image_path = os.path.realpath(image_path)
286      if not os.path.isfile(image_path):
287        logger.getLogger().LogError('"{0}" is not a file.'.format(image_path))
288        return 1
289
290      chromeos_root = None
291      if options.chromeos_root:
292        chromeos_root = options.chromeos_root
293      else:
294        ## Deduce chromeos root from image
295        t = image_path
296        while t != '/':
297          if misc.IsChromeOsTree(t):
298            break
299          t = os.path.dirname(t)
300        if misc.IsChromeOsTree(t):
301          chromeos_root = t
302
303      if not chromeos_root:
304        logger.GetLogger().LogError(
305            'Please provide a valid chromeos root via --chromeos_root')
306        return 1
307
308      image = CrosImage(image_path, chromeos_root, options.no_unmount)
309
310      if options.mount_basename:
311        mount_basename = '{basename}.{index}'.format(
312            basename=options.mount_basename, index=i)
313      else:
314        mount_basename = None
315
316      if image.MountImage(mount_basename):
317        images.append(image)
318        image.FindElfFiles()
319
320    if len(images) == 2:
321      image_comparator = ImageComparator(images, options.diff_file)
322      result = image_comparator.CompareImages()
323  finally:
324    for image in images:
325      image.UnmountImage()
326    if image_comparator:
327      image_comparator.Cleanup()
328
329  return 0 if result else 1
330
331
332if __name__ == '__main__':
333  Main(sys.argv)
334