download_images.py revision f2a3ef46f75d2196a93d3ed27f4d1fcf22b54fbe 
1
2# Copyright (c) 2014, 2015 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import ast
7import os
8
9import test_flag
10
11from utils import command_executer
12
13
14class MissingImage(Exception):
15  """Raised when the requested image does not exist in gs://"""
16
17
18class ImageDownloader(object):
19
20  def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None):
21    self._logger = logger_to_use
22    self.log_level = log_level
23    self._ce = cmd_exec or command_executer.GetCommandExecuter(
24        self._logger,
25        log_level=self.log_level)
26
27  def _GetBuildID(self, chromeos_root, xbuddy_label):
28    # Get the translation of the xbuddy_label into the real Google Storage
29    # image name.
30    command = ('cd ~/trunk/src/third_party/toolchain-utils/crosperf; '
31               "python translate_xbuddy.py '%s'" % xbuddy_label)
32    retval, build_id_tuple_str, _ = self._ce.ChrootRunCommandWOutput(
33        chromeos_root, command)
34    if not build_id_tuple_str:
35      raise MissingImage("Unable to find image for '%s'" % xbuddy_label)
36
37    build_id_tuple = ast.literal_eval(build_id_tuple_str)
38    build_id = build_id_tuple[0]
39
40    return build_id
41
42  def _DownloadImage(self, chromeos_root, build_id, image_name):
43    if self.log_level == 'average':
44      self._logger.LogOutput('Preparing to download %s image to local '
45                             'directory.' % build_id)
46
47    # Make sure the directory for downloading the image exists.
48    download_path = os.path.join(chromeos_root, 'chroot/tmp', 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  def Run(self, chromeos_root, xbuddy_label):
85    build_id = self._GetBuildID(chromeos_root, xbuddy_label)
86    image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz'
87                  % build_id)
88
89    # Verify that image exists for build_id, before attempting to
90    # download it.
91    status = 0
92    if not test_flag.GetTestMode():
93      cmd = 'gsutil ls %s' % image_name
94      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
95    if status != 0:
96      raise MissingImage('Cannot find official image: %s.' % image_name)
97    image_path = self._DownloadImage(chromeos_root, build_id, image_name)
98    retval = 0
99    if image_path:
100      retval = self._UncompressImage(chromeos_root, build_id)
101    else:
102      retval = 1
103
104    if retval == 0 and self.log_level != 'quiet':
105      self._logger.LogOutput('Using image from %s.' % image_path)
106
107    return retval, image_path
108