1 #!/usr/bin/env python 2 # 3 # Copyright 2016 - The Android Open Source Project 4 # 5 # Licensed under the Apache License, Version 2.0 (the "License"); 6 # you may not use this file except in compliance with the License. 7 # You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 17 """Public Device Driver APIs. 18 19 This module provides public device driver APIs that can be called 20 as a Python library. 21 22 TODO(fdeng): The following APIs have not been implemented 23 - RebootAVD(ip): 24 - RegisterSshPubKey(username, key): 25 - UnregisterSshPubKey(username, key): 26 - CleanupStaleImages(): 27 - CleanupStaleDevices(): 28 """ 29 30 import datetime 31 import logging 32 import os 33 34 import dateutil.parser 35 import dateutil.tz 36 37 from acloud.public import avd 38 from acloud.public import errors 39 from acloud.public import report 40 from acloud.internal import constants 41 from acloud.internal.lib import auth 42 from acloud.internal.lib import android_build_client 43 from acloud.internal.lib import android_compute_client 44 from acloud.internal.lib import gstorage_client 45 from acloud.internal.lib import utils 46 47 logger = logging.getLogger(__name__) 48 49 ALL_SCOPES = " ".join([android_build_client.AndroidBuildClient.SCOPE, 50 gstorage_client.StorageClient.SCOPE, 51 android_compute_client.AndroidComputeClient.SCOPE]) 52 53 MAX_BATCH_CLEANUP_COUNT = 100 54 55 56 class AndroidVirtualDevicePool(object): 57 """A class that manages a pool of devices.""" 58 59 def __init__(self, cfg, devices=None): 60 self._devices = devices or [] 61 self._cfg = cfg 62 credentials = auth.CreateCredentials(cfg, ALL_SCOPES) 63 self._build_client = android_build_client.AndroidBuildClient( 64 credentials) 65 self._storage_client = gstorage_client.StorageClient(credentials) 66 self._compute_client = android_compute_client.AndroidComputeClient( 67 cfg, credentials) 68 69 def _CreateGceImageWithBuildInfo(self, build_target, build_id): 70 """Creates a Gce image using build from Launch Control. 71 72 Clone avd-system.tar.gz of a build to a cache storage bucket 73 using launch control api. And then create a Gce image. 74 75 Args: 76 build_target: Target name, e.g. "gce_x86-userdebug" 77 build_id: Build id, a string, e.g. "2263051", "P2804227" 78 79 Returns: 80 String, name of the Gce image that has been created. 81 """ 82 logger.info("Creating a new gce image using build: build_id %s, " 83 "build_target %s", build_id, build_target) 84 disk_image_id = utils.GenerateUniqueName( 85 suffix=self._cfg.disk_image_name) 86 self._build_client.CopyTo( 87 build_target, 88 build_id, 89 artifact_name=self._cfg.disk_image_name, 90 destination_bucket=self._cfg.storage_bucket_name, 91 destination_path=disk_image_id) 92 disk_image_url = self._storage_client.GetUrl( 93 self._cfg.storage_bucket_name, disk_image_id) 94 try: 95 image_name = self._compute_client.GenerateImageName(build_target, 96 build_id) 97 self._compute_client.CreateImage(image_name=image_name, 98 source_uri=disk_image_url) 99 finally: 100 self._storage_client.Delete(self._cfg.storage_bucket_name, 101 disk_image_id) 102 return image_name 103 104 def _CreateGceImageWithLocalFile(self, local_disk_image): 105 """Create a Gce image with a local image file. 106 107 The local disk image can be either a tar.gz file or a 108 raw vmlinux image. 109 e.g. /tmp/avd-system.tar.gz or /tmp/android_system_disk_syslinux.img 110 If a raw vmlinux image is provided, it will be archived into a tar.gz file. 111 112 The final tar.gz file will be uploaded to a cache bucket in storage. 113 114 Args: 115 local_disk_image: string, path to a local disk image, 116 117 Returns: 118 String, name of the Gce image that has been created. 119 120 Raises: 121 DriverError: if a file with an unexpected extension is given. 122 """ 123 logger.info("Creating a new gce image from a local file %s", 124 local_disk_image) 125 with utils.TempDir() as tempdir: 126 if local_disk_image.endswith(self._cfg.disk_raw_image_extension): 127 dest_tar_file = os.path.join(tempdir, 128 self._cfg.disk_image_name) 129 utils.MakeTarFile( 130 src_dict={local_disk_image: self._cfg.disk_raw_image_name}, 131 dest=dest_tar_file) 132 local_disk_image = dest_tar_file 133 elif not local_disk_image.endswith(self._cfg.disk_image_extension): 134 raise errors.DriverError( 135 "Wrong local_disk_image type, must be a *%s file or *%s file" 136 % (self._cfg.disk_raw_image_extension, 137 self._cfg.disk_image_extension)) 138 139 disk_image_id = utils.GenerateUniqueName( 140 suffix=self._cfg.disk_image_name) 141 self._storage_client.Upload( 142 local_src=local_disk_image, 143 bucket_name=self._cfg.storage_bucket_name, 144 object_name=disk_image_id, 145 mime_type=self._cfg.disk_image_mime_type) 146 disk_image_url = self._storage_client.GetUrl( 147 self._cfg.storage_bucket_name, disk_image_id) 148 try: 149 image_name = self._compute_client.GenerateImageName() 150 self._compute_client.CreateImage(image_name=image_name, 151 source_uri=disk_image_url) 152 finally: 153 self._storage_client.Delete(self._cfg.storage_bucket_name, 154 disk_image_id) 155 return image_name 156 157 def CreateDevices(self, 158 num, 159 build_target=None, 160 build_id=None, 161 gce_image=None, 162 local_disk_image=None, 163 cleanup=True, 164 extra_data_disk_size_gb=None, 165 precreated_data_image=None): 166 """Creates |num| devices for given build_target and build_id. 167 168 - If gce_image is provided, will use it to create an instance. 169 - If local_disk_image is provided, will upload it to a temporary 170 caching storage bucket which is defined by user as |storage_bucket_name| 171 And then create an gce image with it; and then create an instance. 172 - If build_target and build_id are provided, will clone the disk image 173 via launch control to the temporary caching storage bucket. 174 And then create an gce image with it; and then create an instance. 175 176 Args: 177 num: Number of devices to create. 178 build_target: Target name, e.g. "gce_x86-userdebug" 179 build_id: Build id, a string, e.g. "2263051", "P2804227" 180 gce_image: string, if given, will use this image 181 instead of creating a new one. 182 implies cleanup=False. 183 local_disk_image: string, path to a local disk image, e.g. 184 /tmp/avd-system.tar.gz 185 cleanup: boolean, if True clean up compute engine image after creating 186 the instance. 187 extra_data_disk_size_gb: Integer, size of extra disk, or None. 188 precreated_data_image: A string, the image to use for the extra disk. 189 190 Raises: 191 errors.DriverError: If no source is specified for image creation. 192 """ 193 if gce_image: 194 # GCE image is provided, we can directly move to instance creation. 195 logger.info("Using existing gce image %s", gce_image) 196 image_name = gce_image 197 cleanup = False 198 elif local_disk_image: 199 image_name = self._CreateGceImageWithLocalFile(local_disk_image) 200 elif build_target and build_id: 201 image_name = self._CreateGceImageWithBuildInfo(build_target, 202 build_id) 203 else: 204 raise errors.DriverError( 205 "Invalid image source, must specify one of the following: gce_image, " 206 "local_disk_image, or build_target and build id.") 207 208 # Create GCE instances. 209 try: 210 for _ in range(num): 211 instance = self._compute_client.GenerateInstanceName( 212 build_target, build_id) 213 extra_disk_name = None 214 if extra_data_disk_size_gb > 0: 215 extra_disk_name = self._compute_client.GetDataDiskName( 216 instance) 217 self._compute_client.CreateDisk(extra_disk_name, 218 precreated_data_image, 219 extra_data_disk_size_gb) 220 self._compute_client.CreateInstance(instance, image_name, 221 extra_disk_name) 222 ip = self._compute_client.GetInstanceIP(instance) 223 self.devices.append(avd.AndroidVirtualDevice( 224 ip=ip, instance_name=instance)) 225 finally: 226 if cleanup: 227 self._compute_client.DeleteImage(image_name) 228 229 def DeleteDevices(self): 230 """Deletes devices. 231 232 Returns: 233 A tuple, (deleted, failed, error_msgs) 234 deleted: A list of names of instances that have been deleted. 235 faild: A list of names of instances that we fail to delete. 236 error_msgs: A list of failure messages. 237 """ 238 instance_names = [device.instance_name for device in self._devices] 239 return self._compute_client.DeleteInstances(instance_names, 240 self._cfg.zone) 241 242 def WaitForBoot(self): 243 """Waits for all devices to boot up. 244 245 Returns: 246 A dictionary that contains all the failures. 247 The key is the name of the instance that fails to boot, 248 the value is an errors.DeviceBootTimeoutError object. 249 """ 250 failures = {} 251 for device in self._devices: 252 try: 253 self._compute_client.WaitForBoot(device.instance_name) 254 except errors.DeviceBootTimeoutError as e: 255 failures[device.instance_name] = e 256 return failures 257 258 @property 259 def devices(self): 260 """Returns a list of devices in the pool. 261 262 Returns: 263 A list of devices in the pool. 264 """ 265 return self._devices 266 267 268 def _AddDeletionResultToReport(report_obj, deleted, failed, error_msgs, 269 resource_name): 270 """Adds deletion result to a Report object. 271 272 This function will add the following to report.data. 273 "deleted": [ 274 {"name": "resource_name", "type": "resource_name"}, 275 ], 276 "failed": [ 277 {"name": "resource_name", "type": "resource_name"}, 278 ], 279 This function will append error_msgs to report.errors. 280 281 Args: 282 report_obj: A Report object. 283 deleted: A list of names of the resources that have been deleted. 284 failed: A list of names of the resources that we fail to delete. 285 error_msgs: A list of error message strings to be added to the report. 286 resource_name: A string, representing the name of the resource. 287 """ 288 for name in deleted: 289 report_obj.AddData(key="deleted", 290 value={"name": name, 291 "type": resource_name}) 292 for name in failed: 293 report_obj.AddData(key="failed", 294 value={"name": name, 295 "type": resource_name}) 296 report_obj.AddErrors(error_msgs) 297 if failed or error_msgs: 298 report_obj.SetStatus(report.Status.FAIL) 299 300 301 def _FetchSerialLogsFromDevices(compute_client, instance_names, output_file, 302 port): 303 """Fetch serial logs from a port for a list of devices to a local file. 304 305 Args: 306 compute_client: An object of android_compute_client.AndroidComputeClient 307 instance_names: A list of instance names. 308 output_file: A path to a file ending with "tar.gz" 309 port: The number of serial port to read from, 0 for serial output, 1 for 310 logcat. 311 """ 312 with utils.TempDir() as tempdir: 313 src_dict = {} 314 for instance_name in instance_names: 315 serial_log = compute_client.GetSerialPortOutput( 316 instance=instance_name, port=port) 317 file_name = "%s.log" % instance_name 318 file_path = os.path.join(tempdir, file_name) 319 src_dict[file_path] = file_name 320 with open(file_path, "w") as f: 321 f.write(serial_log.encode("utf-8")) 322 utils.MakeTarFile(src_dict, output_file) 323 324 325 def _CreateSshKeyPairIfNecessary(cfg): 326 """Create ssh key pair if necessary. 327 328 Args: 329 cfg: An Acloudconfig instance. 330 331 Raises: 332 error.DriverError: If it falls into an unexpected condition. 333 """ 334 if not cfg.ssh_public_key_path: 335 logger.warning("ssh_public_key_path is not specified in acloud config. " 336 "Project-wide public key will " 337 "be used when creating AVD instances. " 338 "Please ensure you have the correct private half of " 339 "a project-wide public key if you want to ssh into the " 340 "instances after creation.") 341 elif cfg.ssh_public_key_path and not cfg.ssh_private_key_path: 342 logger.warning("Only ssh_public_key_path is specified in acloud config," 343 " but ssh_private_key_path is missing. " 344 "Please ensure you have the correct private half " 345 "if you want to ssh into the instances after creation.") 346 elif cfg.ssh_public_key_path and cfg.ssh_private_key_path: 347 utils.CreateSshKeyPairIfNotExist( 348 cfg.ssh_private_key_path, cfg.ssh_public_key_path) 349 else: 350 # Should never reach here. 351 raise errors.DriverError( 352 "Unexpected error in _CreateSshKeyPairIfNecessary") 353 354 355 def CreateAndroidVirtualDevices(cfg, 356 build_target=None, 357 build_id=None, 358 num=1, 359 gce_image=None, 360 local_disk_image=None, 361 cleanup=True, 362 serial_log_file=None, 363 logcat_file=None): 364 """Creates one or multiple android devices. 365 366 Args: 367 cfg: An AcloudConfig instance. 368 build_target: Target name, e.g. "gce_x86-userdebug" 369 build_id: Build id, a string, e.g. "2263051", "P2804227" 370 num: Number of devices to create. 371 gce_image: string, if given, will use this gce image 372 instead of creating a new one. 373 implies cleanup=False. 374 local_disk_image: string, path to a local disk image, e.g. 375 /tmp/avd-system.tar.gz 376 cleanup: boolean, if True clean up compute engine image and 377 disk image in storage after creating the instance. 378 serial_log_file: A path to a file where serial output should 379 be saved to. 380 logcat_file: A path to a file where logcat logs should be saved. 381 382 Returns: 383 A Report instance. 384 """ 385 r = report.Report(command="create") 386 credentials = auth.CreateCredentials(cfg, ALL_SCOPES) 387 compute_client = android_compute_client.AndroidComputeClient(cfg, 388 credentials) 389 try: 390 _CreateSshKeyPairIfNecessary(cfg) 391 device_pool = AndroidVirtualDevicePool(cfg) 392 device_pool.CreateDevices( 393 num, 394 build_target, 395 build_id, 396 gce_image, 397 local_disk_image, 398 cleanup, 399 extra_data_disk_size_gb=cfg.extra_data_disk_size_gb, 400 precreated_data_image=cfg.precreated_data_image_map.get( 401 cfg.extra_data_disk_size_gb)) 402 failures = device_pool.WaitForBoot() 403 # Write result to report. 404 for device in device_pool.devices: 405 device_dict = {"ip": device.ip, 406 "instance_name": device.instance_name} 407 if device.instance_name in failures: 408 r.AddData(key="devices_failing_boot", value=device_dict) 409 r.AddError(str(failures[device.instance_name])) 410 else: 411 r.AddData(key="devices", value=device_dict) 412 if failures: 413 r.SetStatus(report.Status.BOOT_FAIL) 414 else: 415 r.SetStatus(report.Status.SUCCESS) 416 417 # Dump serial and logcat logs. 418 if serial_log_file: 419 _FetchSerialLogsFromDevices( 420 compute_client, 421 instance_names=[d.instance_name for d in device_pool.devices], 422 port=constants.DEFAULT_SERIAL_PORT, 423 output_file=serial_log_file) 424 if logcat_file: 425 _FetchSerialLogsFromDevices( 426 compute_client, 427 instance_names=[d.instance_name for d in device_pool.devices], 428 port=constants.LOGCAT_SERIAL_PORT, 429 output_file=logcat_file) 430 except errors.DriverError as e: 431 r.AddError(str(e)) 432 r.SetStatus(report.Status.FAIL) 433 return r 434 435 436 def DeleteAndroidVirtualDevices(cfg, instance_names): 437 """Deletes android devices. 438 439 Args: 440 cfg: An AcloudConfig instance. 441 instance_names: A list of names of the instances to delete. 442 443 Returns: 444 A Report instance. 445 """ 446 r = report.Report(command="delete") 447 credentials = auth.CreateCredentials(cfg, ALL_SCOPES) 448 compute_client = android_compute_client.AndroidComputeClient(cfg, 449 credentials) 450 try: 451 deleted, failed, error_msgs = compute_client.DeleteInstances( 452 instance_names, cfg.zone) 453 _AddDeletionResultToReport( 454 r, deleted, 455 failed, error_msgs, 456 resource_name="instance") 457 if r.status == report.Status.UNKNOWN: 458 r.SetStatus(report.Status.SUCCESS) 459 except errors.DriverError as e: 460 r.AddError(str(e)) 461 r.SetStatus(report.Status.FAIL) 462 return r 463 464 465 def _FindOldItems(items, cut_time, time_key): 466 """Finds items from |items| whose timestamp is earlier than |cut_time|. 467 468 Args: 469 items: A list of items. Each item is a dictionary represent 470 the properties of the item. It should has a key as noted 471 by time_key. 472 cut_time: A datetime.datatime object. 473 time_key: String, key for the timestamp. 474 475 Returns: 476 A list of those from |items| whose timestamp is earlier than cut_time. 477 """ 478 cleanup_list = [] 479 for item in items: 480 t = dateutil.parser.parse(item[time_key]) 481 if t < cut_time: 482 cleanup_list.append(item) 483 return cleanup_list 484 485 486 def Cleanup(cfg, expiration_mins): 487 """Cleans up stale gce images, gce instances, and disk images in storage. 488 489 Args: 490 cfg: An AcloudConfig instance. 491 expiration_mins: Integer, resources older than |expiration_mins| will 492 be cleaned up. 493 494 Returns: 495 A Report instance. 496 """ 497 r = report.Report(command="cleanup") 498 try: 499 cut_time = (datetime.datetime.now(dateutil.tz.tzlocal()) - 500 datetime.timedelta(minutes=expiration_mins)) 501 logger.info( 502 "Cleaning up any gce images/instances and cached build artifacts." 503 "in google storage that are older than %s", cut_time) 504 credentials = auth.CreateCredentials(cfg, ALL_SCOPES) 505 compute_client = android_compute_client.AndroidComputeClient( 506 cfg, credentials) 507 storage_client = gstorage_client.StorageClient(credentials) 508 509 # Cleanup expired instances 510 items = compute_client.ListInstances(zone=cfg.zone) 511 cleanup_list = [ 512 item["name"] 513 for item in _FindOldItems(items, cut_time, "creationTimestamp") 514 ] 515 logger.info("Found expired instances: %s", cleanup_list) 516 for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): 517 result = compute_client.DeleteInstances( 518 instances=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT], 519 zone=cfg.zone) 520 _AddDeletionResultToReport(r, *result, resource_name="instance") 521 522 # Cleanup expired images 523 items = compute_client.ListImages() 524 skip_list = cfg.precreated_data_image_map.viewvalues() 525 cleanup_list = [ 526 item["name"] 527 for item in _FindOldItems(items, cut_time, "creationTimestamp") 528 if item["name"] not in skip_list 529 ] 530 logger.info("Found expired images: %s", cleanup_list) 531 for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): 532 result = compute_client.DeleteImages( 533 image_names=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT]) 534 _AddDeletionResultToReport(r, *result, resource_name="image") 535 536 # Cleanup expired disks 537 # Disks should have been attached to instances with autoDelete=True. 538 # However, sometimes disks may not be auto deleted successfully. 539 items = compute_client.ListDisks(zone=cfg.zone) 540 cleanup_list = [ 541 item["name"] 542 for item in _FindOldItems(items, cut_time, "creationTimestamp") 543 if not item.get("users") 544 ] 545 logger.info("Found expired disks: %s", cleanup_list) 546 for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): 547 result = compute_client.DeleteDisks( 548 disk_names=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT], 549 zone=cfg.zone) 550 _AddDeletionResultToReport(r, *result, resource_name="disk") 551 552 # Cleanup expired google storage 553 items = storage_client.List(bucket_name=cfg.storage_bucket_name) 554 cleanup_list = [ 555 item["name"] 556 for item in _FindOldItems(items, cut_time, "timeCreated") 557 ] 558 logger.info("Found expired cached artifacts: %s", cleanup_list) 559 for i in range(0, len(cleanup_list), MAX_BATCH_CLEANUP_COUNT): 560 result = storage_client.DeleteFiles( 561 bucket_name=cfg.storage_bucket_name, 562 object_names=cleanup_list[i:i + MAX_BATCH_CLEANUP_COUNT]) 563 _AddDeletionResultToReport( 564 r, *result, resource_name="cached_build_artifact") 565 566 # Everything succeeded, write status to report. 567 if r.status == report.Status.UNKNOWN: 568 r.SetStatus(report.Status.SUCCESS) 569 except errors.DriverError as e: 570 r.AddError(str(e)) 571 r.SetStatus(report.Status.FAIL) 572 return r 573 574 575 def AddSshRsa(cfg, user, ssh_rsa_path): 576 """Add public ssh rsa key to the project. 577 578 Args: 579 cfg: An AcloudConfig instance. 580 user: the name of the user which the key belongs to. 581 ssh_rsa_path: The absolute path to public rsa key. 582 583 Returns: 584 A Report instance. 585 """ 586 r = report.Report(command="sshkey") 587 try: 588 credentials = auth.CreateCredentials(cfg, ALL_SCOPES) 589 compute_client = android_compute_client.AndroidComputeClient( 590 cfg, credentials) 591 compute_client.AddSshRsa(user, ssh_rsa_path) 592 r.SetStatus(report.Status.SUCCESS) 593 except errors.DriverError as e: 594 r.AddError(str(e)) 595 r.SetStatus(report.Status.FAIL) 596 return r 597 598 599 def CheckAccess(cfg): 600 """Check if user has access. 601 602 Args: 603 cfg: An AcloudConfig instance. 604 """ 605 credentials = auth.CreateCredentials(cfg, ALL_SCOPES) 606 compute_client = android_compute_client.AndroidComputeClient( 607 cfg, credentials) 608 logger.info("Checking if user has access to project %s", cfg.project) 609 if not compute_client.CheckAccess(): 610 logger.error("User does not have access to project %s", cfg.project) 611 # Print here so that command line user can see it. 612 print "Looks like you do not have access to %s. " % cfg.project 613 if cfg.project in cfg.no_project_access_msg_map: 614 print cfg.no_project_access_msg_map[cfg.project] 615
