lib/image_extractor.py

# Copyright (c) 2012 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.

"""Module containing class to extract the latest image for a build."""

import distutils.version
import logging
import os
import re
import shutil

import chromite.lib.cros_build_lib as chromite_build_lib


class ImageExtractor(object):
  """Class used to get the latest image for the board."""
  # Archive directory the buildbot stores images.
  LOCAL_ARCHIVE = '/b/archive'
  # The image we want to test.
  IMAGE_TO_EXTRACT = 'chromiumos_test_image.bin'
  # Archive directory in the src tree to keep latest archived image after
  # we've unzipped them.
  SRC_ARCHIVE_DIR = 'latest_image'

  def __init__(self, build_config):
    """Initializes a extractor for the build_config."""
    self.archive = os.path.join(self.LOCAL_ARCHIVE, build_config)

  def GetLatestImage(self, target_version):
    """Gets the last image archived for the board.

    Args:
      target_version:  The version that is being tested.  The archive
        directory may be being populated with the results of this version
        while we're running so we shouldn't use it as the last image archived.
    """
    logging.info('Searching for previously generated images in %s ... ',
                 self.archive)
    if os.path.exists(self.archive):
      my_re = re.compile(r'R\d+-(\d+)\.(\d+)\.(\d+).*')
      filelist = []
      target_lv = distutils.version.LooseVersion(target_version)
      for filename in os.listdir(self.archive):
        lv = distutils.version.LooseVersion(filename)
        if my_re.match(filename):
          zip_image = os.path.join(self.archive, filename, 'image.zip')
          if lv < target_lv and os.path.exists(zip_image):
            filelist.append(lv)
          elif not filename.startswith(target_version):
            logging.error('Version in archive dir is too new: %s', filename)
      if filelist:
        return os.path.join(self.archive, str(max(filelist)))

    logging.warn('Could not find a previously generated image on this host.')
    return None

  def UnzipImage(self, image_dir):
    """Unzips the image.zip from the image_dir and returns the image.

    This method unzips the image under SRC_ARCHIVE_DIR along with its version
    string. In order to save time, if it is attempting
    to re-unzip the same image with the same version string, it uses the
    cached image in SRC_ARCHIVE_DIR. It determines the version string based
    on the basename of the image_dir.

    Returns: the path to the image.bin file after it has been unzipped.
    Raises: MissingImageZipException if there is nothing to unzip within
      the image_dir.
    """
    version_string = os.path.basename(image_dir)
    cached_dir = os.path.join(ImageExtractor.SRC_ARCHIVE_DIR, version_string)
    cached_image = os.path.abspath(os.path.join(
        cached_dir, ImageExtractor.IMAGE_TO_EXTRACT))
    # If we previously unzipped the image, we're done.
    if os.path.exists(cached_image):
      logging.info('Re-using image with version %s that we previously '
                   'unzipped to %s.', version_string, cached_image)
    else:
      # Cached image for version not found. Unzipping image from archive.
      if os.path.exists(ImageExtractor.SRC_ARCHIVE_DIR):
        logging.info('Removing previously archived images from %s',
                     ImageExtractor.SRC_ARCHIVE_DIR)
        shutil.rmtree(ImageExtractor.SRC_ARCHIVE_DIR)

      os.makedirs(cached_dir)
      zip_path = os.path.join(image_dir, 'image.zip')
      logging.info('Unzipping image from %s to %s', zip_path, cached_dir)
      chromite_build_lib.RunCommand(['unzip', '-d', cached_dir, zip_path],
                                    print_cmd=False)

    return cached_image