download_images.py revision e627fd61c2edba668eb2af8221892286b13f05a3 
1#!/usr/bin/python
2
3# Copyright (c) 2014, 2015 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7import ast
8import os
9
10import test_flag
11
12from utils import command_executer
13
14class MissingImage(Exception):
15  """Raised when the requested image does not exist in gs://"""
16
17class ImageDownloader(object):
18
19  def __init__(self, logger_to_use=None, log_level="verbose",
20               cmd_exec=None):
21    self._logger = logger_to_use
22    self.log_level = log_level
23    self._ce = cmd_exec or command_executer.GetCommandExecuter(self._logger,
24                                                   log_level = self.log_level)
25
26  def _GetBuildID (self, chromeos_root, xbuddy_label):
27    # Get the translation of the xbuddy_label into the real Google Storage
28    # image name.
29    command = ("cd ~/trunk/src/third_party/toolchain-utils/crosperf; "
30               "python translate_xbuddy.py '%s'" % xbuddy_label)
31    retval, build_id_tuple_str, _ = self._ce.ChrootRunCommand(chromeos_root,
32                                                          command, True)
33    if not build_id_tuple_str:
34      raise MissingImage ("Unable to find image for '%s'" % xbuddy_label)
35
36    build_id_tuple = ast.literal_eval(build_id_tuple_str)
37    build_id = build_id_tuple[0]
38
39    return build_id
40
41  def _DownloadImage(self, chromeos_root, build_id, image_name):
42    if self.log_level == "average":
43      self._logger.LogOutput ("Preparing to download %s image to local "
44                              "directory." % build_id)
45
46    # Make sure the directory for downloading the image exists.
47    download_path = os.path.join(chromeos_root, "chroot/tmp",
48                                 build_id)
49    image_path = os.path.join(download_path, "chromiumos_test_image.bin")
50    if not os.path.exists(download_path):
51      os.makedirs(download_path)
52
53    # Check to see if the image has already been downloaded.  If not,
54    # download the image.
55    status = 0
56    if not os.path.exists(image_path):
57      command = "gsutil cp %s /tmp/%s" % (image_name, build_id)
58
59      if self.log_level != "verbose":
60        self._logger.LogOutput ("CMD: %s" % command)
61      status = self._ce.ChrootRunCommand(chromeos_root, command)
62
63    if status == 0:
64      return image_path
65    else:
66      return None
67
68  def _UncompressImage(self, chromeos_root, build_id):
69    # Check to see if the file has already been uncompresssed, etc.
70    if os.path.exists(os.path.join(chromeos_root, "chroot/tmp", build_id,
71                                   "chromiumos_test_image.bin")):
72      return 0
73
74    # Uncompress and untar the downloaded image.
75    command = ("cd /tmp/%s ;unxz chromiumos_test_image.tar.xz; "
76               "tar -xvf chromiumos_test_image.tar" % build_id)
77    if self.log_level != "verbose":
78      self._logger.LogOutput("CMD: %s" % command)
79      print("(Uncompressing and un-tarring may take a couple of minutes..."
80            "please be patient.)")
81    retval = self._ce.ChrootRunCommand(chromeos_root, command)
82    return retval
83
84
85  def Run(self, chromeos_root, xbuddy_label):
86    build_id = self._GetBuildID(chromeos_root, xbuddy_label)
87    image_name = ("gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz"
88                  % build_id)
89
90    # Verify that image exists for build_id, before attempting to
91    # download it.
92    status = 0
93    if not test_flag.GetTestMode():
94      cmd = "gsutil ls %s" % image_name
95      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
96    if status != 0:
97      raise MissingImage("Cannot find official image: %s." % image_name)
98    image_path = self._DownloadImage(chromeos_root, build_id, image_name)
99    retval = 0
100    if image_path:
101      retval = self._UncompressImage(chromeos_root, build_id)
102    else:
103      retval = 1
104
105    if retval == 0 and self.log_level != "quiet":
106      self._logger.LogOutput("Using image from %s." % image_path)
107
108    return retval, image_path
109