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