• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# -*- coding: utf-8 -*-
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
6"""Download images from Cloud Storage."""
7
8from __future__ import print_function
9
10import ast
11import os
12
13import test_flag
14
15from cros_utils import command_executer
16
17GS_UTIL = 'src/chromium/depot_tools/gsutil.py'
18
19
20class MissingImage(Exception):
21  """Raised when the requested image does not exist in gs://"""
22
23
24class MissingFile(Exception):
25  """Raised when the requested file does not exist in gs://"""
26
27
28class RunCommandExceptionHandler(object):
29  """Handle Exceptions from calls to RunCommand"""
30
31  def __init__(self, logger_to_use, log_level, cmd_exec, command):
32    self.logger = logger_to_use
33    self.log_level = log_level
34    self.ce = cmd_exec
35    self.cleanup_command = command
36
37  def HandleException(self, _, e):
38    # Exception handler, Run specified command
39    if self.log_level != 'verbose' and self.cleanup_command is not None:
40      self.logger.LogOutput('CMD: %s' % self.cleanup_command)
41    if self.cleanup_command is not None:
42      _ = self.ce.RunCommand(self.cleanup_command)
43    # Raise exception again
44    raise e
45
46
47class ImageDownloader(object):
48  """Download images from Cloud Storage."""
49
50  def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None):
51    self._logger = logger_to_use
52    self.log_level = log_level
53    self._ce = cmd_exec or command_executer.GetCommandExecuter(
54        self._logger, log_level=self.log_level)
55
56  def GetBuildID(self, chromeos_root, xbuddy_label):
57    # Get the translation of the xbuddy_label into the real Google Storage
58    # image name.
59    command = ('cd /mnt/host/source/src/third_party/toolchain-utils/crosperf; '
60               "./translate_xbuddy.py '%s'" % xbuddy_label)
61    _, build_id_tuple_str, _ = self._ce.ChrootRunCommandWOutput(
62        chromeos_root, command)
63    if not build_id_tuple_str:
64      raise MissingImage("Unable to find image for '%s'" % xbuddy_label)
65
66    build_id_tuple = ast.literal_eval(build_id_tuple_str)
67    build_id = build_id_tuple[0]
68
69    return build_id
70
71  def DownloadImage(self, chromeos_root, build_id, image_name):
72    if self.log_level == 'average':
73      self._logger.LogOutput('Preparing to download %s image to local '
74                             'directory.' % build_id)
75
76    # Make sure the directory for downloading the image exists.
77    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
78    image_path = os.path.join(download_path, 'chromiumos_test_image.bin')
79    if not os.path.exists(download_path):
80      os.makedirs(download_path)
81
82    # Check to see if the image has already been downloaded.  If not,
83    # download the image.
84    if not os.path.exists(image_path):
85      gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
86      command = '%s cp %s %s' % (gsutil_cmd, image_name, download_path)
87
88      if self.log_level != 'verbose':
89        self._logger.LogOutput('CMD: %s' % command)
90      status = self._ce.RunCommand(command)
91      downloaded_image_name = os.path.join(download_path,
92                                           'chromiumos_test_image.tar.xz')
93      if status != 0 or not os.path.exists(downloaded_image_name):
94        raise MissingImage('Cannot download image: %s.' % downloaded_image_name)
95
96    return image_path
97
98  def UncompressImage(self, chromeos_root, build_id):
99    # Check to see if the file has already been uncompresssed, etc.
100    if os.path.exists(
101        os.path.join(chromeos_root, 'chroot/tmp', build_id,
102                     'chromiumos_test_image.bin')):
103      return
104
105    # Uncompress and untar the downloaded image.
106    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
107    command = ('cd %s ; tar -Jxf chromiumos_test_image.tar.xz ' % download_path)
108    # Cleanup command for exception handler
109    clean_cmd = ('cd %s ; rm -f chromiumos_test_image.bin ' % download_path)
110    exception_handler = RunCommandExceptionHandler(self._logger, self.log_level,
111                                                   self._ce, clean_cmd)
112    if self.log_level != 'verbose':
113      self._logger.LogOutput('CMD: %s' % command)
114      print('(Uncompressing and un-tarring may take a couple of minutes...'
115            'please be patient.)')
116    retval = self._ce.RunCommand(
117        command, except_handler=exception_handler.HandleException)
118    if retval != 0:
119      if self.log_level != 'verbose':
120        self._logger.LogOutput('CMD: %s' % clean_cmd)
121        print('(Removing file chromiumos_test_image.bin.)')
122      # Remove partially uncompressed file
123      _ = self._ce.RunCommand(clean_cmd)
124      # Raise exception for failure to uncompress
125      raise MissingImage('Cannot uncompress image: %s.' % build_id)
126
127    # Remove compressed image
128    command = ('cd %s ; rm -f chromiumos_test_image.tar.xz; ' % download_path)
129    if self.log_level != 'verbose':
130      self._logger.LogOutput('CMD: %s' % command)
131      print('(Removing file chromiumos_test_image.tar.xz.)')
132    # try removing file, its ok to have an error, print if encountered
133    retval = self._ce.RunCommand(command)
134    if retval != 0:
135      print('(Warning: Could not remove file chromiumos_test_image.tar.xz .)')
136
137  def DownloadSingleFile(self, chromeos_root, build_id, package_file_name):
138    # Verify if package files exist
139    status = 0
140    gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
141                       (build_id, package_file_name))
142    gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
143    if not test_flag.GetTestMode():
144      cmd = '%s ls %s' % (gsutil_cmd, gs_package_name)
145      status = self._ce.RunCommand(cmd)
146    if status != 0:
147      raise MissingFile('Cannot find package file: %s.' % package_file_name)
148
149    if self.log_level == 'average':
150      self._logger.LogOutput('Preparing to download %s package to local '
151                             'directory.' % package_file_name)
152
153    # Make sure the directory for downloading the package exists.
154    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
155    package_path = os.path.join(download_path, package_file_name)
156    if not os.path.exists(download_path):
157      os.makedirs(download_path)
158
159    # Check to see if the package file has already been downloaded.  If not,
160    # download it.
161    if not os.path.exists(package_path):
162      command = '%s cp %s %s' % (gsutil_cmd, gs_package_name, download_path)
163
164      if self.log_level != 'verbose':
165        self._logger.LogOutput('CMD: %s' % command)
166      status = self._ce.RunCommand(command)
167      if status != 0 or not os.path.exists(package_path):
168        raise MissingFile('Cannot download package: %s .' % package_path)
169
170  def UncompressSingleFile(self, chromeos_root, build_id, package_file_name,
171                           uncompress_cmd):
172    # Uncompress file
173    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
174    command = ('cd %s ; %s %s' %
175               (download_path, uncompress_cmd, package_file_name))
176
177    if self.log_level != 'verbose':
178      self._logger.LogOutput('CMD: %s' % command)
179      print('(Uncompressing file %s .)' % package_file_name)
180    retval = self._ce.RunCommand(command)
181    if retval != 0:
182      raise MissingFile('Cannot uncompress file: %s.' % package_file_name)
183    # Remove uncompressed downloaded file
184    command = ('cd %s ; rm -f %s' % (download_path, package_file_name))
185    if self.log_level != 'verbose':
186      self._logger.LogOutput('CMD: %s' % command)
187      print('(Removing processed file %s .)' % package_file_name)
188    # try removing file, its ok to have an error, print if encountered
189    retval = self._ce.RunCommand(command)
190    if retval != 0:
191      print('(Warning: Could not remove file %s .)' % package_file_name)
192
193  def VerifyFileExists(self, chromeos_root, build_id, package_file):
194    # Quickly verify if the files are there
195    status = 0
196    gs_package_name = ('gs://chromeos-image-archive/%s/%s' %
197                       (build_id, package_file))
198    gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
199    if not test_flag.GetTestMode():
200      cmd = '%s ls %s' % (gsutil_cmd, gs_package_name)
201      if self.log_level != 'verbose':
202        self._logger.LogOutput('CMD: %s' % cmd)
203      status = self._ce.RunCommand(cmd)
204      if status != 0:
205        print('(Warning: Could not find file %s )' % gs_package_name)
206        return 1
207    # Package exists on server
208    return 0
209
210  def DownloadAutotestFiles(self, chromeos_root, build_id):
211    # Download autest package files (3 files)
212    autotest_packages_name = ('autotest_packages.tar')
213    autotest_server_package_name = ('autotest_server_package.tar.bz2')
214    autotest_control_files_name = ('control_files.tar')
215
216    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
217    # Autotest directory relative path wrt chroot
218    autotest_rel_path = os.path.join('/tmp', build_id, 'autotest_files')
219    # Absolute Path to download files
220    autotest_path = os.path.join(chromeos_root, 'chroot/tmp', build_id,
221                                 'autotest_files')
222
223    if not os.path.exists(autotest_path):
224      # Quickly verify if the files are present on server
225      # If not, just exit with warning
226      status = self.VerifyFileExists(chromeos_root, build_id,
227                                     autotest_packages_name)
228      if status != 0:
229        default_autotest_dir = '/mnt/host/source/src/third_party/autotest/files'
230        print('(Warning: Could not find autotest packages .)\n'
231              '(Warning: Defaulting autotest path to %s .' %
232              default_autotest_dir)
233        return default_autotest_dir
234
235      # Files exist on server, download and uncompress them
236      self.DownloadSingleFile(chromeos_root, build_id, autotest_packages_name)
237      self.DownloadSingleFile(chromeos_root, build_id,
238                              autotest_server_package_name)
239      self.DownloadSingleFile(chromeos_root, build_id,
240                              autotest_control_files_name)
241
242      self.UncompressSingleFile(chromeos_root, build_id, autotest_packages_name,
243                                'tar -xf ')
244      self.UncompressSingleFile(chromeos_root, build_id,
245                                autotest_server_package_name, 'tar -jxf ')
246      self.UncompressSingleFile(chromeos_root, build_id,
247                                autotest_control_files_name, 'tar -xf ')
248      # Rename created autotest directory to autotest_files
249      command = ('cd %s ; mv autotest autotest_files' % download_path)
250      if self.log_level != 'verbose':
251        self._logger.LogOutput('CMD: %s' % command)
252        print('(Moving downloaded autotest files to autotest_files)')
253      retval = self._ce.RunCommand(command)
254      if retval != 0:
255        raise MissingFile('Could not create directory autotest_files')
256
257    return autotest_rel_path
258
259  def DownloadDebugFile(self, chromeos_root, build_id):
260    # Download autest package files (3 files)
261    debug_archive_name = 'debug.tgz'
262
263    download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id)
264    # Debug directory relative path wrt chroot
265    debug_rel_path = os.path.join('/tmp', build_id, 'debug_files')
266    # Debug path to download files
267    debug_path = os.path.join(chromeos_root, 'chroot/tmp', build_id,
268                              'debug_files')
269
270    if not os.path.exists(debug_path):
271      # Quickly verify if the file is present on server
272      # If not, just exit with warning
273      status = self.VerifyFileExists(chromeos_root, build_id,
274                                     debug_archive_name)
275      if status != 0:
276        self._logger.LogOutput('WARNING: Could not find debug archive on gs')
277        return ''
278
279      # File exists on server, download and uncompress it
280      self.DownloadSingleFile(chromeos_root, build_id, debug_archive_name)
281
282      self.UncompressSingleFile(chromeos_root, build_id, debug_archive_name,
283                                'tar -xf ')
284      # Extract and move debug files into the proper location.
285      debug_dir = 'debug_files/usr/lib'
286      command = ('cd %s ; mkdir -p %s; mv debug %s' %
287                 (download_path, debug_dir, debug_dir))
288      if self.log_level != 'verbose':
289        self._logger.LogOutput('CMD: %s' % command)
290        print('Moving downloaded debug files to %s' % debug_dir)
291      retval = self._ce.RunCommand(command)
292      if retval != 0:
293        raise MissingFile('Could not create directory %s' %
294                          os.path.join(debug_dir, 'debug'))
295
296    return debug_rel_path
297
298  def Run(self, chromeos_root, xbuddy_label, autotest_path, debug_path,
299          download_debug):
300    build_id = self.GetBuildID(chromeos_root, xbuddy_label)
301    image_name = (
302        'gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz' %
303        build_id)
304
305    # Verify that image exists for build_id, before attempting to
306    # download it.
307    status = 0
308    if not test_flag.GetTestMode():
309      gsutil_cmd = os.path.join(chromeos_root, GS_UTIL)
310      cmd = '%s ls %s' % (gsutil_cmd, image_name)
311      status = self._ce.RunCommand(cmd)
312    if status != 0:
313      raise MissingImage('Cannot find official image: %s.' % image_name)
314
315    image_path = self.DownloadImage(chromeos_root, build_id, image_name)
316    self.UncompressImage(chromeos_root, build_id)
317
318    if self.log_level != 'quiet':
319      self._logger.LogOutput('Using image from %s.' % image_path)
320
321    if autotest_path == '':
322      autotest_path = self.DownloadAutotestFiles(chromeos_root, build_id)
323
324    if debug_path == '' and download_debug:
325      debug_path = self.DownloadDebugFile(chromeos_root, build_id)
326
327    return image_path, autotest_path, debug_path
328