#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2019 The ChromiumOS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Diff 2 chromiumos images by comparing each elf file.

   The script diffs every *ELF* files by dissembling every *executable*
   section, which means it is not a FULL elf differ.

   A simple usage example -
     chromiumos_image_diff.py --image1 image-path-1 --image2 image-path-2

   Note that image path should be inside the chroot, if not (ie, image is
   downloaded from web), please specify a chromiumos checkout via
   "--chromeos_root".

   And this script should be executed outside chroot.
"""

from __future__ import print_function

__author__ = 'shenhan@google.com (Han Shen)'

import argparse
import os
import re
import sys
import tempfile

import image_chromeos
from cros_utils import command_executer
from cros_utils import logger
from cros_utils import misc


class CrosImage(object):
  """A cros image object."""

  def __init__(self, image, chromeos_root, no_unmount):
    self.image = image
    self.chromeos_root = chromeos_root
    self.mounted = False
    self._ce = command_executer.GetCommandExecuter()
    self.logger = logger.GetLogger()
    self.elf_files = []
    self.no_unmount = no_unmount
    self.unmount_script = ''
    self.stateful = ''
    self.rootfs = ''

  def MountImage(self, mount_basename):
    """Mount/unpack the image."""

    if mount_basename:
      self.rootfs = '/tmp/{0}.rootfs'.format(mount_basename)
      self.stateful = '/tmp/{0}.stateful'.format(mount_basename)
      self.unmount_script = '/tmp/{0}.unmount.sh'.format(mount_basename)
    else:
      self.rootfs = tempfile.mkdtemp(
          suffix='.rootfs', prefix='chromiumos_image_diff')
      ## rootfs is like /tmp/tmpxyz012.rootfs.
      match = re.match(r'^(.*)\.rootfs$', self.rootfs)
      basename = match.group(1)
      self.stateful = basename + '.stateful'
      os.mkdir(self.stateful)
      self.unmount_script = '{0}.unmount.sh'.format(basename)

    self.logger.LogOutput('Mounting "{0}" onto "{1}" and "{2}"'.format(
        self.image, self.rootfs, self.stateful))
    ## First of all creating an unmount image
    self.CreateUnmountScript()
    command = image_chromeos.GetImageMountCommand(self.image, self.rootfs,
                                                  self.stateful)
    rv = self._ce.RunCommand(command, print_to_console=True)
    self.mounted = (rv == 0)
    if not self.mounted:
      self.logger.LogError('Failed to mount "{0}" onto "{1}" and "{2}".'.format(
          self.image, self.rootfs, self.stateful))
    return self.mounted

  def CreateUnmountScript(self):
    command = ('sudo umount {r}/usr/local {r}/usr/share/oem '
               '{r}/var {r}/mnt/stateful_partition {r}; sudo umount {s} ; '
               'rmdir {r} ; rmdir {s}\n').format(
                   r=self.rootfs, s=self.stateful)
    f = open(self.unmount_script, 'w', encoding='utf-8')
    f.write(command)
    f.close()
    self._ce.RunCommand(
        'chmod +x {}'.format(self.unmount_script), print_to_console=False)
    self.logger.LogOutput('Created an unmount script - "{0}"'.format(
        self.unmount_script))

  def UnmountImage(self):
    """Unmount the image and delete mount point."""

    self.logger.LogOutput('Unmounting image "{0}" from "{1}" and "{2}"'.format(
        self.image, self.rootfs, self.stateful))
    if self.mounted:
      command = 'bash "{0}"'.format(self.unmount_script)
      if self.no_unmount:
        self.logger.LogOutput(('Please unmount manually - \n'
                               '\t bash "{0}"'.format(self.unmount_script)))
      else:
        if self._ce.RunCommand(command, print_to_console=True) == 0:
          self._ce.RunCommand('rm {0}'.format(self.unmount_script))
          self.mounted = False
          self.rootfs = None
          self.stateful = None
          self.unmount_script = None

    return not self.mounted

  def FindElfFiles(self):
    """Find all elf files for the image.

    Returns:
      Always true
    """

    self.logger.LogOutput('Finding all elf files in "{0}" ...'.format(
        self.rootfs))
    # Note '\;' must be prefixed by 'r'.
    command = ('find "{0}" -type f -exec '
               'bash -c \'file -b "{{}}" | grep -q "ELF"\''
               r' \; '
               r'-exec echo "{{}}" \;').format(self.rootfs)
    self.logger.LogCmd(command)
    _, out, _ = self._ce.RunCommandWOutput(command, print_to_console=False)
    self.elf_files = out.splitlines()
    self.logger.LogOutput('Total {0} elf files found.'.format(
        len(self.elf_files)))
    return True


class ImageComparator(object):
  """A class that wraps comparsion actions."""

  def __init__(self, images, diff_file):
    self.images = images
    self.logger = logger.GetLogger()
    self.diff_file = diff_file
    self.tempf1 = None
    self.tempf2 = None

  def Cleanup(self):
    if self.tempf1 and self.tempf2:
      command_executer.GetCommandExecuter().RunCommand('rm {0} {1}'.format(
          self.tempf1, self.tempf2))
      logger.GetLogger('Removed "{0}" and "{1}".'.format(
          self.tempf1, self.tempf2))

  def CheckElfFileSetEquality(self):
    """Checking whether images have exactly number of elf files."""

    self.logger.LogOutput('Checking elf file equality ...')
    i1 = self.images[0]
    i2 = self.images[1]
    t1 = i1.rootfs + '/'
    elfset1 = {e.replace(t1, '') for e in i1.elf_files}
    t2 = i2.rootfs + '/'
    elfset2 = {e.replace(t2, '') for e in i2.elf_files}
    dif1 = elfset1.difference(elfset2)
    msg = None
    if dif1:
      msg = 'The following files are not in "{image}" - "{rootfs}":\n'.format(
          image=i2.image, rootfs=i2.rootfs)
      for d in dif1:
        msg += '\t' + d + '\n'
    dif2 = elfset2.difference(elfset1)
    if dif2:
      msg = 'The following files are not in "{image}" - "{rootfs}":\n'.format(
          image=i1.image, rootfs=i1.rootfs)
      for d in dif2:
        msg += '\t' + d + '\n'
    if msg:
      self.logger.LogError(msg)
      return False
    return True

  def CompareImages(self):
    """Do the comparsion work."""

    if not self.CheckElfFileSetEquality():
      return False

    mismatch_list = []
    match_count = 0
    i1 = self.images[0]
    i2 = self.images[1]
    self.logger.LogOutput('Start comparing {0} elf file by file ...'.format(
        len(i1.elf_files)))
    ## Note - i1.elf_files and i2.elf_files have exactly the same entries here.

    ## Create 2 temp files to be used for all disassembed files.
    handle, self.tempf1 = tempfile.mkstemp()
    os.close(handle)  # We do not need the handle
    handle, self.tempf2 = tempfile.mkstemp()
    os.close(handle)

    cmde = command_executer.GetCommandExecuter()
    for elf1 in i1.elf_files:
      tmp_rootfs = i1.rootfs + '/'
      f1 = elf1.replace(tmp_rootfs, '')
      full_path1 = elf1
      full_path2 = elf1.replace(i1.rootfs, i2.rootfs)

      if full_path1 == full_path2:
        self.logger.LogError(
            "Error:  We're comparing the SAME file - {0}".format(f1))
        continue

      command = (
          'objdump -d "{f1}" > {tempf1} ; '
          'objdump -d "{f2}" > {tempf2} ; '
          # Remove path string inside the dissemble
          "sed -i 's!{rootfs1}!!g' {tempf1} ; "
          "sed -i 's!{rootfs2}!!g' {tempf2} ; "
          'diff {tempf1} {tempf2} 1>/dev/null 2>&1').format(
              f1=full_path1,
              f2=full_path2,
              rootfs1=i1.rootfs,
              rootfs2=i2.rootfs,
              tempf1=self.tempf1,
              tempf2=self.tempf2)
      ret = cmde.RunCommand(command, print_to_console=False)
      if ret != 0:
        self.logger.LogOutput('*** Not match - "{0}" "{1}"'.format(
            full_path1, full_path2))
        mismatch_list.append(f1)
        if self.diff_file:
          command = ('echo "Diffs of disassemble of \"{f1}\" and \"{f2}\"" '
                     '>> {diff_file} ; diff {tempf1} {tempf2} '
                     '>> {diff_file}').format(
                         f1=full_path1,
                         f2=full_path2,
                         diff_file=self.diff_file,
                         tempf1=self.tempf1,
                         tempf2=self.tempf2)
          cmde.RunCommand(command, print_to_console=False)
      else:
        match_count += 1
    ## End of comparing every elf files.

    if not mismatch_list:
      self.logger.LogOutput(
          '** COOL, ALL {0} BINARIES MATCHED!! **'.format(match_count))
      return True

    mismatch_str = 'Found {0} mismatch:\n'.format(len(mismatch_list))
    for b in mismatch_list:
      mismatch_str += '\t' + b + '\n'

    self.logger.LogOutput(mismatch_str)
    return False


def Main(argv):
  """The main function."""

  command_executer.InitCommandExecuter()
  images = []

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--no_unmount',
      action='store_true',
      dest='no_unmount',
      default=False,
      help='Do not unmount after finish, this is useful for debugging.')
  parser.add_argument(
      '--chromeos_root',
      dest='chromeos_root',
      default=None,
      action='store',
      help=('[Optional] Specify a chromeos tree instead of '
            'deducing it from image path so that we can compare '
            '2 images that are downloaded.'))
  parser.add_argument(
      '--mount_basename',
      dest='mount_basename',
      default=None,
      action='store',
      help=('Specify a meaningful name for the mount point. With this being '
            'set, the mount points would be "/tmp/mount_basename.x.rootfs" '
            ' and "/tmp/mount_basename.x.stateful". (x is 1 or 2).'))
  parser.add_argument(
      '--diff_file',
      dest='diff_file',
      default=None,
      help='Dumping all the diffs (if any) to the diff file')
  parser.add_argument(
      '--image1',
      dest='image1',
      default=None,
      required=True,
      help=('Image 1 file name.'))
  parser.add_argument(
      '--image2',
      dest='image2',
      default=None,
      required=True,
      help=('Image 2 file name.'))
  options = parser.parse_args(argv[1:])

  if options.mount_basename and options.mount_basename.find('/') >= 0:
    logger.GetLogger().LogError(
        '"--mount_basename" must be a name, not a path.')
    parser.print_help()
    return 1

  result = False
  image_comparator = None
  try:
    for i, image_path in enumerate([options.image1, options.image2], start=1):
      image_path = os.path.realpath(image_path)
      if not os.path.isfile(image_path):
        logger.GetLogger().LogError('"{0}" is not a file.'.format(image_path))
        return 1

      chromeos_root = None
      if options.chromeos_root:
        chromeos_root = options.chromeos_root
      else:
        ## Deduce chromeos root from image
        t = image_path
        while t != '/':
          if misc.IsChromeOsTree(t):
            break
          t = os.path.dirname(t)
        if misc.IsChromeOsTree(t):
          chromeos_root = t

      if not chromeos_root:
        logger.GetLogger().LogError(
            'Please provide a valid chromeos root via --chromeos_root')
        return 1

      image = CrosImage(image_path, chromeos_root, options.no_unmount)

      if options.mount_basename:
        mount_basename = '{basename}.{index}'.format(
            basename=options.mount_basename, index=i)
      else:
        mount_basename = None

      if image.MountImage(mount_basename):
        images.append(image)
        image.FindElfFiles()

    if len(images) == 2:
      image_comparator = ImageComparator(images, options.diff_file)
      result = image_comparator.CompareImages()
  finally:
    for image in images:
      image.UnmountImage()
    if image_comparator:
      image_comparator.Cleanup()

  return 0 if result else 1


if __name__ == '__main__':
  Main(sys.argv)