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