#!/usr/bin/python2
#
# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This module generates update payloads for testing in parallel.

This module generates update payloads in parallel using the devserver. After
running this module, test payloads are generated and left in the devserver
cache. In addition, this module produces a serialized dictionary stored
with the target image that contains a mapping from the update payload name
to the path it is stored in the devserver cache.  This dictionary can then be
used by other testing scripts i.e. au_test_harness, to locate and use these
payloads for testing in virtual machines.

FOR USE OUTSIDE CHROOT ONLY.
"""

from __future__ import print_function

import functools
import optparse
import os
import pickle
import re
import shutil
import sys
import tempfile

import constants
sys.path.append(constants.CROSUTILS_LIB_DIR)
sys.path.append(constants.CROS_PLATFORM_ROOT)
sys.path.append(constants.SOURCE_ROOT)

from chromite.lib import cros_build_lib
from chromite.lib import cros_logging as logging
from chromite.lib import dev_server_wrapper
from chromite.lib import locking
from chromite.lib import osutils
from chromite.lib import parallel
from chromite.lib import path_util
from chromite.lib import sudo
from chromite.lib import timeout_util
from crostestutils.au_test_harness import cros_au_test_harness
from crostestutils.generate_test_payloads import payload_generation_exception
from crostestutils.lib import image_extractor
from crostestutils.lib import public_key_manager
from crostestutils.lib import test_helper


class InvalidDevserverOutput(Exception):
  """If we are unable to parse devserver output, this is raised."""


class UpdatePayload(object):
  """Wrapper around an update payload.

  This class defines an update payload that should be generated.  The only
  required variable to be set is |target|.  If the base image is set to None,
  this defines a full update payload to the target image.

  Variables:
    target: Payload to this image.
    base: If not None, a delta payload from this image.
    key: If set, signed payload using this private key.
    archive: If set, this payload should be archived.
    archive_stateful: If set and archive is set, archive the stateful tarball
      for the target image.
  """
  NAME_SPLITTER = '_'

  def __init__(self, target, base, key=None, archive=False,
               archive_stateful=False, for_vm=False):
    self.base = base
    self.target = target
    self.key = key
    self.archive = archive
    self.archive_stateful = archive_stateful
    self.for_vm = for_vm

  def GetNameForBin(self):
    """Returns the path we should name an archived payload."""
    real_target = os.path.realpath(self.target)
    board, target_os_version, _ = real_target.split('/')[-3:]
    prefix = 'chromeos'
    suffix = 'dev.bin'
    if self.base:
      real_base = os.path.realpath(self.base)
      base_os_version, _ = real_base.split('/')[-2:]
      name = self.NAME_SPLITTER.join([base_os_version, target_os_version, board,
                                      'delta'])
    else:
      name = self.NAME_SPLITTER.join([target_os_version, board, 'full'])

    return self.NAME_SPLITTER.join([prefix, name, suffix])

  def UpdateId(self):
    """Generates a unique update id the test harness can understand."""
    return dev_server_wrapper.GenerateUpdateId(self.target, self.base,
                                               self.key, self.for_vm)

  def __str__(self):
    my_repr = self.target
    if self.base:
      my_repr = self.base + '->' + my_repr

    if self.key:
      my_repr = my_repr + '+' + self.key

    if self.for_vm:
      my_repr = my_repr + '+' + 'for_vm'

    return my_repr

  def __eq__(self, other):
    return str(self) == str(other)

  def __hash__(self):
    return hash(str(self))


class UpdatePayloadGenerator(object):
  """Class responsible for generating update payloads."""
  CHROOT_PATH_TO_DEVSERVER_CACHE = 'var/lib/devserver/static/cache'

  def __init__(self, options):
    """Initializes a generator object from parsed options.

    Args:
      options: Parsed options from main().
    """
    self.target = options.target
    self.base = options.base
    self.target_signed = None  # Set later when creating the image.

    # For vm tests we use the _qemu name for the images.  Regardless of vm or
    # non vm, these no_vm names are guaranteed to be non-qemu base/target names.
    self.base_no_vm = self.base
    self.target_no_vm = self.target

    # Keys.
    self.public_key = options.public_key
    self.private_key = options.private_key

    # Affect what payloads we create.
    self.board = options.board
    self.basic_suite = options.basic_suite
    self.full_suite = options.full_suite
    self.payloads = set([])
    self.full_payload = options.full_payload
    self.nplus1_archive_dir = options.nplus1_archive_dir

    self.jobs = options.jobs
    self.nplus1 = options.nplus1

    self.vm = _ShouldGenerateVM(options)

  def _AddUpdatePayload(self, target, base, key=None, archive=False,
                        archive_stateful=False, for_vm=False):
    """Adds a new required update payload.  If base is None, a full payload."""
    self.payloads.add(UpdatePayload(target, base, key, archive,
                                    archive_stateful, for_vm))

  def GenerateImagesForTesting(self):
    # All vm testing requires a VM'ized target.
    if self.vm:
      self.target = test_helper.CreateVMImage(self.target, self.board)

    if self.full_suite:
      if self.public_key:
        self.target_signed = self.target + '.signed'
        if not os.path.exists(self.target_signed):
          logging.info('Creating a signed image for signed payload test.')
          shutil.copy(self.target, self.target_signed)

        public_key_manager.PublicKeyManager(self.target_signed,
                                            self.public_key).AddKeyToImage()

      # The full suite may not have a VM image produced for the test image yet.
      # Ensure this is created.
      self.base = test_helper.CreateVMImage(self.base, self.board)

  def GeneratePayloadRequirements(self):
    """Generate Payload Requirements for AUTestHarness and NPlus1 Testing."""
    if self.full_suite:
      # N-1->N.
      self._AddUpdatePayload(self.target, self.base, for_vm=self.vm)

      # N->N after N-1->N.
      self._AddUpdatePayload(self.target, self.target, for_vm=self.vm)

      # N->N From VM base.
      self._AddUpdatePayload(self.target, self.target, for_vm=self.vm)

      # Need a signed payload for the signed payload test.
      if self.target_signed:
        self._AddUpdatePayload(self.target_signed, self.target_signed,
                               self.private_key, for_vm=self.vm)

    if self.basic_suite:
      # Update image to itself from VM base.
      self._AddUpdatePayload(self.target, self.target, for_vm=self.vm)

    # Add deltas for m minus 1 to n and n to n.
    if self.nplus1:
      self._AddUpdatePayload(self.target_no_vm, self.base_no_vm, archive=True)
      self._AddUpdatePayload(self.target_no_vm, self.target_no_vm, archive=True)

    # Add the full payload.
    if self.nplus1 or self.full_payload:
      self._AddUpdatePayload(self.target_no_vm, None, archive=True,
                             archive_stateful=True)


  def GeneratePayloads(self):
    """Iterates through payload requirements and generates them.

    This is the main method of this class.  It iterates through payloads
    it needs, generates them, and builds a Cache that can be used by the
    test harness to reference these payloads.

    Returns:
      The cache as a Python dict.
    """

    def GeneratePayload(payload, log_file):
      """Returns the error code from generating an update with the devserver."""
      # Base command.
      command = ['start_devserver', '--pregenerate_update', '--exit']

      in_chroot_key = in_chroot_base = None
      in_chroot_target = path_util.ToChrootPath(payload.target)
      if payload.base:
        in_chroot_base = path_util.ToChrootPath(payload.base)

      if payload.key:
        in_chroot_key = path_util.ToChrootPath(payload.key)

      command.append('--image=%s' % in_chroot_target)
      if payload.base:
        command.append('--src_image=%s' % in_chroot_base)
      if payload.key:
        command.append('--private_key=%s' % in_chroot_key)

      if payload.base:
        debug_message = 'delta payload from %s to %s' % (payload.base,
                                                         payload.target)
      else:
        debug_message = 'full payload to %s' % payload.target

      if payload.for_vm:
        debug_message += ' and not patching the kernel.'

      if in_chroot_key:
        debug_message = 'Generating a signed %s' % debug_message
      else:
        debug_message = 'Generating an unsigned %s' % debug_message

      logging.info(debug_message)
      try:
        with timeout_util.Timeout(constants.MAX_TIMEOUT_SECONDS):
          cros_build_lib.RunCommand(command, log_stdout_to_file=log_file,
                                    combine_stdout_stderr=True,
                                    enter_chroot=True, print_cmd=False,
                                    cwd=constants.SOURCE_ROOT)
      except (timeout_util.TimeoutError, cros_build_lib.RunCommandError):
        # Print output first, then re-raise the exception.
        if os.path.isfile(log_file):
          logging.error(osutils.ReadFile(log_file))
        raise

    def ProcessOutput(log_files):
      """Processes results from the log files of GeneratePayload invocations.

      Args:
        log_files: A list of filename strings with stored logs.

      Returns:
        An array of cache entries from the log files.

      Raises:
        payload_generation_exception.PayloadGenerationException: Raises this
          exception if we failed to parse the devserver output to find the
          location of the update path.
      """
      # Looking for this line in the output.
      key_line_re = re.compile(r'^PREGENERATED_UPDATE=([\w/./+]+)')
      return_array = []
      for log_file in log_files:
        with open(log_file) as f:
          for line in f:
            match = key_line_re.search(line)
            if match:
              # Convert cache/label/update.gz -> update/cache/label.
              path_to_update_gz = match.group(1).rstrip()
              path_to_update_dir = path_to_update_gz.rpartition(
                  '/update.gz')[0]

              # Check that we could actually parse the directory correctly.
              if not path_to_update_dir:
                raise payload_generation_exception.PayloadGenerationException(
                    'Payload generated but failed to parse cache directory.')

              return_array.append('/'.join(['update', path_to_update_dir]))
              break
          else:
            logging.error('Could not find PREGENERATED_UPDATE in log:')
            f.seek(0)
            for line in f:
              logging.error('  log: %s', line)
            # This is not a recoverable error.
            raise InvalidDevserverOutput('Could not parse devserver log')

      return return_array

    jobs = []
    log_files = []
    # Generate list of paylods and list of log files.
    for payload in self.payloads:
      fd, log_file = tempfile.mkstemp('GenerateVMUpdate')
      os.close(fd)  # Just want filename so close file immediately.

      jobs.append(functools.partial(GeneratePayload, payload, log_file))
      log_files.append(log_file)

    # Run update generation code and wait for output.
    logging.info('Generating updates required for this test suite in parallel.')
    try:
      parallel.RunParallelSteps(jobs, max_parallel=self.jobs)
    except parallel.BackgroundFailure as ex:
      logging.error(ex)
      raise payload_generation_exception.PayloadGenerationException(
          'Failed to generate a required update.')

    results = ProcessOutput(log_files)

    # Build the dictionary from our id's and returned cache paths.
    cache_dictionary = {}
    for index, payload in enumerate(self.payloads):
      # Path return is of the form update/cache/directory.
      update_path = results[index]
      cache_dictionary[payload.UpdateId()] = update_path
      # Archive payload to payload directory.
      if payload.archive and self.nplus1_archive_dir:
        # Only need directory as we know the rest.
        path_to_payload_dir = os.path.join(
            constants.SOURCE_ROOT, 'chroot',
            self.CHROOT_PATH_TO_DEVSERVER_CACHE, os.path.basename(update_path))
        payload_path = os.path.join(path_to_payload_dir, 'update.gz')
        archive_path = os.path.join(self.nplus1_archive_dir,
                                    payload.GetNameForBin())
        logging.info('Archiving %s to %s.', payload.GetNameForBin(),
                     archive_path)
        shutil.copyfile(payload_path, archive_path)
        if payload.archive_stateful:
          stateful_path = os.path.join(path_to_payload_dir, 'stateful.tgz')
          archive_path = os.path.join(self.nplus1_archive_dir, 'stateful.tgz')
          logging.info('Archiving stateful payload from %s to %s',
                       payload.GetNameForBin(), archive_path)
          shutil.copyfile(stateful_path, archive_path)

    return cache_dictionary

  def DumpCacheToDisk(self, cache):
    """Dumps the cache to the same folder as the images."""
    if not self.basic_suite and not self.full_suite:
      logging.info('Not dumping payload cache to disk as payloads for the '
                   'test harness were not requested.')
    else:
      path_to_dump = os.path.dirname(self.target)
      cache_file = os.path.join(path_to_dump, cros_au_test_harness.CACHE_FILE)

      logging.info('Dumping %s', cache_file)
      with open(cache_file, 'w') as file_handle:
        pickle.dump(cache, file_handle)


def _ShouldGenerateVM(options):
  """Returns true if we will need a VM version of our images."""
  # This is a combination of options.vm and whether or not we are generating
  # payloads for vm testing.
  return options.vm and (options.basic_suite or options.full_suite)


def CheckOptions(parser, options):
  """Checks that given options are valid.

  Args:
    parser: Parser used to parse options.
    options: Parse options from OptionParser.
  """
  if not options.target or not os.path.isfile(options.target):
    parser.error('Target image must exist.')

  # Determine the base image. If latest_from_config specified, find the latest
  # image from the given config. If it doesn't exist, use the target image.
  target_version = os.path.realpath(options.target).rsplit('/', 2)[-2]
  if options.base_latest_from_dir:
    # Extract the latest build.
    extractor = image_extractor.ImageExtractor(options.base_latest_from_dir,
                                               os.path.basename(options.target))
    latest_image_dir = extractor.GetLatestImage(target_version)
    if latest_image_dir:
      options.base = extractor.UnzipImage(latest_image_dir)
    else:
      logging.warning('No previous image.zip found in local archive.')

  if not options.base:
    logging.info('Using target image as base image.')
    options.base = options.target

  if not os.path.isfile(options.base):
    parser.error('Base image must exist.')

  if options.private_key:
    if not os.path.isfile(options.private_key):
      parser.error('Private key must exist.')

    if not os.path.isfile(options.public_key):
      parser.error('Public key must exist.')

  if _ShouldGenerateVM(options):
    if not options.board:
      parser.error('Board must be set to generate update '
                   'payloads for vm.')

  if options.full_payload or options.nplus1:
    if not options.nplus1_archive_dir:
      parser.error('Must specify an archive directory if nplus1 or '
                   'full payload are specified.')


def main():
  test_helper.SetupCommonLoggingFormat()
  parser = optparse.OptionParser()

  # Options related to which payloads to generate.
  parser.add_option('--basic_suite', default=False, action='store_true',
                    help='Prepare to run the basic au test suite.')
  parser.add_option('--full_suite', default=False, action='store_true',
                    help='Prepare to run the full au test suite.')
  parser.add_option('--full_payload', default=False, action='store_true',
                    help='Generate the full update payload and store it in '
                    'the nplus1 archive dir.')
  parser.add_option('--nplus1', default=False, action='store_true',
                    help='Produce nplus1 updates for testing in lab and store '
                    'them in the nplus1 archive dir.')
  parser.add_option('--nplus1_archive_dir', default=None,
                    help='Archive nplus1 updates into this directory.')

  # Options related to how to generate test payloads for the test harness.
  parser.add_option('--novm', default=True, action='store_false', dest='vm',
                    help='Test Harness payloads will not be tested in a VM.')
  parser.add_option('--private_key',
                    help='Private key to sign payloads for test harness.')
  parser.add_option('--public_key',
                    help='Public key to verify payloads for test harness.')

  # Options related to the images to test.
  parser.add_option('--board', help='Board used for the images.')
  parser.add_option('--base', help='Image we want to test updates from.')
  parser.add_option('--base_latest_from_dir', help='Ignore the base '
                    'option and use the latest image from the specified '
                    'directory as the base image. If none exists, default to '
                    'target image.')
  parser.add_option('--target', help='Image we want to test updates to.')

  # Miscellaneous options.
  parser.add_option('--jobs', default=test_helper.CalculateDefaultJobs(),
                    type=int,
                    help='Number of payloads to generate in parallel.')

  options = parser.parse_args()[0]
  CheckOptions(parser, options)
  if options.nplus1_archive_dir and not os.path.exists(
      options.nplus1_archive_dir):
    os.makedirs(options.nplus1_archive_dir)

  # Don't allow this code to be run more than once at a time.
  lock_path = os.path.join(os.path.dirname(__file__), '.lock_file')
  with locking.FileLock(lock_path, 'generate payloads lock') as lock:
    lock.write_lock()
    with sudo.SudoKeepAlive():
      generator = UpdatePayloadGenerator(options)
      generator.GenerateImagesForTesting()
      generator.GeneratePayloadRequirements()
      cache = generator.GeneratePayloads()
      generator.DumpCacheToDisk(cache)


if __name__ == '__main__':
  main()