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 logging 32import os 33 34from acloud import errors 35from acloud.public import avd 36from acloud.public import report 37from acloud.public.actions import common_operations 38from acloud.internal import constants 39from acloud.internal.lib import auth 40from acloud.internal.lib import android_build_client 41from acloud.internal.lib import android_compute_client 42from acloud.internal.lib import gstorage_client 43from acloud.internal.lib import utils 44from acloud.internal.lib.adb_tools import AdbTools 45 46 47logger = logging.getLogger(__name__) 48 49MAX_BATCH_CLEANUP_COUNT = 100 50 51_SSH_USER = "root" 52 53 54# pylint: disable=invalid-name 55class AndroidVirtualDevicePool(): 56 """A class that manages a pool of devices.""" 57 58 def __init__(self, cfg, devices=None): 59 self._devices = devices or [] 60 self._cfg = cfg 61 credentials = auth.CreateCredentials(cfg) 62 self._build_client = android_build_client.AndroidBuildClient( 63 credentials) 64 self._storage_client = gstorage_client.StorageClient(credentials) 65 self._compute_client = android_compute_client.AndroidComputeClient( 66 cfg, credentials) 67 68 @utils.TimeExecute("Creating GCE image") 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. "aosp_cf_x86_64_phone-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 @utils.TimeExecute("Creating GCE image") 105 def _CreateGceImageWithLocalFile(self, local_disk_image): 106 """Create a Gce image with a local image file. 107 108 The local disk image can be either a tar.gz file or a 109 raw vmlinux image. 110 e.g. /tmp/avd-system.tar.gz or /tmp/android_system_disk_syslinux.img 111 If a raw vmlinux image is provided, it will be archived into a tar.gz file. 112 113 The final tar.gz file will be uploaded to a cache bucket in storage. 114 115 Args: 116 local_disk_image: string, path to a local disk image, 117 118 Returns: 119 String, name of the Gce image that has been created. 120 121 Raises: 122 DriverError: if a file with an unexpected extension is given. 123 """ 124 logger.info("Creating a new gce image from a local file %s", 125 local_disk_image) 126 with utils.TempDir() as tempdir: 127 if local_disk_image.endswith(self._cfg.disk_raw_image_extension): 128 dest_tar_file = os.path.join(tempdir, 129 self._cfg.disk_image_name) 130 utils.MakeTarFile( 131 src_dict={local_disk_image: self._cfg.disk_raw_image_name}, 132 dest=dest_tar_file) 133 local_disk_image = dest_tar_file 134 elif not local_disk_image.endswith(self._cfg.disk_image_extension): 135 raise errors.DriverError( 136 "Wrong local_disk_image type, must be a *%s file or *%s file" 137 % (self._cfg.disk_raw_image_extension, 138 self._cfg.disk_image_extension)) 139 140 disk_image_id = utils.GenerateUniqueName( 141 suffix=self._cfg.disk_image_name) 142 self._storage_client.Upload( 143 local_src=local_disk_image, 144 bucket_name=self._cfg.storage_bucket_name, 145 object_name=disk_image_id, 146 mime_type=self._cfg.disk_image_mime_type) 147 disk_image_url = self._storage_client.GetUrl( 148 self._cfg.storage_bucket_name, disk_image_id) 149 try: 150 image_name = self._compute_client.GenerateImageName() 151 self._compute_client.CreateImage(image_name=image_name, 152 source_uri=disk_image_url) 153 finally: 154 self._storage_client.Delete(self._cfg.storage_bucket_name, 155 disk_image_id) 156 return image_name 157 158 # pylint: disable=too-many-locals 159 def CreateDevices(self, 160 num, 161 build_target=None, 162 build_id=None, 163 gce_image=None, 164 local_disk_image=None, 165 cleanup=True, 166 extra_data_disk_size_gb=None, 167 precreated_data_image=None, 168 avd_spec=None, 169 extra_scopes=None): 170 """Creates |num| devices for given build_target and build_id. 171 172 - If gce_image is provided, will use it to create an instance. 173 - If local_disk_image is provided, will upload it to a temporary 174 caching storage bucket which is defined by user as |storage_bucket_name| 175 And then create an gce image with it; and then create an instance. 176 - If build_target and build_id are provided, will clone the disk image 177 via launch control to the temporary caching storage bucket. 178 And then create an gce image with it; and then create an instance. 179 180 Args: 181 num: Number of devices to create. 182 build_target: Target name, e.g. "aosp_cf_x86_64_phone-userdebug" 183 build_id: Build id, a string, e.g. "2263051", "P2804227" 184 gce_image: string, if given, will use this image 185 instead of creating a new one. 186 implies cleanup=False. 187 local_disk_image: string, path to a local disk image, e.g. 188 /tmp/avd-system.tar.gz 189 cleanup: boolean, if True clean up compute engine image after creating 190 the instance. 191 extra_data_disk_size_gb: Integer, size of extra disk, or None. 192 precreated_data_image: A string, the image to use for the extra disk. 193 avd_spec: AVDSpec object for pass hw_property. 194 extra_scopes: A list of extra scopes given to the new instance. 195 196 Raises: 197 errors.DriverError: If no source is specified for image creation. 198 """ 199 if gce_image: 200 # GCE image is provided, we can directly move to instance creation. 201 logger.info("Using existing gce image %s", gce_image) 202 image_name = gce_image 203 cleanup = False 204 elif local_disk_image: 205 image_name = self._CreateGceImageWithLocalFile(local_disk_image) 206 elif build_target and build_id: 207 image_name = self._CreateGceImageWithBuildInfo(build_target, 208 build_id) 209 else: 210 raise errors.DriverError( 211 "Invalid image source, must specify one of the following: gce_image, " 212 "local_disk_image, or build_target and build id.") 213 214 # Create GCE instances. 215 try: 216 for _ in range(num): 217 instance = self._compute_client.GenerateInstanceName( 218 build_target, build_id) 219 extra_disk_name = None 220 if extra_data_disk_size_gb > 0: 221 extra_disk_name = self._compute_client.GetDataDiskName( 222 instance) 223 self._compute_client.CreateDisk(extra_disk_name, 224 precreated_data_image, 225 extra_data_disk_size_gb, 226 disk_type=avd_spec.disk_type) 227 self._compute_client.CreateInstance( 228 instance=instance, 229 image_name=image_name, 230 extra_disk_name=extra_disk_name, 231 avd_spec=avd_spec, 232 extra_scopes=extra_scopes) 233 ip = self._compute_client.GetInstanceIP(instance) 234 self.devices.append(avd.AndroidVirtualDevice( 235 ip=ip, instance_name=instance)) 236 finally: 237 if cleanup: 238 self._compute_client.DeleteImage(image_name) 239 240 def DeleteDevices(self): 241 """Deletes devices. 242 243 Returns: 244 A tuple, (deleted, failed, error_msgs) 245 deleted: A list of names of instances that have been deleted. 246 faild: A list of names of instances that we fail to delete. 247 error_msgs: A list of failure messages. 248 """ 249 instance_names = [device.instance_name for device in self._devices] 250 return self._compute_client.DeleteInstances(instance_names, 251 self._cfg.zone) 252 253 @utils.TimeExecute("Waiting for AVD to boot") 254 def WaitForBoot(self): 255 """Waits for all devices to boot up. 256 257 Returns: 258 A dictionary that contains all the failures. 259 The key is the name of the instance that fails to boot, 260 the value is an errors.DeviceBoottError object. 261 """ 262 failures = {} 263 for device in self._devices: 264 try: 265 self._compute_client.WaitForBoot(device.instance_name) 266 except errors.DeviceBootError as e: 267 failures[device.instance_name] = e 268 return failures 269 270 @property 271 def devices(self): 272 """Returns a list of devices in the pool. 273 274 Returns: 275 A list of devices in the pool. 276 """ 277 return self._devices 278 279 280def AddDeletionResultToReport(report_obj, deleted, failed, error_msgs, 281 resource_name): 282 """Adds deletion result to a Report object. 283 284 This function will add the following to report.data. 285 "deleted": [ 286 {"name": "resource_name", "type": "resource_name"}, 287 ], 288 "failed": [ 289 {"name": "resource_name", "type": "resource_name"}, 290 ], 291 This function will append error_msgs to report.errors. 292 293 Args: 294 report_obj: A Report object. 295 deleted: A list of names of the resources that have been deleted. 296 failed: A list of names of the resources that we fail to delete. 297 error_msgs: A list of error message strings to be added to the report. 298 resource_name: A string, representing the name of the resource. 299 """ 300 for name in deleted: 301 report_obj.AddData(key="deleted", 302 value={"name": name, 303 "type": resource_name}) 304 for name in failed: 305 report_obj.AddData(key="failed", 306 value={"name": name, 307 "type": resource_name}) 308 report_obj.AddErrors(error_msgs) 309 if failed or error_msgs: 310 report_obj.SetStatus(report.Status.FAIL) 311 312 313def _FetchSerialLogsFromDevices(compute_client, instance_names, output_file, 314 port): 315 """Fetch serial logs from a port for a list of devices to a local file. 316 317 Args: 318 compute_client: An object of android_compute_client.AndroidComputeClient 319 instance_names: A list of instance names. 320 output_file: A path to a file ending with "tar.gz" 321 port: The number of serial port to read from, 0 for serial output, 1 for 322 logcat. 323 """ 324 with utils.TempDir() as tempdir: 325 src_dict = {} 326 for instance_name in instance_names: 327 serial_log = compute_client.GetSerialPortOutput( 328 instance=instance_name, port=port) 329 file_name = "%s.log" % instance_name 330 file_path = os.path.join(tempdir, file_name) 331 src_dict[file_path] = file_name 332 with open(file_path, "w") as f: 333 f.write(serial_log.encode("utf-8")) 334 utils.MakeTarFile(src_dict, output_file) 335 336 337# pylint: disable=too-many-locals 338def CreateGCETypeAVD(cfg, 339 build_target=None, 340 build_id=None, 341 num=1, 342 gce_image=None, 343 local_disk_image=None, 344 cleanup=True, 345 serial_log_file=None, 346 autoconnect=False, 347 report_internal_ip=False, 348 avd_spec=None): 349 """Creates one or multiple gce android devices. 350 351 Args: 352 cfg: An AcloudConfig instance. 353 build_target: Target name, e.g. "aosp_cf_x86_64_phone-userdebug" 354 build_id: Build id, a string, e.g. "2263051", "P2804227" 355 num: Number of devices to create. 356 gce_image: string, if given, will use this gce image 357 instead of creating a new one. 358 implies cleanup=False. 359 local_disk_image: string, path to a local disk image, e.g. 360 /tmp/avd-system.tar.gz 361 cleanup: boolean, if True clean up compute engine image and 362 disk image in storage after creating the instance. 363 serial_log_file: A path to a file where serial output should 364 be saved to. 365 autoconnect: Create ssh tunnel(s) and adb connect after device creation. 366 report_internal_ip: Boolean to report the internal ip instead of 367 external ip. 368 avd_spec: AVDSpec object for pass hw_property. 369 370 Returns: 371 A Report instance. 372 """ 373 r = report.Report(command="create") 374 credentials = auth.CreateCredentials(cfg) 375 compute_client = android_compute_client.AndroidComputeClient(cfg, 376 credentials) 377 try: 378 common_operations.CreateSshKeyPairIfNecessary(cfg) 379 device_pool = AndroidVirtualDevicePool(cfg) 380 device_pool.CreateDevices( 381 num, 382 build_target, 383 build_id, 384 gce_image, 385 local_disk_image, 386 cleanup, 387 extra_data_disk_size_gb=cfg.extra_data_disk_size_gb, 388 precreated_data_image=cfg.precreated_data_image_map.get( 389 cfg.extra_data_disk_size_gb), 390 avd_spec=avd_spec, 391 extra_scopes=cfg.extra_scopes) 392 failures = device_pool.WaitForBoot() 393 # Write result to report. 394 for device in device_pool.devices: 395 ip = (device.ip.internal if report_internal_ip 396 else device.ip.external) 397 device_dict = { 398 "ip": ip, 399 "instance_name": device.instance_name 400 } 401 if autoconnect: 402 forwarded_ports = utils.AutoConnect( 403 ip_addr=ip, 404 rsa_key_file=cfg.ssh_private_key_path, 405 target_vnc_port=constants.GCE_VNC_PORT, 406 target_adb_port=constants.GCE_ADB_PORT, 407 ssh_user=_SSH_USER, 408 client_adb_port=avd_spec.client_adb_port, 409 extra_args_ssh_tunnel=cfg.extra_args_ssh_tunnel) 410 device_dict[constants.VNC_PORT] = forwarded_ports.vnc_port 411 device_dict[constants.ADB_PORT] = forwarded_ports.adb_port 412 if avd_spec.unlock_screen: 413 AdbTools(forwarded_ports.adb_port).AutoUnlockScreen() 414 if device.instance_name in failures: 415 r.AddData(key="devices_failing_boot", value=device_dict) 416 r.AddError(str(failures[device.instance_name])) 417 else: 418 r.AddData(key="devices", value=device_dict) 419 if failures: 420 r.SetStatus(report.Status.BOOT_FAIL) 421 else: 422 r.SetStatus(report.Status.SUCCESS) 423 424 # Dump serial logs. 425 if serial_log_file: 426 _FetchSerialLogsFromDevices( 427 compute_client, 428 instance_names=[d.instance_name for d in device_pool.devices], 429 port=constants.DEFAULT_SERIAL_PORT, 430 output_file=serial_log_file) 431 except errors.DriverError as e: 432 r.AddError(str(e)) 433 r.SetStatus(report.Status.FAIL) 434 return r 435 436 437def DeleteAndroidVirtualDevices(cfg, instance_names, default_report=None): 438 """Deletes android devices. 439 440 Args: 441 cfg: An AcloudConfig instance. 442 instance_names: A list of names of the instances to delete. 443 default_report: A initialized Report instance. 444 445 Returns: 446 A Report instance. 447 """ 448 # delete, failed, error_msgs are used to record result. 449 deleted = [] 450 failed = [] 451 error_msgs = [] 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 zone_instances = compute_client.GetZonesByInstances(instance_names) 458 459 try: 460 for zone, instances in zone_instances.items(): 461 deleted_ins, failed_ins, error_ins = compute_client.DeleteInstances( 462 instances, zone) 463 deleted.extend(deleted_ins) 464 failed.extend(failed_ins) 465 error_msgs.extend(error_ins) 466 AddDeletionResultToReport( 467 r, deleted, 468 failed, error_msgs, 469 resource_name="instance") 470 if r.status == report.Status.UNKNOWN: 471 r.SetStatus(report.Status.SUCCESS) 472 except errors.DriverError as e: 473 r.AddError(str(e)) 474 r.SetStatus(report.Status.FAIL) 475 return r 476 477 478def CheckAccess(cfg): 479 """Check if user has access. 480 481 Args: 482 cfg: An AcloudConfig instance. 483 """ 484 credentials = auth.CreateCredentials(cfg) 485 compute_client = android_compute_client.AndroidComputeClient( 486 cfg, credentials) 487 logger.info("Checking if user has access to project %s", cfg.project) 488 if not compute_client.CheckAccess(): 489 logger.error("User does not have access to project %s", cfg.project) 490 # Print here so that command line user can see it. 491 print("Looks like you do not have access to %s. " % cfg.project) 492 if cfg.project in cfg.no_project_access_msg_map: 493 print(cfg.no_project_access_msg_map[cfg.project]) 494