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 15GS_UTIL = 'chromium/tools/depot_tools/gsutil.py' 16 17 18class MissingImage(Exception): 19 """Raised when the requested image does not exist in gs://""" 20 21 22class MissingFile(Exception): 23 """Raised when the requested file does not exist in gs://""" 24 25 26class RunCommandExceptionHandler(object): 27 """Handle Exceptions from calls to RunCommand""" 28 29 def __init__(self, logger_to_use, log_level, cmd_exec, command): 30 self.logger = logger_to_use 31 self.log_level = log_level 32 self.ce = cmd_exec 33 self.cleanup_command = command 34 35 def HandleException(self, _, e): 36 # Exception handler, Run specified command 37 if self.log_level != 'verbose' and self.cleanup_command is not None: 38 self.logger.LogOutput('CMD: %s' % self.cleanup_command) 39 if self.cleanup_command is not None: 40 _ = self.ce.RunCommand(self.cleanup_command) 41 # Raise exception again 42 raise e 43 44 45class ImageDownloader(object): 46 """Download images from Cloud Storage.""" 47 48 def __init__(self, logger_to_use=None, log_level='verbose', cmd_exec=None): 49 self._logger = logger_to_use 50 self.log_level = log_level 51 self._ce = cmd_exec or command_executer.GetCommandExecuter( 52 self._logger, log_level=self.log_level) 53 54 def GetBuildID(self, chromeos_root, xbuddy_label): 55 # Get the translation of the xbuddy_label into the real Google Storage 56 # image name. 57 command = ('cd ~/trunk/src/third_party/toolchain-utils/crosperf; ' 58 "python translate_xbuddy.py '%s'" % xbuddy_label) 59 _, build_id_tuple_str, _ = self._ce.ChrootRunCommandWOutput(chromeos_root, 60 command) 61 if not build_id_tuple_str: 62 raise MissingImage("Unable to find image for '%s'" % xbuddy_label) 63 64 build_id_tuple = ast.literal_eval(build_id_tuple_str) 65 build_id = build_id_tuple[0] 66 67 return build_id 68 69 def DownloadImage(self, chromeos_root, build_id, image_name): 70 if self.log_level == 'average': 71 self._logger.LogOutput('Preparing to download %s image to local ' 72 'directory.' % build_id) 73 74 # Make sure the directory for downloading the image exists. 75 download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) 76 image_path = os.path.join(download_path, 'chromiumos_test_image.bin') 77 if not os.path.exists(download_path): 78 os.makedirs(download_path) 79 80 # Check to see if the image has already been downloaded. If not, 81 # download the image. 82 if not os.path.exists(image_path): 83 gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) 84 command = '%s cp %s %s' % (gsutil_cmd, image_name, download_path) 85 86 if self.log_level != 'verbose': 87 self._logger.LogOutput('CMD: %s' % command) 88 status = self._ce.RunCommand(command) 89 downloaded_image_name = os.path.join(download_path, 90 'chromiumos_test_image.tar.xz') 91 if status != 0 or not os.path.exists(downloaded_image_name): 92 raise MissingImage('Cannot download image: %s.' % downloaded_image_name) 93 94 return image_path 95 96 def UncompressImage(self, chromeos_root, build_id): 97 # Check to see if the file has already been uncompresssed, etc. 98 if os.path.exists( 99 os.path.join(chromeos_root, 'chroot/tmp', build_id, 100 'chromiumos_test_image.bin')): 101 return 102 103 # Uncompress and untar the downloaded image. 104 download_path = os.path.join(chromeos_root, 'chroot/tmp', build_id) 105 command = ('cd %s ; tar -Jxf chromiumos_test_image.tar.xz ' % download_path) 106 # Cleanup command for exception handler 107 clean_cmd = ('cd %s ; rm -f chromiumos_test_image.bin ' % download_path) 108 exception_handler = RunCommandExceptionHandler(self._logger, self.log_level, 109 self._ce, clean_cmd) 110 if self.log_level != 'verbose': 111 self._logger.LogOutput('CMD: %s' % command) 112 print('(Uncompressing and un-tarring may take a couple of minutes...' 113 'please be patient.)') 114 retval = self._ce.RunCommand( 115 command, except_handler=exception_handler.HandleException) 116 if retval != 0: 117 if self.log_level != 'verbose': 118 self._logger.LogOutput('CMD: %s' % clean_cmd) 119 print('(Removing file chromiumos_test_image.bin.)') 120 # Remove partially uncompressed file 121 _ = self._ce.RunCommand(clean_cmd) 122 # Raise exception for failure to uncompress 123 raise MissingImage('Cannot uncompress image: %s.' % build_id) 124 125 # Remove compressed image 126 command = ('cd %s ; rm -f chromiumos_test_image.tar.xz; ' % download_path) 127 if self.log_level != 'verbose': 128 self._logger.LogOutput('CMD: %s' % command) 129 print('(Removing file chromiumos_test_image.tar.xz.)') 130 # try removing file, its ok to have an error, print if encountered 131 retval = self._ce.RunCommand(command) 132 if retval != 0: 133 print('(Warning: Could not remove file chromiumos_test_image.tar.xz .)') 134 135 def DownloadSingleAutotestFile(self, chromeos_root, build_id, 136 package_file_name): 137 # Verify if package files exist 138 status = 0 139 gs_package_name = ('gs://chromeos-image-archive/%s/%s' % 140 (build_id, package_file_name)) 141 gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) 142 if not test_flag.GetTestMode(): 143 cmd = '%s ls %s' % (gsutil_cmd, gs_package_name) 144 status = self._ce.RunCommand(cmd) 145 if status != 0: 146 raise MissingFile('Cannot find autotest package file: %s.' % 147 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 UncompressSingleAutotestFile(self, chromeos_root, build_id, 171 package_file_name, 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 autotest 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 autotest 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 VerifyAutotestFilesExist(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.VerifyAutotestFilesExist(chromeos_root, build_id, 227 autotest_packages_name) 228 if status != 0: 229 default_autotest_dir = '~/trunk/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.DownloadSingleAutotestFile(chromeos_root, build_id, 237 autotest_packages_name) 238 self.DownloadSingleAutotestFile(chromeos_root, build_id, 239 autotest_server_package_name) 240 self.DownloadSingleAutotestFile(chromeos_root, build_id, 241 autotest_control_files_name) 242 243 self.UncompressSingleAutotestFile(chromeos_root, build_id, 244 autotest_packages_name, 'tar -xvf ') 245 self.UncompressSingleAutotestFile(chromeos_root, build_id, 246 autotest_server_package_name, 247 'tar -jxvf ') 248 self.UncompressSingleAutotestFile(chromeos_root, build_id, 249 autotest_control_files_name, 250 'tar -xvf ') 251 # Rename created autotest directory to autotest_files 252 command = ('cd %s ; mv autotest autotest_files' % download_path) 253 if self.log_level != 'verbose': 254 self._logger.LogOutput('CMD: %s' % command) 255 print('(Moving downloaded autotest files to autotest_files)') 256 retval = self._ce.RunCommand(command) 257 if retval != 0: 258 raise MissingFile('Could not create directory autotest_files') 259 260 return autotest_rel_path 261 262 def Run(self, chromeos_root, xbuddy_label, autotest_path): 263 build_id = self.GetBuildID(chromeos_root, xbuddy_label) 264 image_name = ('gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz' 265 % build_id) 266 267 # Verify that image exists for build_id, before attempting to 268 # download it. 269 status = 0 270 if not test_flag.GetTestMode(): 271 gsutil_cmd = os.path.join(chromeos_root, GS_UTIL) 272 cmd = '%s ls %s' % (gsutil_cmd, image_name) 273 status = self._ce.RunCommand(cmd) 274 if status != 0: 275 raise MissingImage('Cannot find official image: %s.' % image_name) 276 277 image_path = self.DownloadImage(chromeos_root, build_id, image_name) 278 self.UncompressImage(chromeos_root, build_id) 279 280 if self.log_level != 'quiet': 281 self._logger.LogOutput('Using image from %s.' % image_path) 282 283 if autotest_path == '': 284 autotest_path = self.DownloadAutotestFiles(chromeos_root, build_id) 285 286 return image_path, autotest_path 287