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(object):
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_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_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                self._compute_client.CreateInstance(
227                    instance=instance,
228                    image_name=image_name,
229                    extra_disk_name=extra_disk_name,
230                    avd_spec=avd_spec,
231                    extra_scopes=extra_scopes)
232                ip = self._compute_client.GetInstanceIP(instance)
233                self.devices.append(avd.AndroidVirtualDevice(
234                    ip=ip, instance_name=instance))
235        finally:
236            if cleanup:
237                self._compute_client.DeleteImage(image_name)
238
239    def DeleteDevices(self):
240        """Deletes devices.
241
242        Returns:
243            A tuple, (deleted, failed, error_msgs)
244            deleted: A list of names of instances that have been deleted.
245            faild: A list of names of instances that we fail to delete.
246            error_msgs: A list of failure messages.
247        """
248        instance_names = [device.instance_name for device in self._devices]
249        return self._compute_client.DeleteInstances(instance_names,
250                                                    self._cfg.zone)
251
252    @utils.TimeExecute("Waiting for AVD to boot")
253    def WaitForBoot(self):
254        """Waits for all devices to boot up.
255
256        Returns:
257            A dictionary that contains all the failures.
258            The key is the name of the instance that fails to boot,
259            the value is an errors.DeviceBoottError object.
260        """
261        failures = {}
262        for device in self._devices:
263            try:
264                self._compute_client.WaitForBoot(device.instance_name)
265            except errors.DeviceBootError as e:
266                failures[device.instance_name] = e
267        return failures
268
269    @property
270    def devices(self):
271        """Returns a list of devices in the pool.
272
273        Returns:
274            A list of devices in the pool.
275        """
276        return self._devices
277
278
279def AddDeletionResultToReport(report_obj, deleted, failed, error_msgs,
280                              resource_name):
281    """Adds deletion result to a Report object.
282
283    This function will add the following to report.data.
284      "deleted": [
285         {"name": "resource_name", "type": "resource_name"},
286       ],
287      "failed": [
288         {"name": "resource_name", "type": "resource_name"},
289       ],
290    This function will append error_msgs to report.errors.
291
292    Args:
293        report_obj: A Report object.
294        deleted: A list of names of the resources that have been deleted.
295        failed: A list of names of the resources that we fail to delete.
296        error_msgs: A list of error message strings to be added to the report.
297        resource_name: A string, representing the name of the resource.
298    """
299    for name in deleted:
300        report_obj.AddData(key="deleted",
301                           value={"name": name,
302                                  "type": resource_name})
303    for name in failed:
304        report_obj.AddData(key="failed",
305                           value={"name": name,
306                                  "type": resource_name})
307    report_obj.AddErrors(error_msgs)
308    if failed or error_msgs:
309        report_obj.SetStatus(report.Status.FAIL)
310
311
312def _FetchSerialLogsFromDevices(compute_client, instance_names, output_file,
313                                port):
314    """Fetch serial logs from a port for a list of devices to a local file.
315
316    Args:
317        compute_client: An object of android_compute_client.AndroidComputeClient
318        instance_names: A list of instance names.
319        output_file: A path to a file ending with "tar.gz"
320        port: The number of serial port to read from, 0 for serial output, 1 for
321              logcat.
322    """
323    with utils.TempDir() as tempdir:
324        src_dict = {}
325        for instance_name in instance_names:
326            serial_log = compute_client.GetSerialPortOutput(
327                instance=instance_name, port=port)
328            file_name = "%s.log" % instance_name
329            file_path = os.path.join(tempdir, file_name)
330            src_dict[file_path] = file_name
331            with open(file_path, "w") as f:
332                f.write(serial_log.encode("utf-8"))
333        utils.MakeTarFile(src_dict, output_file)
334
335
336# pylint: disable=too-many-locals
337def CreateGCETypeAVD(cfg,
338                     build_target=None,
339                     build_id=None,
340                     num=1,
341                     gce_image=None,
342                     local_disk_image=None,
343                     cleanup=True,
344                     serial_log_file=None,
345                     autoconnect=False,
346                     report_internal_ip=False,
347                     avd_spec=None):
348    """Creates one or multiple gce android devices.
349
350    Args:
351        cfg: An AcloudConfig instance.
352        build_target: Target name, e.g. "aosp_cf_x86_phone-userdebug"
353        build_id: Build id, a string, e.g. "2263051", "P2804227"
354        num: Number of devices to create.
355        gce_image: string, if given, will use this gce image
356                   instead of creating a new one.
357                   implies cleanup=False.
358        local_disk_image: string, path to a local disk image, e.g.
359                          /tmp/avd-system.tar.gz
360        cleanup: boolean, if True clean up compute engine image and
361                 disk image in storage after creating the instance.
362        serial_log_file: A path to a file where serial output should
363                         be saved to.
364        autoconnect: Create ssh tunnel(s) and adb connect after device creation.
365        report_internal_ip: Boolean to report the internal ip instead of
366                            external ip.
367        avd_spec: AVDSpec object for pass hw_property.
368
369    Returns:
370        A Report instance.
371    """
372    r = report.Report(command="create")
373    credentials = auth.CreateCredentials(cfg)
374    compute_client = android_compute_client.AndroidComputeClient(cfg,
375                                                                 credentials)
376    try:
377        common_operations.CreateSshKeyPairIfNecessary(cfg)
378        device_pool = AndroidVirtualDevicePool(cfg)
379        device_pool.CreateDevices(
380            num,
381            build_target,
382            build_id,
383            gce_image,
384            local_disk_image,
385            cleanup,
386            extra_data_disk_size_gb=cfg.extra_data_disk_size_gb,
387            precreated_data_image=cfg.precreated_data_image_map.get(
388                cfg.extra_data_disk_size_gb),
389            avd_spec=avd_spec,
390            extra_scopes=cfg.extra_scopes)
391        failures = device_pool.WaitForBoot()
392        # Write result to report.
393        for device in device_pool.devices:
394            ip = (device.ip.internal if report_internal_ip
395                  else device.ip.external)
396            device_dict = {
397                "ip": ip,
398                "instance_name": device.instance_name
399            }
400            if autoconnect:
401                forwarded_ports = utils.AutoConnect(
402                    ip_addr=ip,
403                    rsa_key_file=cfg.ssh_private_key_path,
404                    target_vnc_port=constants.GCE_VNC_PORT,
405                    target_adb_port=constants.GCE_ADB_PORT,
406                    ssh_user=_SSH_USER,
407                    client_adb_port=avd_spec.client_adb_port,
408                    extra_args_ssh_tunnel=cfg.extra_args_ssh_tunnel)
409                device_dict[constants.VNC_PORT] = forwarded_ports.vnc_port
410                device_dict[constants.ADB_PORT] = forwarded_ports.adb_port
411                if avd_spec.unlock_screen:
412                    AdbTools(forwarded_ports.adb_port).AutoUnlockScreen()
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 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    except errors.DriverError as e:
431        r.AddError(str(e))
432        r.SetStatus(report.Status.FAIL)
433    return r
434
435
436def DeleteAndroidVirtualDevices(cfg, instance_names, default_report=None):
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        default_report: A initialized Report instance.
443
444    Returns:
445        A Report instance.
446    """
447    # delete, failed, error_msgs are used to record result.
448    deleted = []
449    failed = []
450    error_msgs = []
451
452    r = default_report if default_report else report.Report(command="delete")
453    credentials = auth.CreateCredentials(cfg)
454    compute_client = android_compute_client.AndroidComputeClient(cfg,
455                                                                 credentials)
456    zone_instances = compute_client.GetZonesByInstances(instance_names)
457
458    try:
459        for zone, instances in zone_instances.items():
460            deleted_ins, failed_ins, error_ins = compute_client.DeleteInstances(
461                instances, zone)
462            deleted.extend(deleted_ins)
463            failed.extend(failed_ins)
464            error_msgs.extend(error_ins)
465        AddDeletionResultToReport(
466            r, deleted,
467            failed, error_msgs,
468            resource_name="instance")
469        if r.status == report.Status.UNKNOWN:
470            r.SetStatus(report.Status.SUCCESS)
471    except errors.DriverError as e:
472        r.AddError(str(e))
473        r.SetStatus(report.Status.FAIL)
474    return r
475
476
477def CheckAccess(cfg):
478    """Check if user has access.
479
480    Args:
481         cfg: An AcloudConfig instance.
482    """
483    credentials = auth.CreateCredentials(cfg)
484    compute_client = android_compute_client.AndroidComputeClient(
485        cfg, credentials)
486    logger.info("Checking if user has access to project %s", cfg.project)
487    if not compute_client.CheckAccess():
488        logger.error("User does not have access to project %s", cfg.project)
489        # Print here so that command line user can see it.
490        print("Looks like you do not have access to %s. " % cfg.project)
491        if cfg.project in cfg.no_project_access_msg_map:
492            print(cfg.no_project_access_msg_map[cfg.project])
493