tools/gonk_tools/flash.py

# Copyright 2023 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Flash Gonk via DFU."""

import argparse
from pathlib import Path
import shutil
import signal
import subprocess
import sys

import pw_cli.color

from gonk_tools.firmware_files import bundled_bin_path

_COLOR = pw_cli.color.colors()


def _parse_args():
    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument(
        '--bin-file',
        help='Binary to flash',
        type=Path,
    )
    return parser.parse_args()


def pyfu_flash(bin_file: Path) -> int:
    pyfu_util_binary = shutil.which('pyfu-usb')
    if not pyfu_util_binary or not Path(pyfu_util_binary).is_file():
        raise FileNotFoundError('Unable to find "pyfu-usb"')

    pyfu_flash_args = [
        pyfu_util_binary,
        '--verbose',
        '--device',
        '0483:df11',
        '--address',
        '0x08000000',
        '--download',
        str(bin_file),
    ]

    print(
        _COLOR.cyan('Running ==>'),
        ' '.join(pyfu_flash_args),
    )

    # Ignore Ctrl-C to allow pyfu-usb to handle normally.
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    return subprocess.run(pyfu_flash_args, check=False).returncode


def dfu_flash(bin_file: Path) -> int:
    dfu_util_binary = shutil.which('dfu-util')
    if not dfu_util_binary or not Path(dfu_util_binary).is_file():
        raise FileNotFoundError('Unable to find "dfu-util"')

    dfu_flash_args = [
        dfu_util_binary,
        '--device',
        '0483:df11',
        '--dfuse-address',
        '0x08000000:leave',
        '--alt',
        '0',
        '--download',
        str(bin_file),
    ]

    print(
        _COLOR.cyan('Running ==>'),
        ' '.join(dfu_flash_args),
    )

    # Ignore Ctrl-C to allow pyfu-usb to handle normally.
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    return subprocess.run(dfu_flash_args, check=False).returncode


def main(
    bin_file: Path | None = None,
) -> int:
    """Flash Gonk via pyfu-usb."""

    if not bin_file:
        bin_file = bundled_bin_path()

    if not bin_file or not bin_file.is_file():
        raise FileNotFoundError(
            f'Unable to find the binary file to flash: "{bin_file}"'
        )

    pyfu_flash_returncode = pyfu_flash(bin_file)
    if pyfu_flash_returncode == 0:
        return 0

    print(
        _COLOR.red('ERROR:'),
        'Unable to flash with pyfu-util; Falling back to "dfu-util"',
    )
    return dfu_flash(bin_file)


def flash() -> None:
    sys.exit(main(**vars(_parse_args())))


if __name__ == '__main__':
    flash()