download_images.py revision c7ed2bca6f5dbf0cc6ee7cd1e7d30ea3119adbe8 
1# Copyright (c) 2014, 2015 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import ast
6import os
7
8import test_flag
9
10from utils import command_executer
11
12
13class MissingImage(Exception):
14  """Raised when the requested image does not exist in gs://"""
15
16
17class ImageDownloader(object):
18
19  def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None):
20    self._logger = logger_to_use
21    self.log_level = log_level
22    self._ce = cmd_exec or command_executer.GetCommandExecuter(
23        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.ChrootRunCommandWOutput(
32        chromeos_root, command)
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', build_id)
48    image_path = os.path.join(download_path, 'chromiumos_test_image.bin')
49    if not os.path.exists(download_path):
50      os.makedirs(download_path)
51
52    # Check to see if the image has already been downloaded.  If not,
53    # download the image.
54    if not os.path.exists(image_path):
55      command = 'gsutil cp %s /tmp/%s' % (image_name, build_id)
56
57      if self.log_level != 'verbose':
58        self._logger.LogOutput('CMD: %s' % command)
59      status = self._ce.ChrootRunCommand(chromeos_root, command)
60      if status != 0 or not os.path.exists(image_path):
61        raise MissingImage('Cannot download image: %s.' % image_name)
62
63    return image_path
64
65  def UncompressImage(self, chromeos_root, build_id):
66    # Check to see if the file has already been uncompresssed, etc.
67    if os.path.exists(os.path.join(chromeos_root, 'chroot/tmp', build_id,
68                                   'chromiumos_test_image.bin')):
69      return
70
71    # Uncompress and untar the downloaded image.
72    command = ('cd /tmp/%s ;unxz chromiumos_test_image.tar.xz; '
73               'tar -xvf chromiumos_test_image.tar' % build_id)
74    if self.log_level != 'verbose':
75      self._logger.LogOutput('CMD: %s' % command)
76      print('(Uncompressing and un-tarring may take a couple of minutes...'
77            'please be patient.)')
78    retval = self._ce.ChrootRunCommand(chromeos_root, command)
79    if retval != 0:
80      raise MissingImage('Cannot uncompress image: %s.' % build_id)
81
82  def Run(self, chromeos_root, xbuddy_label):
83    build_id = self.GetBuildID(chromeos_root, xbuddy_label)
84    image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz'
85                  % build_id)
86
87    # Verify that image exists for build_id, before attempting to
88    # download it.
89    status = 0
90    if not test_flag.GetTestMode():
91      cmd = 'gsutil ls %s' % image_name
92      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
93    if status != 0:
94      raise MissingImage('Cannot find official image: %s.' % image_name)
95
96    image_path = self.DownloadImage(chromeos_root, build_id, image_name)
97    self.UncompressImage(chromeos_root, build_id)
98
99    if self.log_level != 'quiet':
100      self._logger.LogOutput('Using image from %s.' % image_path)
101
102    return image_path
103