#!/usr/bin/python
#
# 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.
"""

import logging
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)
import cros_build_lib as cros_lib

from crostestutils.au_test_harness import cros_au_test_harness
from crostestutils.generate_test_payloads import payload_generation_exception
from crostestutils.lib import dev_server_wrapper
from crostestutils.lib import parallel_test_job
from crostestutils.lib import public_key_manager
from crostestutils.lib import test_helper


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):
    self.base = base
    self.target = target
    self.key = key
    self.archive = archive
    self.archive_stateful = archive_stateful

  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)

  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

    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.full_suite = options.full_suite
    self.payloads = set([])
    self.nplus1_archive_dir = options.nplus1_archive_dir

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

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

  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()

      self.base = test_helper.CreateVMImage(self.base, self.board)

  def GeneratePayloadRequirements(self):
    """Generate Payload Requirements for AUTestHarness and NPlus1 Testing."""

    def AddPayloadsForAUTestHarness():
      if self.full_suite:
        self._AddUpdatePayload(self.target_no_vm, self.base)
        self._AddUpdatePayload(self.base_no_vm, self.target_no_vm)
        self._AddUpdatePayload(self.target_no_vm, self.target)

        # Need a signed payload for the signed payload test.
        if self.target_signed:
          self._AddUpdatePayload(self.target_no_vm, self.target_signed,
                                 self.private_key)
      else:
        self._AddUpdatePayload(self.target_no_vm, self.target)

    def AddNPlus1Updates():
      self._AddUpdatePayload(self.target_no_vm, self.base_no_vm, archive=True)
      self._AddUpdatePayload(self.target_no_vm, self.target_no_vm, archive=True)
      self._AddUpdatePayload(self.target_no_vm, None, archive=True,
                             archive_stateful=True)

    if self.nplus1: AddNPlus1Updates()
    AddPayloadsForAUTestHarness()

  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 = ['sudo', 'start_devserver', '--pregenerate_update', '--exit']

      in_chroot_key = None
      in_chroot_base = None
      in_chroot_target = cros_lib.ReinterpretPathForChroot(payload.target)
      if payload.base:
        in_chroot_base = cros_lib.ReinterpretPathForChroot(payload.base)

      if payload.key:
        in_chroot_key = cros_lib.ReinterpretPathForChroot(payload.key)

      command.append('--image=%s' % in_chroot_target)
      if payload.base: command.append('--src_image=%s' % in_chroot_base)
      if self.vm: command.append('--for_vm')
      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 in_chroot_key:
        debug_message = 'Generating a signed %s' % debug_message
      else:
        debug_message = 'Generating an unsigned %s' % debug_message

      logging.info(debug_message)
      return cros_lib.RunCommand(command, enter_chroot=True, print_cmd=False,
                                 cwd=cros_lib.GetCrosUtilsPath(),
                                 log_to_file=log_file, error_ok=True,
                                 exit_code=True)

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

      Args:
        log_files:  A list of filename strings with stored logs.
        return_codes: An equally sized list of return codes pertaining to
          the processes that ran that generated those logs.
      Returns:
        An array of cache entries from the log files.
      Raises:
        payload_generation_exception.PayloadGenerationException: Returns this
          exception if the devserver returned an error code when producing
          an update payload OR 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('^PREGENERATED_UPDATE=([\w/./+]+)')
      return_array = []
      for log_file, return_code in zip(log_files, return_codes):
        with open(log_file) as log_file_handle:
          output = log_file_handle.read()

        if return_code != 0:
          logging.error(output)
          raise payload_generation_exception.PayloadGenerationException(
              'Failed to generate a required update.')
        else:
          for line in output.splitlines():
            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.000
              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 update string in log.')

      return return_array

    payloads = []
    jobs = []
    args = []
    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.

      payloads.append(payload)
      jobs.append(GeneratePayload)
      args.append((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.')
    error_codes = parallel_test_job.RunParallelJobs(self.jobs, jobs, args)
    results = ProcessOutput(log_files, error_codes)

    # Build the dictionary from our id's and returned cache paths.
    cache_dictionary = {}
    for index, payload in enumerate(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."""
    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 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.')

  if not options.base:
    logging.info('Base image not specified. Using target 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 options.vm:
    if not options.board:
      parser.error('Board must be set to generate update '
                   'payloads for vm.')


def main():
  test_helper.SetupCommonLoggingFormat()
  parser = optparse.OptionParser()
  parser.add_option('--base', help='Image we want to test updates from.')
  parser.add_option('--board', help='Board used for the images.')
  parser.add_option('--clean', default=False, dest='clean', action='store_true',
                    help='Clean cache of previous payloads')
  parser.add_option('--full_suite', default=False, action='store_true',
                    help='Prepare to run the full au test suite.')
  parser.add_option('--jobs', default=test_helper.CalculateDefaultJobs(),
                    type=int,
                    help='Number of payloads to generate in parallel.')
  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('--nplus1', default=False, action='store_true',
                    help='Produce nplus1 updates for testing in lab.')
  parser.add_option('--nplus1_archive_dir', default=None,
                    help='If set, archive nplu1 updates in this directory.')
  parser.add_option('--private_key',
                    help='Private key to sign payloads.')
  parser.add_option('--public_key',
                    help='Public key to verify signed payloads.')
  parser.add_option('--target', help='Image we want to test updates to.')
  options = parser.parse_args()[0]
  CheckOptions(parser, options)

  if options.clean:
    dev_server_wrapper.DevServerWrapper.WipePayloadCache()

  if options.nplus1_archive_dir and not os.path.exists(
      options.nplus1_archive_dir):
    os.makedirs(options.nplus1_archive_dir)

  generator = UpdatePayloadGenerator(options)
  generator.GenerateImagesForTesting()
  generator.GeneratePayloadRequirements()
  cache = generator.GeneratePayloads()
  generator.DumpCacheToDisk(cache)


if __name__ == '__main__':
  main()