1#!/usr/bin/env python3
2#
3#   Copyright 2019 - 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
17import os
18import logging
19import psutil
20import socket
21import tarfile
22import tempfile
23import time
24import usbinfo
25
26from acts import utils
27from acts.controllers.fuchsia_lib.ssh import FuchsiaSSHError
28from acts.libs.proc import job
29from acts.utils import get_fuchsia_mdns_ipv6_address
30
31MDNS_LOOKUP_RETRY_MAX = 3
32FASTBOOT_TIMEOUT = 30
33AFTER_FLASH_BOOT_TIME = 30
34WAIT_FOR_EXISTING_FLASH_TO_FINISH_SEC = 360
35PROCESS_CHECK_WAIT_TIME_SEC = 30
36
37FUCHSIA_SDK_URL = "gs://fuchsia-sdk/development"
38FUCHSIA_RELEASE_TESTING_URL = "gs://fuchsia-release-testing/images"
39
40
41def flash(fuchsia_device, use_ssh=False,
42          fuchsia_reconnect_after_reboot_time=5):
43    """A function to flash, not pave, a fuchsia_device
44
45    Args:
46        fuchsia_device: An ACTS fuchsia_device
47
48    Returns:
49        True if successful.
50    """
51    if not fuchsia_device.authorized_file:
52        raise ValueError('A ssh authorized_file must be present in the '
53                         'ACTS config to flash fuchsia_devices.')
54    # This is the product type from the fx set command.
55    # Do 'fx list-products' to see options in Fuchsia source tree.
56    if not fuchsia_device.product_type:
57        raise ValueError('A product type must be specified to flash '
58                         'fuchsia_devices.')
59    # This is the board type from the fx set command.
60    # Do 'fx list-boards' to see options in Fuchsia source tree.
61    if not fuchsia_device.board_type:
62        raise ValueError('A board type must be specified to flash '
63                         'fuchsia_devices.')
64    if not fuchsia_device.build_number:
65        fuchsia_device.build_number = 'LATEST'
66    if not fuchsia_device.mdns_name:
67        raise ValueError(
68            'Either fuchsia_device mdns_name must be specified or '
69            'ip must be the mDNS name to be able to flash.')
70
71    file_to_download = None
72    image_archive_path = None
73    image_path = None
74
75    if not fuchsia_device.specific_image:
76        product_build = fuchsia_device.product_type
77        if fuchsia_device.build_type:
78            product_build = f'{product_build}_{fuchsia_device.build_type}'
79        if 'LATEST' in fuchsia_device.build_number:
80            sdk_version = 'sdk'
81            if 'LATEST_F' in fuchsia_device.build_number:
82                f_branch = fuchsia_device.build_number.split('LATEST_F', 1)[1]
83                sdk_version = f'f{f_branch}_sdk'
84            file_to_download = (
85                f'{FUCHSIA_RELEASE_TESTING_URL}/'
86                f'{sdk_version}-{product_build}.{fuchsia_device.board_type}-release.tgz'
87            )
88        else:
89            # Must be a fully qualified build number (e.g. 5.20210721.4.1215)
90            file_to_download = (
91                f'{FUCHSIA_SDK_URL}/{fuchsia_device.build_number}/images/'
92                f'{product_build}.{fuchsia_device.board_type}-release.tgz')
93    elif 'gs://' in fuchsia_device.specific_image:
94        file_to_download = fuchsia_device.specific_image
95    elif os.path.isdir(fuchsia_device.specific_image):
96        image_path = fuchsia_device.specific_image
97    elif tarfile.is_tarfile(fuchsia_device.specific_image):
98        image_archive_path = fuchsia_device.specific_image
99    else:
100        raise ValueError(
101            f'Invalid specific_image "{fuchsia_device.specific_image}"')
102
103    if image_path:
104        reboot_to_bootloader(fuchsia_device, use_ssh,
105                             fuchsia_reconnect_after_reboot_time)
106        logging.info(
107            f'Flashing {fuchsia_device.mdns_name} with {image_path} using authorized keys "{fuchsia_device.authorized_file}".'
108        )
109        run_flash_script(fuchsia_device, image_path)
110    else:
111        suffix = fuchsia_device.board_type
112        with tempfile.TemporaryDirectory(suffix=suffix) as image_path:
113            if file_to_download:
114                logging.info(f'Downloading {file_to_download} to {image_path}')
115                job.run(f'gsutil cp {file_to_download} {image_path}')
116                image_archive_path = os.path.join(
117                    image_path, os.path.basename(file_to_download))
118
119            if image_archive_path:
120                # Use tar command instead of tarfile.extractall, as it takes too long.
121                job.run(f'tar xfvz {image_archive_path} -C {image_path}',
122                        timeout=120)
123
124            reboot_to_bootloader(fuchsia_device, use_ssh,
125                                 fuchsia_reconnect_after_reboot_time)
126
127            logging.info(
128                f'Flashing {fuchsia_device.mdns_name} with {image_archive_path} using authorized keys "{fuchsia_device.authorized_file}".'
129            )
130            run_flash_script(fuchsia_device, image_path)
131    return True
132
133
134def reboot_to_bootloader(fuchsia_device,
135                         use_ssh=False,
136                         fuchsia_reconnect_after_reboot_time=5):
137    if use_ssh:
138        logging.info('Sending reboot command via SSH to '
139                     'get into bootloader.')
140        # Sending this command will put the device in fastboot
141        # but it does not guarantee the device will be in fastboot
142        # after this command.  There is no check so if there is an
143        # expectation of the device being in fastboot, then some
144        # other check needs to be done.
145        try:
146            fuchsia_device.ssh.run(
147                'dm rb', timeout_sec=fuchsia_reconnect_after_reboot_time)
148        except FuchsiaSSHError as e:
149            if 'closed by remote host' not in e.result.stderr:
150                raise e
151    else:
152        pass
153        ## Todo: Add elif for SL4F if implemented in SL4F
154
155    time_counter = 0
156    while time_counter < FASTBOOT_TIMEOUT:
157        logging.info('Checking to see if fuchsia_device(%s) SN: %s is in '
158                     'fastboot. (Attempt #%s Timeout: %s)' %
159                     (fuchsia_device.mdns_name, fuchsia_device.serial_number,
160                      str(time_counter + 1), FASTBOOT_TIMEOUT))
161        for usb_device in usbinfo.usbinfo():
162            if (usb_device['iSerialNumber'] == fuchsia_device.serial_number
163                    and usb_device['iProduct'] == 'USB_download_gadget'):
164                logging.info(
165                    'fuchsia_device(%s) SN: %s is in fastboot.' %
166                    (fuchsia_device.mdns_name, fuchsia_device.serial_number))
167                time_counter = FASTBOOT_TIMEOUT
168        time_counter = time_counter + 1
169        if time_counter == FASTBOOT_TIMEOUT:
170            for fail_usb_device in usbinfo.usbinfo():
171                logging.debug(fail_usb_device)
172            raise TimeoutError(
173                'fuchsia_device(%s) SN: %s '
174                'never went into fastboot' %
175                (fuchsia_device.mdns_name, fuchsia_device.serial_number))
176        time.sleep(1)
177
178    end_time = time.time() + WAIT_FOR_EXISTING_FLASH_TO_FINISH_SEC
179    # Attempt to wait for existing flashing process to finish
180    while time.time() < end_time:
181        flash_process_found = False
182        for proc in psutil.process_iter():
183            if "bash" in proc.name() and "flash.sh" in proc.cmdline():
184                logging.info(
185                    "Waiting for existing flash.sh process to complete.")
186                time.sleep(PROCESS_CHECK_WAIT_TIME_SEC)
187                flash_process_found = True
188        if not flash_process_found:
189            break
190
191
192def run_flash_script(fuchsia_device, flash_dir):
193    try:
194        flash_output = job.run(
195            f'bash {flash_dir}/flash.sh --ssh-key={fuchsia_device.authorized_file} -s {fuchsia_device.serial_number}',
196            timeout=120)
197        logging.debug(flash_output.stderr)
198    except job.TimeoutError as err:
199        raise TimeoutError(err)
200
201    logging.info('Waiting %s seconds for device'
202                 ' to come back up after flashing.' % AFTER_FLASH_BOOT_TIME)
203    time.sleep(AFTER_FLASH_BOOT_TIME)
204    logging.info('Updating device to new IP addresses.')
205    mdns_ip = None
206    for retry_counter in range(MDNS_LOOKUP_RETRY_MAX):
207        mdns_ip = get_fuchsia_mdns_ipv6_address(fuchsia_device.mdns_name)
208        if mdns_ip:
209            break
210        else:
211            time.sleep(1)
212    if mdns_ip and utils.is_valid_ipv6_address(mdns_ip):
213        logging.info('IP for fuchsia_device(%s) changed from %s to %s' %
214                     (fuchsia_device.mdns_name, fuchsia_device.ip, mdns_ip))
215        fuchsia_device.ip = mdns_ip
216        fuchsia_device.address = "http://[{}]:{}".format(
217            fuchsia_device.ip, fuchsia_device.sl4f_port)
218    else:
219        raise ValueError('Invalid IP: %s after flashing.' %
220                         fuchsia_device.mdns_name)
221
222
223def wait_for_port(host: str, port: int, timeout_sec: int = 5) -> None:
224    """Wait for the host to start accepting connections on the port.
225
226    Some services take some time to start. Call this after launching the service
227    to avoid race conditions.
228
229    Args:
230        host: IP of the running service.
231        port: Port of the running service.
232        timeout_sec: Seconds to wait until raising TimeoutError
233
234    Raises:
235        TimeoutError: when timeout_sec has expired without a successful
236            connection to the service
237    """
238    timeout = time.perf_counter() + timeout_sec
239    while True:
240        try:
241            with socket.create_connection((host, port), timeout=timeout_sec):
242                return
243        except ConnectionRefusedError as e:
244            if time.perf_counter() > timeout:
245                raise TimeoutError(
246                    f'Waited over {timeout_sec}s for the service to start '
247                    f'accepting connections at {host}:{port}') from e
248