download_images.py revision 6ffbb92f0913f990b6a5c7ef706aee9bfb4faa9e 
1#!/usr/bin/python
2
3# Copyright (c) 2014 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
10from utils import command_executer
11
12class ImageDownloader(object):
13
14  def __init__(self, logger_to_use=None, log_level="verbose",
15               cmd_exec=None):
16    self._logger = logger_to_use
17    self.log_level = log_level
18    self._ce = cmd_exec or command_executer.GetCommandExecuter(self._logger,
19                                                   log_level = self.log_level)
20
21  def _GetBuildID (self, chromeos_root, xbuddy_label):
22    # Get the translation of the xbuddy_label into the real Google Storage
23    # image name.
24    command = ("cd ~/trunk/src/third_party/toolchain-utils/crosperf; "
25               "python translate_xbuddy.py '%s'" % xbuddy_label)
26    retval, build_id_tuple_str, _ = self._ce.ChrootRunCommand(chromeos_root,
27                                                          command, True)
28    build_id_tuple = ast.literal_eval(build_id_tuple_str)
29    build_id = build_id_tuple[0]
30
31    return build_id
32
33  def _DownloadImage(self, chromeos_root, build_id):
34    if self.log_level == "average":
35      self._logger.LogOutput ("Preparing to download %s image to local directory." % build_id)
36
37    # Make sure the directory for downloading the image exists.
38    download_path = os.path.join(chromeos_root, "chroot/tmp",
39                                 build_id)
40    image_path = os.path.join(download_path, "chromiumos_test_image.bin")
41    if not os.path.exists(download_path):
42      os.makedirs(download_path)
43
44    # Check to see if the image has already been downloaded.  If not,
45    # download the image.
46    status = 0
47    if not os.path.exists(image_path):
48      command = ("gsutil cp gs://chromeos-image-archive/%s"
49                 "/chromiumos_test_image.tar.xz /tmp/%s" % (build_id,
50                                                            build_id))
51
52      if self.log_level != "verbose":
53        self._logger.LogOutput ("CMD: %s" % command)
54      status = self._ce.ChrootRunCommand(chromeos_root, command)
55
56    if status == 0:
57      return image_path
58    else:
59      return None
60
61  def _UncompressImage(self, chromeos_root, build_id):
62    # Check to see if the file has already been uncompresssed, etc.
63    if os.path.exists(os.path.join(chromeos_root, "chroot/tmp", build_id,
64                                   "chromiumos_test_image.bin")):
65      return 0
66
67    # Uncompress and untar the downloaded image.
68    command = ("cd /tmp/%s ;unxz chromiumos_test_image.tar.xz; "
69               "tar -xvf chromiumos_test_image.tar" % build_id)
70    if self.log_level != "verbose":
71      self._logger.LogOutput("CMD: %s" % command)
72      print("(Uncompressing and un-tarring may take a couple of minutes..."
73            "please be patient.)")
74    retval = self._ce.ChrootRunCommand(chromeos_root, command)
75    return retval
76
77
78  def Run(self, chromeos_root, xbuddy_label):
79    build_id = self._GetBuildID(chromeos_root, xbuddy_label)
80
81
82    retval = 0
83    image_path = self._DownloadImage(chromeos_root, build_id)
84    if image_path:
85      retval = self._UncompressImage(chromeos_root, build_id)
86    else:
87      retval = 1
88
89    if retval == 0 and self.log_level != "quiet":
90      self._logger.LogOutput("Using image from %s." % image_path)
91
92    return retval, image_path
93