download_images.py revision 4b133961b76c2cb8bc58f0ea2cded9e3438ffb6f 
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"""Download images from Cloud Storage."""
5
6from __future__ import print_function
7
8import ast
9import os
10
11import test_flag
12
13from cros_utils import command_executer
14
15
16class MissingImage(Exception):
17  """Raised when the requested image does not exist in gs://"""
18
19
20class MissingFile(Exception):
21  """Raised when the requested file does not exist in gs://"""
22
23
24class ImageDownloader(object):
25  """Download images from Cloud Storage."""
26
27  def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None):
28    self._logger = logger_to_use
29    self.log_level = log_level
30    self._ce = cmd_exec or command_executer.GetCommandExecuter(
31        self._logger, log_level=self.log_level)
32
33  def GetBuildID(self, chromeos_root, xbuddy_label):
34    # Get the translation of the xbuddy_label into the real Google Storage
35    # image name.
36    command = ('cd ~/trunk/src/third_party/toolchain-utils/crosperf; '
37               "python translate_xbuddy.py '%s'" % xbuddy_label)
38    _, build_id_tuple_str, _ = self._ce.ChrootRunCommandWOutput(chromeos_root,
39                                                                command)
40    if not build_id_tuple_str:
41      raise MissingImage("Unable to find image for '%s'" % xbuddy_label)
42
43    build_id_tuple = ast.literal_eval(build_id_tuple_str)
44    build_id = build_id_tuple[0]
45
46    return build_id
47
48  def DownloadImage(self, chromeos_root, build_id, image_name):
49    if self.log_level == 'average':
50      self._logger.LogOutput('Preparing to download %s image to local '
51                             'directory.' % build_id)
52
53    # Make sure the directory for downloading the image exists.
54    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
55    image_path = os.path.join(download_path, 'chromiumos_test_image.bin')
56    if not os.path.exists(download_path):
57      os.makedirs(download_path)
58
59    # Check to see if the image has already been downloaded.  If not,
60    # download the image.
61    if not os.path.exists(image_path):
62      command = 'gsutil cp %s /tmp/%s' % (image_name, build_id)
63
64      if self.log_level != 'verbose':
65        self._logger.LogOutput('CMD: %s' % command)
66      status = self._ce.ChrootRunCommand(chromeos_root, command)
67      downloaded_image_name = os.path.join(download_path,
68                                           'chromiumos_test_image.tar.xz')
69      if status != 0 or not os.path.exists(downloaded_image_name):
70        raise MissingImage('Cannot download image: %s.' % downloaded_image_name)
71
72    return image_path
73
74  def UncompressImage(self, chromeos_root, build_id):
75    # Check to see if the file has already been uncompresssed, etc.
76    if os.path.exists(
77        os.path.join(chromeos_root, 'chroot/tmp', build_id,
78                     'chromiumos_test_image.bin')):
79      return
80
81    # Uncompress and untar the downloaded image.
82    command = ('cd /tmp/%s ;unxz chromiumos_test_image.tar.xz; '
83               'tar -xvf chromiumos_test_image.tar' % build_id)
84    if self.log_level != 'verbose':
85      self._logger.LogOutput('CMD: %s' % command)
86      print('(Uncompressing and un-tarring may take a couple of minutes...'
87            'please be patient.)')
88    retval = self._ce.ChrootRunCommand(chromeos_root, command)
89    if retval != 0:
90      raise MissingImage('Cannot uncompress image: %s.' % build_id)
91
92  def DownloadSingleAutotestFile(self, chromeos_root, build_id,
93                                 package_file_name):
94    # Verify if package files exist
95    status = 0
96    gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
97                       (build_id, package_file_name))
98    if not test_flag.GetTestMode():
99      cmd = 'gsutil ls %s' % gs_package_name
100      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
101    if status != 0:
102      raise MissingFile('Cannot find autotest package file: %s.' %
103                        package_file_name)
104
105    if self.log_level == 'average':
106      self._logger.LogOutput('Preparing to download %s package to local '
107                             'directory.' % package_file_name)
108
109    # Make sure the directory for downloading the package exists.
110    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
111    package_path = os.path.join(download_path, package_file_name)
112    if not os.path.exists(download_path):
113      os.makedirs(download_path)
114
115    # Check to see if the package file has already been downloaded.  If not,
116    # download it.
117    if not os.path.exists(package_path):
118      command = 'gsutil cp %s /tmp/%s' % (gs_package_name, build_id)
119
120      if self.log_level != 'verbose':
121        self._logger.LogOutput('CMD: %s' % command)
122      status = self._ce.ChrootRunCommand(chromeos_root, command)
123      if status != 0 or not os.path.exists(package_path):
124        raise MissingFile('Cannot download package: %s .' % package_path)
125
126  def UncompressSingleAutotestFile(self, chromeos_root, build_id,
127                                   package_file_name, uncompress_cmd):
128    # Uncompress file
129    command = ('cd /tmp/%s ; %s %s' %
130               (build_id, uncompress_cmd, package_file_name))
131
132    if self.log_level != 'verbose':
133      self._logger.LogOutput('CMD: %s' % command)
134      print('(Uncompressing autotest file %s .)' % package_file_name)
135    retval = self._ce.ChrootRunCommand(chromeos_root, command)
136    if retval != 0:
137      raise MissingFile('Cannot uncompress file: %s.' % package_file_name)
138    # Remove uncompressed downloaded file
139    command = ('cd /tmp/%s ; rm -f %s' % (build_id, package_file_name))
140    if self.log_level != 'verbose':
141      self._logger.LogOutput('CMD: %s' % command)
142      print('(Removing processed autotest file %s .)' % package_file_name)
143    # try removing file, its ok to have an error, print if encountered
144    retval = self._ce.ChrootRunCommand(chromeos_root, command)
145    if retval != 0:
146      print('(Warning: Could not remove file %s .)' % package_file_name)
147
148  def DownloadAutotestFiles(self, chromeos_root, build_id):
149    # Download autest package files (3 files)
150    autotest_packages_name = ('autotest_packages.tar')
151    autotest_server_package_name = ('autotest_server_package.tar.bz2')
152    autotest_control_files_name = ('control_files.tar')
153
154    # Autotest directory relative path wrt chroot
155    autotest_rel_path = os.path.join('/tmp', build_id, 'autotest_files')
156    # Absolute Path to download files
157    autotest_path = os.path.join(chromeos_root, 'chroot/tmp', build_id,
158                                 'autotest_files')
159
160    if not os.path.exists(autotest_path):
161      self.DownloadSingleAutotestFile(chromeos_root, build_id,
162                                      autotest_packages_name)
163      self.DownloadSingleAutotestFile(chromeos_root, build_id,
164                                      autotest_server_package_name)
165      self.DownloadSingleAutotestFile(chromeos_root, build_id,
166                                      autotest_control_files_name)
167
168      self.UncompressSingleAutotestFile(chromeos_root, build_id,
169                                        autotest_packages_name, 'tar -xvf ')
170      self.UncompressSingleAutotestFile(chromeos_root, build_id,
171                                        autotest_server_package_name,
172                                        'tar -jxvf ')
173      self.UncompressSingleAutotestFile(chromeos_root, build_id,
174                                        autotest_control_files_name,
175                                        'tar -xvf ')
176      # Rename created autotest directory to autotest_files
177      command = ('cd /tmp/%s ; mv autotest autotest_files' % build_id)
178      if self.log_level != 'verbose':
179        self._logger.LogOutput('CMD: %s' % command)
180        print('(Moving downloaded autotest files to autotest_files)')
181      retval = self._ce.ChrootRunCommand(chromeos_root, command)
182      if retval != 0:
183        raise MissingFile('Could not create directory autotest_files')
184
185    return autotest_rel_path
186
187  def Run(self, chromeos_root, xbuddy_label, autotest_path):
188    build_id = self.GetBuildID(chromeos_root, xbuddy_label)
189    image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz'
190                  % build_id)
191
192    # Verify that image exists for build_id, before attempting to
193    # download it.
194    status = 0
195    if not test_flag.GetTestMode():
196      cmd = 'gsutil ls %s' % image_name
197      status = self._ce.ChrootRunCommand(chromeos_root, cmd)
198    if status != 0:
199      raise MissingImage('Cannot find official image: %s.' % image_name)
200
201    image_path = self.DownloadImage(chromeos_root, build_id, image_name)
202    self.UncompressImage(chromeos_root, build_id)
203
204    if self.log_level != 'quiet':
205      self._logger.LogOutput('Using image from %s.' % image_path)
206
207    if autotest_path == '':
208      autotest_path = self.DownloadAutotestFiles(chromeos_root, build_id)
209
210    return image_path, autotest_path
211