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 = ( 141 'gs://chromeos-image-archive/%s/%s' % (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 = ( 175 'cd %s ; %s %s' % (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 = ( 197 'gs://chromeos-image-archive/%s/%s' % (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( 231 '(Warning: Could not find autotest packages .)\n' 232 '(Warning: Defaulting autotest path to %s .' % 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 # Rename created autotest directory to autotest_files 285 command = ('cd %s ; mv debug debug_files' % download_path) 286 if self.log_level != 'verbose': 287 self._logger.LogOutput('CMD: %s' % command) 288 print('(Moving downloaded debug files to debug_files)') 289 retval = self._ce.RunCommand(command) 290 if retval != 0: 291 raise MissingFile('Could not create directory debug_files') 292 293 return debug_rel_path 294 295 def Run(self, chromeos_root, xbuddy_label, autotest_path, debug_path, 296 download_debug): 297 build_id = self.GetBuildID(chromeos_root, xbuddy_label) 298 image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz' 299 % build_id) 300 301 # Verify that image exists for build_id, before attempting to 302 # download it. 303 status = 0 304 if not test_flag.GetTestMode(): 305 gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) 306 cmd = '%s ls %s' % (gsutil_cmd, image_name) 307 status = self._ce.RunCommand(cmd) 308 if status != 0: 309 raise MissingImage('Cannot find official image: %s.' % image_name) 310 311 image_path = self.DownloadImage(chromeos_root, build_id, image_name) 312 self.UncompressImage(chromeos_root, build_id) 313 314 if self.log_level != 'quiet': 315 self._logger.LogOutput('Using image from %s.' % image_path) 316 317 if autotest_path == '': 318 autotest_path = self.DownloadAutotestFiles(chromeos_root, build_id) 319 320 if debug_path == '' and download_debug: 321 debug_path = self.DownloadDebugFile(chromeos_root, build_id) 322 323 return image_path, autotest_path, debug_path 324