1#!/usr/bin/env python
2#
3# Copyright 2018 - 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"""Common operations between managing GCE and Cuttlefish devices.
17
18This module provides the common operations between managing GCE (device_driver)
19and Cuttlefish (create_cuttlefish_action) devices. Should not be called
20directly.
21"""
22
23import logging
24import os
25
26from acloud import errors
27from acloud.public import avd
28from acloud.public import report
29from acloud.internal import constants
30from acloud.internal.lib import utils
31from acloud.internal.lib.adb_tools import AdbTools
32
33
34logger = logging.getLogger(__name__)
35
36
37def CreateSshKeyPairIfNecessary(cfg):
38    """Create ssh key pair if necessary.
39
40    Args:
41        cfg: An Acloudconfig instance.
42
43    Raises:
44        error.DriverError: If it falls into an unexpected condition.
45    """
46    if not cfg.ssh_public_key_path:
47        logger.warning(
48            "ssh_public_key_path is not specified in acloud config. "
49            "Project-wide public key will "
50            "be used when creating AVD instances. "
51            "Please ensure you have the correct private half of "
52            "a project-wide public key if you want to ssh into the "
53            "instances after creation.")
54    elif cfg.ssh_public_key_path and not cfg.ssh_private_key_path:
55        logger.warning(
56            "Only ssh_public_key_path is specified in acloud config, "
57            "but ssh_private_key_path is missing. "
58            "Please ensure you have the correct private half "
59            "if you want to ssh into the instances after creation.")
60    elif cfg.ssh_public_key_path and cfg.ssh_private_key_path:
61        utils.CreateSshKeyPairIfNotExist(cfg.ssh_private_key_path,
62                                         cfg.ssh_public_key_path)
63    else:
64        # Should never reach here.
65        raise errors.DriverError(
66            "Unexpected error in CreateSshKeyPairIfNecessary")
67
68
69class DevicePool(object):
70    """A class that manages a pool of virtual devices.
71
72    Attributes:
73        devices: A list of devices in the pool.
74    """
75
76    def __init__(self, device_factory, devices=None):
77        """Constructs a new DevicePool.
78
79        Args:
80            device_factory: A device factory capable of producing a goldfish or
81                cuttlefish device. The device factory must expose an attribute with
82                the credentials that can be used to retrieve information from the
83                constructed device.
84            devices: List of devices managed by this pool.
85        """
86        self._devices = devices or []
87        self._device_factory = device_factory
88        self._compute_client = device_factory.GetComputeClient()
89
90    def CreateDevices(self, num):
91        """Creates |num| devices for given build_target and build_id.
92
93        Args:
94            num: Number of devices to create.
95        """
96        # Create host instances for cuttlefish/goldfish device.
97        # Currently one instance supports only 1 device.
98        for _ in range(num):
99            instance = self._device_factory.CreateInstance()
100            ip = self._compute_client.GetInstanceIP(instance)
101            time_info = self._compute_client.execution_time if hasattr(
102                self._compute_client, "execution_time") else {}
103            self.devices.append(
104                avd.AndroidVirtualDevice(ip=ip, instance_name=instance,
105                                         time_info=time_info))
106
107    @utils.TimeExecute(function_description="Waiting for AVD(s) to boot up",
108                       result_evaluator=utils.BootEvaluator)
109    def WaitForBoot(self, boot_timeout_secs):
110        """Waits for all devices to boot up.
111
112        Args:
113            boot_timeout_secs: Integer, the maximum time in seconds used to
114                               wait for the AVD to boot.
115
116        Returns:
117            A dictionary that contains all the failures.
118            The key is the name of the instance that fails to boot,
119            and the value is an errors.DeviceBootError object.
120        """
121        failures = {}
122        for device in self._devices:
123            try:
124                self._compute_client.WaitForBoot(device.instance_name, boot_timeout_secs)
125            except errors.DeviceBootError as e:
126                failures[device.instance_name] = e
127        return failures
128
129    def UpdateReport(self, reporter):
130        """Update report from compute client.
131
132        Args:
133            reporter: Report object.
134        """
135        reporter.UpdateData(self._compute_client.dict_report)
136
137    def CollectSerialPortLogs(self, output_file,
138                              port=constants.DEFAULT_SERIAL_PORT):
139        """Tar the instance serial logs into specified output_file.
140
141        Args:
142            output_file: String, the output tar file path
143            port: The serial port number to be collected
144        """
145        # For emulator, the serial log is the virtual host serial log.
146        # For GCE AVD device, the serial log is the AVD device serial log.
147        with utils.TempDir() as tempdir:
148            src_dict = {}
149            for device in self._devices:
150                logger.info("Store instance %s serial port %s output to %s",
151                            device.instance_name, port, output_file)
152                serial_log = self._compute_client.GetSerialPortOutput(
153                    instance=device.instance_name, port=port)
154                file_name = "%s_serial_%s.log" % (device.instance_name, port)
155                file_path = os.path.join(tempdir, file_name)
156                src_dict[file_path] = file_name
157                with open(file_path, "w") as f:
158                    f.write(serial_log.encode("utf-8"))
159            utils.MakeTarFile(src_dict, output_file)
160
161    def SetDeviceBuildInfo(self):
162        """Add devices build info."""
163        for device in self._devices:
164            device.build_info = self._device_factory.GetBuildInfoDict()
165
166    @property
167    def devices(self):
168        """Returns a list of devices in the pool.
169
170        Returns:
171            A list of devices in the pool.
172        """
173        return self._devices
174
175# pylint: disable=too-many-locals,unused-argument,too-many-branches
176def CreateDevices(command, cfg, device_factory, num, avd_type,
177                  report_internal_ip=False, autoconnect=False,
178                  serial_log_file=None, client_adb_port=None,
179                  boot_timeout_secs=None, unlock_screen=False,
180                  wait_for_boot=True, connect_webrtc=False):
181    """Create a set of devices using the given factory.
182
183    Main jobs in create devices.
184        1. Create GCE instance: Launch instance in GCP(Google Cloud Platform).
185        2. Starting up AVD: Wait device boot up.
186
187    Args:
188        command: The name of the command, used for reporting.
189        cfg: An AcloudConfig instance.
190        device_factory: A factory capable of producing a single device.
191        num: The number of devices to create.
192        avd_type: String, the AVD type(cuttlefish, goldfish...).
193        report_internal_ip: Boolean to report the internal ip instead of
194                            external ip.
195        serial_log_file: String, the file path to tar the serial logs.
196        autoconnect: Boolean, whether to auto connect to device.
197        client_adb_port: Integer, Specify port for adb forwarding.
198        boot_timeout_secs: Integer, boot timeout secs.
199        unlock_screen: Boolean, whether to unlock screen after invoke vnc client.
200        wait_for_boot: Boolean, True to check serial log include boot up
201                       message.
202        connect_webrtc: Boolean, whether to auto connect webrtc to device.
203
204    Raises:
205        errors: Create instance fail.
206
207    Returns:
208        A Report instance.
209    """
210    reporter = report.Report(command=command)
211    try:
212        CreateSshKeyPairIfNecessary(cfg)
213        device_pool = DevicePool(device_factory)
214        device_pool.CreateDevices(num)
215        device_pool.SetDeviceBuildInfo()
216        if wait_for_boot:
217            failures = device_pool.WaitForBoot(boot_timeout_secs)
218        else:
219            failures = device_factory.GetFailures()
220
221        if failures:
222            reporter.SetStatus(report.Status.BOOT_FAIL)
223        else:
224            reporter.SetStatus(report.Status.SUCCESS)
225
226        # Collect logs
227        if serial_log_file:
228            device_pool.CollectSerialPortLogs(
229                serial_log_file, port=constants.DEFAULT_SERIAL_PORT)
230
231        device_pool.UpdateReport(reporter)
232        # Write result to report.
233        for device in device_pool.devices:
234            ip = (device.ip.internal if report_internal_ip
235                  else device.ip.external)
236            device_dict = {
237                "ip": ip,
238                "instance_name": device.instance_name
239            }
240            if device.build_info:
241                device_dict.update(device.build_info)
242            if device.time_info:
243                device_dict.update(device.time_info)
244            if autoconnect:
245                forwarded_ports = utils.AutoConnect(
246                    ip_addr=ip,
247                    rsa_key_file=cfg.ssh_private_key_path,
248                    target_vnc_port=utils.AVD_PORT_DICT[avd_type].vnc_port,
249                    target_adb_port=utils.AVD_PORT_DICT[avd_type].adb_port,
250                    ssh_user=constants.GCE_USER,
251                    client_adb_port=client_adb_port,
252                    extra_args_ssh_tunnel=cfg.extra_args_ssh_tunnel)
253                device_dict[constants.VNC_PORT] = forwarded_ports.vnc_port
254                device_dict[constants.ADB_PORT] = forwarded_ports.adb_port
255                if unlock_screen:
256                    AdbTools(forwarded_ports.adb_port).AutoUnlockScreen()
257            if connect_webrtc:
258                utils.EstablishWebRTCSshTunnel(
259                    ip_addr=ip,
260                    rsa_key_file=cfg.ssh_private_key_path,
261                    ssh_user=constants.GCE_USER,
262                    extra_args_ssh_tunnel=cfg.extra_args_ssh_tunnel)
263            if device.instance_name in failures:
264                reporter.AddData(key="devices_failing_boot", value=device_dict)
265                reporter.AddError(str(failures[device.instance_name]))
266            else:
267                reporter.AddData(key="devices", value=device_dict)
268    except errors.DriverError as e:
269        reporter.AddError(str(e))
270        reporter.SetStatus(report.Status.FAIL)
271    return reporter
272