1# Copyright 2019 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""RemoteInstanceDeviceFactory provides basic interface to create a cuttlefish
16device factory."""
17
18import glob
19import logging
20import os
21import shutil
22import subprocess
23import tempfile
24
25from acloud import errors
26from acloud.internal import constants
27from acloud.internal.lib import utils
28from acloud.internal.lib import ssh
29from acloud.public.actions import gce_device_factory
30
31
32logger = logging.getLogger(__name__)
33_ALL_FILES = "*"
34# bootloader and kernel are files required to launch AVD.
35_BOOTLOADER = "bootloader"
36_KERNEL = "kernel"
37_ARTIFACT_FILES = ["*.img", _BOOTLOADER, _KERNEL]
38_HOME_FOLDER = os.path.expanduser("~")
39
40
41class RemoteInstanceDeviceFactory(gce_device_factory.GCEDeviceFactory):
42    """A class that can produce a cuttlefish device.
43
44    Attributes:
45        avd_spec: AVDSpec object that tells us what we're going to create.
46        cfg: An AcloudConfig instance.
47        local_image_artifact: A string, path to local image.
48        cvd_host_package_artifact: A string, path to cvd host package.
49        report_internal_ip: Boolean, True for the internal ip is used when
50                            connecting from another GCE instance.
51        credentials: An oauth2client.OAuth2Credentials instance.
52        compute_client: An object of cvd_compute_client.CvdComputeClient.
53        ssh: An Ssh object.
54    """
55    def __init__(self, avd_spec, local_image_artifact=None,
56                 cvd_host_package_artifact=None):
57        super().__init__(avd_spec, local_image_artifact)
58        self._cvd_host_package_artifact = cvd_host_package_artifact
59
60    # pylint: disable=broad-except
61    def CreateInstance(self):
62        """Create a single configured cuttlefish device.
63
64        GCE:
65        1. Create gcp instance.
66        2. Upload local built artifacts to remote instance or fetch build on
67           remote instance.
68        3. Launch CVD.
69
70        Remote host:
71        1. Init remote host.
72        2. Download the artifacts to local and upload the artifacts to host
73        3. Launch CVD.
74
75        Returns:
76            A string, representing instance name.
77        """
78        if self._avd_spec.instance_type == constants.INSTANCE_TYPE_HOST:
79            instance = self._InitRemotehost()
80            self._ProcessRemoteHostArtifacts()
81            self._LaunchCvd(instance=instance,
82                            decompress_kernel=None,
83                            boot_timeout_secs=self._avd_spec.boot_timeout_secs)
84        else:
85            instance = self._CreateGceInstance()
86            # If instance is failed, no need to go next step.
87            if instance in self.GetFailures():
88                return instance
89            try:
90                self._ProcessArtifacts(self._avd_spec.image_source)
91                self._LaunchCvd(instance=instance,
92                                boot_timeout_secs=self._avd_spec.boot_timeout_secs)
93            except Exception as e:
94                self._SetFailures(instance, e)
95
96        return instance
97
98    def _InitRemotehost(self):
99        """Initialize remote host.
100
101        Determine the remote host instance name, and activate ssh. It need to
102        get the IP address in the common_operation. So need to pass the IP and
103        ssh to compute_client.
104
105        build_target: The format is like "aosp_cf_x86_phone". We only get info
106                      from the user build image file name. If the file name is
107                      not custom format (no "-"), we will use $TARGET_PRODUCT
108                      from environment variable as build_target.
109
110        Returns:
111            A string, representing instance name.
112        """
113        image_name = os.path.basename(
114            self._local_image_artifact) if self._local_image_artifact else ""
115        build_target = (os.environ.get(constants.ENV_BUILD_TARGET) if "-" not
116                        in image_name else image_name.split("-")[0])
117        build_id = self._USER_BUILD
118        if self._avd_spec.image_source == constants.IMAGE_SRC_REMOTE:
119            build_id = self._avd_spec.remote_image[constants.BUILD_ID]
120
121        instance = "%s-%s-%s-%s" % (constants.INSTANCE_TYPE_HOST,
122                                    self._avd_spec.remote_host,
123                                    build_id, build_target)
124        ip = ssh.IP(ip=self._avd_spec.remote_host)
125        self._ssh = ssh.Ssh(
126            ip=ip,
127            user=self._avd_spec.host_user,
128            ssh_private_key_path=(self._avd_spec.host_ssh_private_key_path or
129                                  self._cfg.ssh_private_key_path),
130            extra_args_ssh_tunnel=self._cfg.extra_args_ssh_tunnel,
131            report_internal_ip=self._report_internal_ip)
132        self._compute_client.InitRemoteHost(
133            self._ssh, ip, self._avd_spec.host_user)
134        return instance
135
136    @utils.TimeExecute(function_description="Downloading Android Build artifact")
137    def _DownloadArtifacts(self, extract_path):
138        """Download the CF image artifacts and process them.
139
140        - Download images from the Android Build system.
141        - Download cvd host package from the Android Build system.
142
143        Args:
144            extract_path: String, a path include extracted files.
145
146        Raises:
147            errors.GetRemoteImageError: Fails to download rom images.
148        """
149        cfg = self._avd_spec.cfg
150        build_id = self._avd_spec.remote_image[constants.BUILD_ID]
151        build_branch = self._avd_spec.remote_image[constants.BUILD_BRANCH]
152        build_target = self._avd_spec.remote_image[constants.BUILD_TARGET]
153
154        # Download images with fetch_cvd
155        fetch_cvd = os.path.join(extract_path, constants.FETCH_CVD)
156        self._compute_client.build_api.DownloadFetchcvd(fetch_cvd,
157                                                        cfg.fetch_cvd_version)
158        fetch_cvd_build_args = self._compute_client.build_api.GetFetchBuildArgs(
159            build_id, build_branch, build_target,
160            self._avd_spec.system_build_info.get(constants.BUILD_ID),
161            self._avd_spec.system_build_info.get(constants.BUILD_BRANCH),
162            self._avd_spec.system_build_info.get(constants.BUILD_TARGET),
163            self._avd_spec.kernel_build_info.get(constants.BUILD_ID),
164            self._avd_spec.kernel_build_info.get(constants.BUILD_BRANCH),
165            self._avd_spec.kernel_build_info.get(constants.BUILD_TARGET),
166            self._avd_spec.bootloader_build_info.get(constants.BUILD_ID),
167            self._avd_spec.bootloader_build_info.get(constants.BUILD_BRANCH),
168            self._avd_spec.bootloader_build_info.get(constants.BUILD_TARGET))
169        creds_cache_file = os.path.join(_HOME_FOLDER, cfg.creds_cache_file)
170        fetch_cvd_cert_arg = self._compute_client.build_api.GetFetchCertArg(
171            creds_cache_file)
172        fetch_cvd_args = [fetch_cvd, "-directory=%s" % extract_path,
173                          fetch_cvd_cert_arg]
174        fetch_cvd_args.extend(fetch_cvd_build_args)
175        logger.debug("Download images command: %s", fetch_cvd_args)
176        try:
177            subprocess.check_call(fetch_cvd_args)
178        except subprocess.CalledProcessError as e:
179            raise errors.GetRemoteImageError("Fails to download images: %s" % e)
180
181    def _ProcessRemoteHostArtifacts(self):
182        """Process remote host artifacts.
183
184        - If images source is local, tool will upload images from local site to
185          remote host.
186        - If images source is remote, tool will download images from android
187          build to local and unzip it then upload to remote host, because there
188          is no permission to fetch build rom on the remote host.
189        """
190        self._compute_client.SetStage(constants.STAGE_ARTIFACT)
191        if self._avd_spec.image_source == constants.IMAGE_SRC_LOCAL:
192            self._UploadLocalImageArtifacts(
193                self._local_image_artifact, self._cvd_host_package_artifact,
194                self._avd_spec.local_image_dir)
195        else:
196            try:
197                artifacts_path = tempfile.mkdtemp()
198                logger.debug("Extracted path of artifacts: %s", artifacts_path)
199                self._DownloadArtifacts(artifacts_path)
200                self._UploadRemoteImageArtifacts(artifacts_path)
201            finally:
202                shutil.rmtree(artifacts_path)
203
204    def _ProcessArtifacts(self, image_source):
205        """Process artifacts.
206
207        - If images source is local, tool will upload images from local site to
208          remote instance.
209        - If images source is remote, tool will download images from android
210          build to remote instance. Before download images, we have to update
211          fetch_cvd to remote instance.
212
213        Args:
214            image_source: String, the type of image source is remote or local.
215        """
216        if image_source == constants.IMAGE_SRC_LOCAL:
217            self._UploadLocalImageArtifacts(self._local_image_artifact,
218                                            self._cvd_host_package_artifact,
219                                            self._avd_spec.local_image_dir)
220        elif image_source == constants.IMAGE_SRC_REMOTE:
221            self._compute_client.UpdateFetchCvd()
222            self._FetchBuild(self._avd_spec)
223
224    def _FetchBuild(self, avd_spec):
225        """Download CF artifacts from android build.
226
227        Args:
228            avd_spec: AVDSpec object that tells us what we're going to create.
229        """
230        self._compute_client.FetchBuild(
231            avd_spec.remote_image[constants.BUILD_ID],
232            avd_spec.remote_image[constants.BUILD_BRANCH],
233            avd_spec.remote_image[constants.BUILD_TARGET],
234            avd_spec.system_build_info[constants.BUILD_ID],
235            avd_spec.system_build_info[constants.BUILD_BRANCH],
236            avd_spec.system_build_info[constants.BUILD_TARGET],
237            avd_spec.kernel_build_info[constants.BUILD_ID],
238            avd_spec.kernel_build_info[constants.BUILD_BRANCH],
239            avd_spec.kernel_build_info[constants.BUILD_TARGET],
240            avd_spec.bootloader_build_info[constants.BUILD_ID],
241            avd_spec.bootloader_build_info[constants.BUILD_BRANCH],
242            avd_spec.bootloader_build_info[constants.BUILD_TARGET])
243
244    @utils.TimeExecute(function_description="Processing and uploading local images")
245    def _UploadLocalImageArtifacts(self,
246                                   local_image_zip,
247                                   cvd_host_package_artifact,
248                                   images_dir):
249        """Upload local images and avd local host package to instance.
250
251        There are two ways to upload local images.
252        1. Using local image zip, it would be decompressed by install_zip.sh.
253        2. Using local image directory, this directory contains all images.
254           Images are compressed/decompressed by lzop during upload process.
255
256        Args:
257            local_image_zip: String, path to zip of local images which
258                             build from 'm dist'.
259            cvd_host_package_artifact: String, path to cvd host package.
260            images_dir: String, directory of local images which build
261                        from 'm'.
262        """
263        if local_image_zip:
264            remote_cmd = ("/usr/bin/install_zip.sh . < %s" % local_image_zip)
265            logger.debug("remote_cmd:\n %s", remote_cmd)
266            self._ssh.Run(remote_cmd)
267        else:
268            # Compress image files for faster upload.
269            try:
270                images_path = os.path.join(images_dir, "required_images")
271                with open(images_path, "r") as images:
272                    artifact_files = images.read().splitlines()
273            except IOError:
274                # Older builds may not have a required_images file. In this case
275                # we fall back to *.img.
276                artifact_files = []
277                for file_name in _ARTIFACT_FILES:
278                    artifact_files.extend(
279                        os.path.basename(image) for image in glob.glob(
280                            os.path.join(images_dir, file_name)))
281            cmd = ("tar -cf - --lzop -S -C {images_dir} {artifact_files} | "
282                   "{ssh_cmd} -- tar -xf - --lzop -S".format(
283                       images_dir=images_dir,
284                       artifact_files=" ".join(artifact_files),
285                       ssh_cmd=self._ssh.GetBaseCmd(constants.SSH_BIN)))
286            logger.debug("cmd:\n %s", cmd)
287            ssh.ShellCmdWithRetry(cmd)
288
289        # host_package
290        remote_cmd = ("tar -x -z -f - < %s" % cvd_host_package_artifact)
291        logger.debug("remote_cmd:\n %s", remote_cmd)
292        self._ssh.Run(remote_cmd)
293
294    @utils.TimeExecute(function_description="Uploading remote image artifacts")
295    def _UploadRemoteImageArtifacts(self, images_dir):
296        """Upload remote image artifacts to instance.
297
298        Args:
299            images_dir: String, directory of local artifacts downloaded by fetch_cvd.
300        """
301        artifact_files = [
302            os.path.basename(image)
303            for image in glob.glob(os.path.join(images_dir, _ALL_FILES))
304        ]
305        # TODO(b/182259589): Refactor upload image command into a function.
306        cmd = ("tar -cf - --lzop -S -C {images_dir} {artifact_files} | "
307                "{ssh_cmd} -- tar -xf - --lzop -S".format(
308                    images_dir=images_dir,
309                    artifact_files=" ".join(artifact_files),
310                    ssh_cmd=self._ssh.GetBaseCmd(constants.SSH_BIN)))
311        logger.debug("cmd:\n %s", cmd)
312        ssh.ShellCmdWithRetry(cmd)
313
314    def _LaunchCvd(self, instance, decompress_kernel=None,
315                   boot_timeout_secs=None):
316        """Launch CVD.
317
318        Args:
319            instance: String, instance name.
320            boot_timeout_secs: Integer, the maximum time to wait for the
321                               command to respond.
322        """
323        # TODO(b/140076771) Support kernel image for local image mode.
324        self._compute_client.LaunchCvd(
325            instance,
326            self._avd_spec,
327            self._cfg.extra_data_disk_size_gb,
328            decompress_kernel,
329            boot_timeout_secs)
330
331    def GetBuildInfoDict(self):
332        """Get build info dictionary.
333
334        Returns:
335            A build info dictionary. None for local image case.
336        """
337        if self._avd_spec.image_source == constants.IMAGE_SRC_LOCAL:
338            return None
339        build_info_dict = {
340            key: val for key, val in self._avd_spec.remote_image.items() if val}
341
342        # kernel_target have default value "kernel". If user provide kernel_build_id
343        # or kernel_branch, then start to process kernel image.
344        if (self._avd_spec.kernel_build_info[constants.BUILD_ID]
345                or self._avd_spec.kernel_build_info[constants.BUILD_BRANCH]):
346            build_info_dict.update(
347                {"kernel_%s" % key: val
348                 for key, val in self._avd_spec.kernel_build_info.items() if val}
349            )
350        build_info_dict.update(
351            {"system_%s" % key: val
352             for key, val in self._avd_spec.system_build_info.items() if val}
353        )
354        build_info_dict.update(
355            {"bootloader_%s" % key: val
356             for key, val in self._avd_spec.bootloader_build_info.items() if val}
357        )
358        return build_info_dict
359