1#
2# Copyright (C) 2017 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16"""Class to flash build artifacts onto devices"""
17
18import hashlib
19import logging
20import os
21import resource
22import sys
23import tempfile
24import time
25
26from host_controller import common
27from vts.utils.python.common import cmd_utils
28from vts.utils.python.controllers import android_device
29
30
31class BuildFlasher(object):
32    """Client that manages build flashing.
33
34    Attributes:
35        device: AndroidDevice, the device associated with the client.
36    """
37
38    def __init__(self, serial="", customflasher_path=""):
39        """Initialize the client.
40
41        If serial not provided, find single device connected. Error if
42        zero or > 1 devices connected.
43
44        Args:
45            serial: optional string, serial number for the device.
46            customflasher_path: optional string, set device to use specified
47                                binary to flash a device
48        """
49        if serial != "":
50            self.device = android_device.AndroidDevice(
51                serial, device_callback_port=-1)
52        else:
53            serials = android_device.list_adb_devices()
54            if len(serials) == 0:
55                serials = android_device.list_fastboot_devices()
56                if len(serials) == 0:
57                    raise android_device.AndroidDeviceError(
58                        "ADB and fastboot could not find any target devices.")
59            if len(serials) > 1:
60                logging.info("ADB or fastboot found more than one device: %s",
61                             serials)
62            self.device = android_device.AndroidDevice(
63                serials[0], device_callback_port=-1)
64            if customflasher_path:
65                self.device.SetCustomFlasherPath(customflasher_path)
66
67    def SetSerial(self, serial):
68        """Sets device serial.
69
70        Args:
71            serial: string, a device serial.
72
73        Returns:
74            True if successful; False otherwise.
75        """
76        if not serial:
77            logging.error("no serial is given to BuildFlasher.SetSerial.")
78            return False
79
80        self.device = android_device.AndroidDevice(
81            serial, device_callback_port=-1)
82        return True
83
84    def FlashGSI(self,
85                 system_img,
86                 vbmeta_img=None,
87                 skip_check=False,
88                 skip_vbmeta=False):
89        """Flash the Generic System Image to the device.
90
91        Args:
92            system_img: string, path to GSI
93            vbmeta_img: string, optional, path to vbmeta image for new devices
94            skip_check: boolean, set True to skip adb-based checks when
95                        the DUT is already running its bootloader.
96            skip_vbmeta: bool, whether to skip flashing the vbmeta.img or not.
97                         If the device has the vbmeta slot then flash vbmeta.img
98                         even if the skip_vbmeta is set to True.
99        """
100        if not os.path.exists(system_img):
101            raise ValueError("Couldn't find system image at %s" % system_img)
102        if not skip_check:
103            self.device.adb.wait_for_device()
104            if not self.device.isBootloaderMode:
105                self.device.log.info(self.device.adb.reboot_bootloader())
106        if vbmeta_img is not None:
107            if skip_vbmeta == False or self.device.hasVbmetaSlot:
108                self.device.log.info(
109                    self.device.fastboot.flash('vbmeta', vbmeta_img))
110        self.device.log.info(self.device.fastboot.erase('system'))
111        self.device.log.info(self.device.fastboot.flash('system', system_img))
112        self.device.log.info(self.device.fastboot.erase('metadata'))
113        self.device.log.info(self.device.fastboot._w())
114        self.device.log.info(self.device.fastboot.reboot())
115
116    def Flashall(self, directory):
117        """Flash all images in a directory to the device using flashall.
118
119        Generally the directory is the result of unzipping the .zip from AB.
120        Args:
121            directory: string, path to directory containing images
122        """
123        # fastboot flashall looks for imgs in $ANDROID_PRODUCT_OUT
124        os.environ['ANDROID_PRODUCT_OUT'] = directory
125        self.device.adb.wait_for_device()
126        if not self.device.isBootloaderMode:
127            self.device.log.info(self.device.adb.reboot_bootloader())
128        self.device.log.info(self.device.fastboot.flashall())
129
130    def Flash(self, device_images, skip_vbmeta=False):
131        """Flash the Generic System Image to the device.
132
133        Args:
134            device_images: dict, where the key is partition name and value is
135                           image file path.
136            skip_vbmeta: bool, whether to skip flashing the vbmeta.img or not.
137
138        Returns:
139            True if succesful; False otherwise
140        """
141        if not device_images:
142            logging.warn("Flash skipped because no device image is given.")
143            return False
144
145        if not self.device.isBootloaderMode:
146            self.device.adb.wait_for_device()
147            logging.info("rebooting to bootloader")
148            self.device.log.info(self.device.adb.reboot_bootloader())
149
150        logging.info("checking to flash bootloader.img and radio.img")
151        for partition in ["bootloader", "radio"]:
152            if partition in device_images:
153                image_path = device_images[partition]
154                self.device.log.info("fastboot flash %s %s", partition,
155                                     image_path)
156                self.device.log.info(
157                    self.device.fastboot.flash(partition, image_path))
158                self.device.log.info("fastboot reboot_bootloader")
159                self.device.log.info(self.device.fastboot.reboot_bootloader())
160
161        logging.info("starting to flash vendor and other images...")
162        full_zipfile = False
163        if common.FULL_ZIPFILE in device_images:
164            logging.info("fastboot update %s --skip-reboot",
165                         (device_images[common.FULL_ZIPFILE]))
166            self.device.log.info(
167                self.device.fastboot.update(device_images[common.FULL_ZIPFILE],
168                                            "--skip-reboot"))
169            full_zipfile = True
170
171        for partition, image_path in device_images.iteritems():
172            if partition in (common.FULL_ZIPFILE, common.FULL_ZIPFILE_DIR,
173                             "system", "vbmeta", "bootloader", "radio",
174                             "metadata", "userdata"):
175                continue
176            if full_zipfile and partition in ("vendor", "boot"):
177                logging.info("%s skipped because full zipfile was updated.",
178                             partition)
179                continue
180            if not image_path:
181                self.device.log.warning("%s image is empty", partition)
182                continue
183            self.device.log.info("fastboot flash %s %s", partition, image_path)
184            self.device.log.info(
185                self.device.fastboot.flash(partition, image_path))
186
187        logging.info("starting to flash system and other images...")
188        if "system" in device_images and device_images["system"]:
189            system_img = device_images["system"]
190            vbmeta_img = device_images["vbmeta"] if (
191                "vbmeta" in device_images
192                and device_images["vbmeta"]) else None
193            self.FlashGSI(
194                system_img,
195                vbmeta_img,
196                skip_check=True,
197                skip_vbmeta=skip_vbmeta)
198        else:
199            self.device.log.info(self.device.fastboot.reboot())
200        return True
201
202    def FlashImage(self, device_images, image_partition=None, reboot=False):
203        """Flash specified image(s) to the device.
204
205        Args:
206            device_images: dict, where the key is partition name and value is
207                           image file path.
208            image_partition: string, set to flash only an image in a specified
209                             partition.
210            reboot: boolean, true to reboot the device.
211
212        Returns:
213            True if successful, False otherwise
214        """
215        if not device_images:
216            logging.warn("Flash skipped because no device image is given.")
217            return False
218
219        if not self.device.isBootloaderMode:
220            self.device.adb.wait_for_device()
221            self.device.log.info(self.device.adb.reboot_bootloader())
222
223        for partition, image_path in device_images.iteritems():
224            if image_partition and image_partition != partition:
225                continue
226            if partition.endswith(".img"):
227                partition = partition[:-4]
228            self.device.log.info(
229                self.device.fastboot.flash(partition, image_path))
230        if reboot:
231            self.device.log.info(self.device.fastboot.reboot())
232        return True
233
234    def WaitForDevice(self, timeout_secs=600):
235        """Waits for the device to boot completely.
236
237        Args:
238            timeout_secs: integer, the maximum timeout value for this
239                          operation (unit: seconds).
240
241        Returns:
242            True if device is booted successfully; False otherwise.
243        """
244        return self.device.waitForBootCompletion(timeout=timeout_secs)
245
246    def FlashUsingCustomBinary(self,
247                               device_images,
248                               reboot_mode,
249                               flasher_args,
250                               timeout_secs_for_reboot=900):
251        """Flash the customized image to the device.
252
253        Args:
254            device_images: dict, where the key is partition name and value is
255                           image file path.
256            reboot_mode: string, decides which mode device will reboot into.
257                         ("bootloader"/"download").
258            flasher_args: list of strings, arguments that will be passed to the
259                          flash binary.
260            timeout_secs_for_reboot: integer, the maximum timeout value for
261                                     reboot to flash-able mode(unit: seconds).
262
263        Returns:
264            True if successful; False otherwise.
265        """
266        if not device_images:
267            logging.warn("Flash skipped because no device image is given.")
268            return False
269
270        if not flasher_args:
271            logging.error("No arguments.")
272            return False
273
274        if not self.device.isBootloaderMode:
275            self.device.adb.wait_for_device()
276            logging.info("rebooting to %s mode", reboot_mode)
277            self.device.log.info(self.device.adb.reboot(reboot_mode))
278
279        start = time.time()
280        while not self.device.customflasher._l():
281            if time.time() - start >= timeout_secs_for_reboot:
282                logging.error(
283                    "Timeout while waiting for %s mode boot completion.",
284                    reboot_mode)
285                return False
286            time.sleep(1)
287
288        flasher_output = self.device.customflasher.ExecCustomFlasherCmd(
289            flasher_args[0],
290            " ".join(flasher_args[1:] + [device_images["img"]]))
291        self.device.log.info(flasher_output)
292
293        return True
294
295    def RepackageArtifacts(self, device_images, repackage_form):
296        """Repackage artifacts into a given format.
297
298        Once repackaged, device_images becomes
299        {"img": "path_to_repackaged_image"}
300
301        Args:
302            device_images: dict, where the key is partition name and value is
303                           image file path.
304            repackage_form: string, format to repackage.
305
306        Returns:
307            True if succesful; False otherwise.
308        """
309        if not device_images:
310            logging.warn("Repackage skipped because no device image is given.")
311            return False
312
313        if repackage_form == "tar.md5":
314            tmp_file_name = next(tempfile._get_candidate_names()) + ".tar"
315            tmp_dir_path = os.path.dirname(
316                device_images[device_images.keys()[0]])
317            for img in device_images:
318                if os.path.dirname(device_images[img]) != tmp_dir_path:
319                    os.rename(device_images[img],
320                              os.path.join(tmp_dir_path, img))
321                    device_images[img] = os.path.join(tmp_dir_path, img)
322
323            current_dir = os.getcwd()
324            os.chdir(tmp_dir_path)
325
326            if sys.platform == "linux2":
327                tar_cmd = "tar -cf %s %s" % (tmp_file_name, ' '.join(
328                    (device_images.keys())))
329            else:
330                logging.error("Unsupported OS for the given repackage form.")
331                return False
332            logging.info(tar_cmd)
333            std_out, std_err, err_code = cmd_utils.ExecuteOneShellCommand(
334                tar_cmd)
335            if err_code:
336                logging.error(std_err)
337                return False
338
339            hash_md5 = hashlib.md5()
340            try:
341                with open(tmp_file_name, "rb") as file:
342                    data_chunk = 0
343                    chunk_size = resource.getpagesize()
344                    while data_chunk != b'':
345                        data_chunk = file.read(chunk_size)
346                        hash_md5.update(data_chunk)
347                    hash_ret = hash_md5.hexdigest()
348                with open(tmp_file_name, "a") as file:
349                    file.write("%s  %s" % (hash_ret, tmp_file_name))
350            except IOError as e:
351                logging.error(e.strerror)
352                return False
353
354            device_images.clear()
355            device_images["img"] = os.path.join(tmp_dir_path, tmp_file_name)
356
357            os.chdir(current_dir)
358        else:
359            logging.error(
360                "Please specify correct repackage form: --repackage=%s",
361                repackage_form)
362            return False
363
364        return True
365