1# Copyright (c) 2020 Teslabs Engineering S.L.
2#
3# SPDX-License-Identifier: Apache-2.0
4
5"""Runner for flashing with STM32CubeProgrammer CLI, the official programming
6   utility from ST Microelectronics.
7"""
8
9import argparse
10import functools
11import os
12import platform
13import shlex
14import shutil
15from pathlib import Path
16from typing import ClassVar
17
18from runners.core import RunnerCaps, RunnerConfig, ZephyrBinaryRunner
19
20
21class STM32CubeProgrammerBinaryRunner(ZephyrBinaryRunner):
22    """Runner front-end for STM32CubeProgrammer CLI."""
23
24    _RESET_MODES: ClassVar[dict[str, str]] = {
25        "sw": "SWrst",
26        "hw": "HWrst",
27        "core": "Crst",
28    }
29    """Reset mode argument mappings."""
30
31    def __init__(
32        self,
33        cfg: RunnerConfig,
34        port: str,
35        frequency: int | None,
36        reset_mode: str | None,
37        start_address: int | None,
38        conn_modifiers: str | None,
39        cli: Path | None,
40        use_elf: bool,
41        erase: bool,
42        extload: str | None,
43        tool_opt: list[str],
44    ) -> None:
45        super().__init__(cfg)
46
47        self._port = port
48        self._frequency = frequency
49        self._start_address = start_address
50        self._reset_mode = reset_mode
51        self._conn_modifiers = conn_modifiers
52        self._cli = (
53            cli or STM32CubeProgrammerBinaryRunner._get_stm32cubeprogrammer_path()
54        )
55        self._use_elf = use_elf
56        self._erase = erase
57
58        if extload:
59            p = (
60                STM32CubeProgrammerBinaryRunner._get_stm32cubeprogrammer_path().parent.resolve()
61                / 'ExternalLoader'
62            )
63            self._extload = ['-el', str(p / extload)]
64        else:
65            self._extload = []
66
67        self._tool_opt: list[str] = list()
68        for opts in [shlex.split(opt) for opt in tool_opt]:
69            self._tool_opt += opts
70
71        # add required library loader path to the environment (Linux only)
72        if platform.system() == "Linux":
73            os.environ["LD_LIBRARY_PATH"] = str(self._cli.parent / ".." / "lib")
74
75    @staticmethod
76    def _get_stm32cubeprogrammer_path() -> Path:
77        """Obtain path of the STM32CubeProgrammer CLI tool."""
78
79        if platform.system() == "Linux":
80            cmd = shutil.which("STM32_Programmer_CLI")
81            if cmd is not None:
82                return Path(cmd)
83
84            return (
85                Path.home()
86                / "STMicroelectronics"
87                / "STM32Cube"
88                / "STM32CubeProgrammer"
89                / "bin"
90                / "STM32_Programmer_CLI"
91            )
92
93        if platform.system() == "Windows":
94            cmd = shutil.which("STM32_Programmer_CLI")
95            if cmd is not None:
96                return Path(cmd)
97
98            cli = (
99                Path("STMicroelectronics")
100                / "STM32Cube"
101                / "STM32CubeProgrammer"
102                / "bin"
103                / "STM32_Programmer_CLI.exe"
104            )
105            x86_path = Path(os.environ["PROGRAMFILES(X86)"]) / cli
106            if x86_path.exists():
107                return x86_path
108
109            return Path(os.environ["PROGRAMW6432"]) / cli
110
111        if platform.system() == "Darwin":
112            return (
113                Path("/Applications")
114                / "STMicroelectronics"
115                / "STM32Cube"
116                / "STM32CubeProgrammer"
117                / "STM32CubeProgrammer.app"
118                / "Contents"
119                / "MacOs"
120                / "bin"
121                / "STM32_Programmer_CLI"
122            )
123
124        raise NotImplementedError("Could not determine STM32_Programmer_CLI path")
125
126    @classmethod
127    def name(cls):
128        return "stm32cubeprogrammer"
129
130    @classmethod
131    def capabilities(cls):
132        return RunnerCaps(commands={"flash"}, erase=True, extload=True, tool_opt=True)
133
134    @classmethod
135    def do_add_parser(cls, parser):
136        parser.add_argument(
137            "--port",
138            type=str,
139            required=True,
140            help="Interface identifier, e.g. swd, jtag, /dev/ttyS0...",
141        )
142        parser.add_argument(
143            "--frequency", type=int, required=False, help="Programmer frequency in KHz"
144        )
145        parser.add_argument(
146            "--reset-mode",
147            type=str,
148            required=False,
149            choices=["sw", "hw", "core"],
150            help="Reset mode",
151        )
152        parser.add_argument(
153            "--start-address",
154            # To accept arguments in hex format, a wrapper lambda around int() must be used.
155            # Wrapping the lambda with functools.wraps() makes it so that 'invalid int value'
156            # is displayed when an invalid value is provided for this argument.
157            type=functools.wraps(int)(lambda s: int(s, base=0)),
158            required=False,
159            help="Address where execution should begin after flashing"
160        )
161        parser.add_argument(
162            "--conn-modifiers",
163            type=str,
164            required=False,
165            help="Additional options for the --connect argument",
166        )
167        parser.add_argument(
168            "--cli",
169            type=Path,
170            required=False,
171            help="STM32CubeProgrammer CLI tool path",
172        )
173        parser.add_argument(
174            "--use-elf",
175            action="store_true",
176            required=False,
177            help="Use ELF file when flashing instead of HEX file",
178        )
179
180    @classmethod
181    def extload_help(cls) -> str:
182        return "External Loader for STM32_Programmer_CLI"
183
184    @classmethod
185    def tool_opt_help(cls) -> str:
186        return "Additional options for STM32_Programmer_CLI"
187
188    @classmethod
189    def do_create(
190        cls, cfg: RunnerConfig, args: argparse.Namespace
191    ) -> "STM32CubeProgrammerBinaryRunner":
192        return STM32CubeProgrammerBinaryRunner(
193            cfg,
194            port=args.port,
195            frequency=args.frequency,
196            reset_mode=args.reset_mode,
197            start_address=args.start_address,
198            conn_modifiers=args.conn_modifiers,
199            cli=args.cli,
200            use_elf=args.use_elf,
201            erase=args.erase,
202            extload=args.extload,
203            tool_opt=args.tool_opt,
204        )
205
206    def do_run(self, command: str, **kwargs):
207        if command == "flash":
208            self.flash(**kwargs)
209
210    def flash(self, **kwargs) -> None:
211        self.require(str(self._cli))
212
213        # prepare base command
214        cmd = [str(self._cli)]
215
216        connect_opts = f"port={self._port}"
217        if self._frequency:
218            connect_opts += f" freq={self._frequency}"
219        if self._reset_mode:
220            reset_mode = STM32CubeProgrammerBinaryRunner._RESET_MODES[self._reset_mode]
221            connect_opts += f" reset={reset_mode}"
222        if self._conn_modifiers:
223            connect_opts += f" {self._conn_modifiers}"
224
225        cmd += ["--connect", connect_opts]
226        cmd += self._tool_opt
227        if self._extload:
228            # external loader to come after the tool option in STM32CubeProgrammer
229            cmd += self._extload
230
231        # erase first if requested
232        if self._erase:
233            self.check_call(cmd + ["--erase", "all"])
234
235        # flash image and run application
236        dl_file = self.cfg.elf_file if self._use_elf else self.cfg.hex_file
237        if dl_file is None:
238            raise RuntimeError('cannot flash; no download file was specified')
239        elif not os.path.isfile(dl_file):
240            raise RuntimeError(f'download file {dl_file} does not exist')
241
242        flash_and_run_args = ["--download", dl_file]
243
244        # '--start' is needed to start execution after flash.
245        # The default start address is the beggining of the flash,
246        # but another value can be explicitly specified if desired.
247        flash_and_run_args.append("--start")
248        if self._start_address is not None:
249            flash_and_run_args.append(f"0x{self._start_address:X}")
250
251        self.check_call(cmd + flash_and_run_args)
252