#!/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.

"""GenerateTestPayloads 'quickly' creates payloads needed for testing."""

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):
  """Defines an update payload."""

  def __init__(self, target, base, key=None):
    self.base = base
    self.target = target
    self.key = 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."""

  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 = options.nplus1
    self.vm = options.vm

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

  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)
      self._AddUpdatePayload(self.target_no_vm, self.target_no_vm)
      self._AddUpdatePayload(self.target_no_vm, None)

    AddPayloadsForAUTestHarness()
    if self.nplus1: AddNPlus1Updates()

  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)

      debug_message = 'delta payload from %s to %s' % (payload.base,
                                                       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.
      """
      # 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 blah/blah/update.gz -> update/blah/blah.
              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

    update_ids = []
    jobs = []
    args = []
    log_files = []
    for payload in self.payloads:
      fd, log_file = tempfile.mkstemp('GenerateVMUpdate')
      os.close(fd)  # Just want filename so close file immediately.

      update_id = dev_server_wrapper.GenerateUpdateId(
          payload.target, payload.base, payload.key)
      update_ids.append(update_id)
      jobs.append(GeneratePayload)
      args.append((payload, log_file))
      log_files.append(log_file)

    logging.info('Generating updates required for this test suite in parallel.')
    error_codes = parallel_test_job.RunParallelJobs(
        test_helper.CalculateDefaultJobs(), jobs, args)
    results = ProcessOutput(log_files, error_codes)

    # Build the dictionary from our id's and returned cache paths.
    cache_dictionary = {}
    for index, update_id in enumerate(update_ids):
      cache_dictionary[update_id] = results[index]

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

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


if __name__ == '__main__':
  main()