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"""A client that manages Android compute engine instances.
18
19** AndroidComputeClient **
20
21AndroidComputeClient derives from ComputeClient. It manges a google
22compute engine project that is setup for running Android instances.
23It knows how to create android GCE images and instances.
24
25** Class hierarchy **
26
27  base_cloud_client.BaseCloudApiClient
28                ^
29                |
30       gcompute_client.ComputeClient
31                ^
32                |
33    gcompute_client.AndroidComputeClient
34
35TODO(fdeng):
36  Merge caci/framework/gce_manager.py
37  with this module, update callers of gce_manager.py to use this module.
38"""
39
40import getpass
41import logging
42import os
43import uuid
44
45from acloud.internal.lib import gcompute_client
46from acloud.internal.lib import utils
47from acloud.public import errors
48
49logger = logging.getLogger(__name__)
50
51
52class AndroidComputeClient(gcompute_client.ComputeClient):
53    """Client that manages Anadroid Virtual Device."""
54
55    INSTANCE_NAME_FMT = "ins-{uuid}-{build_id}-{build_target}"
56    IMAGE_NAME_FMT = "img-{uuid}-{build_id}-{build_target}"
57    DATA_DISK_NAME_FMT = "data-{instance}"
58    BOOT_COMPLETED_MSG = "VIRTUAL_DEVICE_BOOT_COMPLETED"
59    BOOT_TIMEOUT_SECS = 5 * 60  # 5 mins, usually it should take ~2 mins
60    BOOT_CHECK_INTERVAL_SECS = 10
61    NAME_LENGTH_LIMIT = 63
62    # If the generated name ends with '-', replace it with REPLACER.
63    REPLACER = "e"
64
65    def __init__(self, acloud_config, oauth2_credentials):
66        """Initialize.
67
68        Args:
69            acloud_config: An AcloudConfig object.
70            oauth2_credentials: An oauth2client.OAuth2Credentials instance.
71        """
72        super(AndroidComputeClient, self).__init__(acloud_config,
73                                                   oauth2_credentials)
74        self._zone = acloud_config.zone
75        self._machine_type = acloud_config.machine_type
76        self._min_machine_size = acloud_config.min_machine_size
77        self._network = acloud_config.network
78        self._orientation = acloud_config.orientation
79        self._resolution = acloud_config.resolution
80        self._metadata = acloud_config.metadata_variable.copy()
81        self._ssh_public_key_path = acloud_config.ssh_public_key_path
82
83    @classmethod
84    def _FormalizeName(cls, name):
85        """Formalize the name to comply with RFC1035.
86
87        The name must be 1-63 characters long and match the regular expression
88        [a-z]([-a-z0-9]*[a-z0-9])? which means the first character must be a
89        lowercase letter, and all following characters must be a dash,
90        lowercase letter, or digit, except the last character, which cannot be
91        a dash.
92
93        Args:
94          name: A string.
95
96        Returns:
97          name: A string that complies with RFC1035.
98        """
99        name = name.replace("_", "-").lower()
100        name = name[:cls.NAME_LENGTH_LIMIT]
101        if name[-1] == "-":
102          name = name[:-1] + cls.REPLACER
103        return name
104
105    def _CheckMachineSize(self):
106        """Check machine size.
107
108        Check if the desired machine type |self._machine_type| meets
109        the requirement of minimum machine size specified as
110        |self._min_machine_size|.
111
112        Raises:
113            errors.DriverError: if check fails.
114        """
115        if self.CompareMachineSize(self._machine_type, self._min_machine_size,
116                                   self._zone) < 0:
117            raise errors.DriverError(
118                "%s does not meet the minimum required machine size %s" %
119                (self._machine_type, self._min_machine_size))
120
121    @classmethod
122    def GenerateImageName(cls, build_target=None, build_id=None):
123        """Generate an image name given build_target, build_id.
124
125        Args:
126            build_target: Target name, e.g. "gce_x86-userdebug"
127            build_id: Build id, a string, e.g. "2263051", "P2804227"
128
129        Returns:
130            A string, representing image name.
131        """
132        if not build_target and not build_id:
133            return "image-" + uuid.uuid4().hex
134        name = cls.IMAGE_NAME_FMT.format(build_target=build_target,
135                                         build_id=build_id,
136                                         uuid=uuid.uuid4().hex[:8])
137        return cls._FormalizeName(name)
138
139    @classmethod
140    def GetDataDiskName(cls, instance):
141        """Get data disk name for an instance.
142
143        Args:
144            instance: An instance_name.
145
146        Returns:
147            The corresponding data disk name.
148        """
149        name = cls.DATA_DISK_NAME_FMT.format(instance=instance)
150        return cls._FormalizeName(name)
151
152    @classmethod
153    def GenerateInstanceName(cls, build_target=None, build_id=None):
154        """Generate an instance name given build_target, build_id.
155
156        Target is not used as instance name has a length limit.
157
158        Args:
159            build_target: Target name, e.g. "gce_x86-userdebug"
160            build_id: Build id, a string, e.g. "2263051", "P2804227"
161
162        Returns:
163            A string, representing instance name.
164        """
165        if not build_target and not build_id:
166            return "instance-" + uuid.uuid4().hex
167        name = cls.INSTANCE_NAME_FMT.format(
168            build_target=build_target,
169            build_id=build_id,
170            uuid=uuid.uuid4().hex[:8]).replace("_", "-").lower()
171        return cls._FormalizeName(name)
172
173    def CreateDisk(self, disk_name, source_image, size_gb):
174        """Create a gce disk.
175
176        Args:
177            disk_name: A string.
178            source_image: A string, name to the image name.
179            size_gb: Integer, size in gigabytes.
180        """
181        if self.CheckDiskExists(disk_name, self._zone):
182            raise errors.DriverError(
183                "Failed to create disk %s, already exists." % disk_name)
184        if source_image and not self.CheckImageExists(source_image):
185            raise errors.DriverError(
186                "Failed to create disk %s, source image %s does not exist." %
187                (disk_name, source_image))
188        super(AndroidComputeClient, self).CreateDisk(disk_name,
189                                                     source_image=source_image,
190                                                     size_gb=size_gb,
191                                                     zone=self._zone)
192
193    def CreateImage(self, image_name, source_uri):
194        """Create a gce image.
195
196        Args:
197            image_name: String, name of the image.
198            source_uri: A full Google Storage URL to the disk image.
199                        e.g. "https://storage.googleapis.com/my-bucket/
200                              avd-system-2243663.tar.gz"
201        """
202        if not self.CheckImageExists(image_name):
203            super(AndroidComputeClient, self).CreateImage(image_name,
204                                                          source_uri)
205
206    def _GetExtraDiskArgs(self, extra_disk_name):
207        """Get extra disk arg for given disk.
208
209        Args:
210            extra_disk_name: Name of the disk.
211
212        Returns:
213            A dictionary of disk args.
214        """
215        return [{
216            "type": "PERSISTENT",
217            "mode": "READ_WRITE",
218            "source": "projects/%s/zones/%s/disks/%s" % (
219                self._project, self._zone, extra_disk_name),
220            "autoDelete": True,
221            "boot": False,
222            "interface": "SCSI",
223            "deviceName": extra_disk_name,
224        }]
225
226    @staticmethod
227    def _LoadSshPublicKey(ssh_public_key_path):
228        """Load the content of ssh public key from a file.
229
230        Args:
231            ssh_public_key_path: String, path to the public key file.
232                               E.g. ~/.ssh/acloud_rsa.pub
233        Returns:
234            String, content of the file.
235
236        Raises:
237            errors.DriverError if the public key file does not exist
238            or the content is not valid.
239        """
240        key_path = os.path.expanduser(ssh_public_key_path)
241        if not os.path.exists(key_path):
242            raise errors.DriverError(
243                "SSH public key file %s does not exist." % key_path)
244
245        with open(key_path) as f:
246            rsa = f.read()
247            rsa = rsa.strip() if rsa else rsa
248            utils.VerifyRsaPubKey(rsa)
249        return rsa
250
251    def CreateInstance(self, instance, image_name, extra_disk_name=None):
252        """Create a gce instance given an gce image.
253
254        Args:
255            instance: instance name.
256            image_name: A string, the name of the GCE image.
257        """
258        self._CheckMachineSize()
259        disk_args = self._GetDiskArgs(instance, image_name)
260        if extra_disk_name:
261            disk_args.extend(self._GetExtraDiskArgs(extra_disk_name))
262        metadata = self._metadata.copy()
263        metadata["cfg_sta_display_resolution"] = self._resolution
264        metadata["t_force_orientation"] = self._orientation
265
266        # Add per-instance ssh key
267        if self._ssh_public_key_path:
268            rsa = self._LoadSshPublicKey(self._ssh_public_key_path)
269            logger.info("ssh_public_key_path is specified in config: %s, "
270                        "will add the key to the instance.",
271                        self._ssh_public_key_path)
272            metadata["sshKeys"] = "%s:%s" % (getpass.getuser(), rsa)
273        else:
274            logger.warning(
275                "ssh_public_key_path is not specified in config, "
276                "only project-wide key will be effective.")
277
278        super(AndroidComputeClient, self).CreateInstance(
279            instance, image_name, self._machine_type, metadata, self._network,
280            self._zone, disk_args)
281
282    def CheckBoot(self, instance):
283        """Check once to see if boot completes.
284
285        Args:
286            instance: string, instance name.
287
288        Returns:
289            True if the BOOT_COMPLETED_MSG appears in serial port output.
290            otherwise False.
291        """
292        try:
293            return self.BOOT_COMPLETED_MSG in self.GetSerialPortOutput(
294                instance=instance, port=1)
295        except errors.HttpError as e:
296            if e.code == 400:
297                logger.debug("CheckBoot: Instance is not ready yet %s",
298                              str(e))
299                return False
300            raise
301
302    def WaitForBoot(self, instance):
303        """Wait for boot to completes or hit timeout.
304
305        Args:
306            instance: string, instance name.
307        """
308        logger.info("Waiting for instance to boot up: %s", instance)
309        timeout_exception = errors.DeviceBootTimeoutError(
310            "Device %s did not finish on boot within timeout (%s secs)" %
311            (instance, self.BOOT_TIMEOUT_SECS)),
312        utils.PollAndWait(func=self.CheckBoot,
313                          expected_return=True,
314                          timeout_exception=timeout_exception,
315                          timeout_secs=self.BOOT_TIMEOUT_SECS,
316                          sleep_interval_secs=self.BOOT_CHECK_INTERVAL_SECS,
317                          instance=instance)
318        logger.info("Instance boot completed: %s", instance)
319
320    def GetInstanceIP(self, instance):
321        """Get Instance IP given instance name.
322
323        Args:
324            instance: String, representing instance name.
325
326        Returns:
327            string, IP of the instance.
328        """
329        return super(AndroidComputeClient, self).GetInstanceIP(instance,
330                                                               self._zone)
331
332    def GetSerialPortOutput(self, instance, port=1):
333        """Get serial port output.
334
335        Args:
336            instance: string, instance name.
337            port: int, which COM port to read from, 1-4, default to 1.
338
339        Returns:
340            String, contents of the output.
341
342        Raises:
343            errors.DriverError: For malformed response.
344        """
345        return super(AndroidComputeClient, self).GetSerialPortOutput(
346            instance, self._zone, port)
347
348    def GetInstanceNamesByIPs(self, ips):
349        """Get Instance names by IPs.
350
351        This function will go through all instances, which
352        could be slow if there are too many instances.  However, currently
353        GCE doesn't support search for instance by IP.
354
355        Args:
356            ips: A set of IPs.
357
358        Returns:
359            A dictionary where key is ip and value is instance name or None
360            if instance is not found for the given IP.
361        """
362        return super(AndroidComputeClient, self).GetInstanceNamesByIPs(
363            ips, self._zone)
364