1# SPDX-FileCopyrightText: 2016-2022 Espressif Systems (Shanghai) CO LTD
2#
3# SPDX-License-Identifier: GPL-2.0-or-later
4# PYTHON_ARGCOMPLETE_OK
5
6import argparse
7import os
8import sys
9from collections import namedtuple
10from io import StringIO
11
12import espefuse.efuse.esp32 as esp32_efuse
13import espefuse.efuse.esp32c2 as esp32c2_efuse
14import espefuse.efuse.esp32c3 as esp32c3_efuse
15import espefuse.efuse.esp32c5 as esp32c5_efuse
16import espefuse.efuse.esp32c5beta3 as esp32c5beta3_efuse
17import espefuse.efuse.esp32c6 as esp32c6_efuse
18import espefuse.efuse.esp32c61 as esp32c61_efuse
19import espefuse.efuse.esp32h2 as esp32h2_efuse
20import espefuse.efuse.esp32h2beta1 as esp32h2beta1_efuse
21import espefuse.efuse.esp32p4 as esp32p4_efuse
22import espefuse.efuse.esp32s2 as esp32s2_efuse
23import espefuse.efuse.esp32s3 as esp32s3_efuse
24import espefuse.efuse.esp32s3beta2 as esp32s3beta2_efuse
25
26import esptool
27
28DefChip = namedtuple("DefChip", ["chip_name", "efuse_lib", "chip_class"])
29
30SUPPORTED_BURN_COMMANDS = [
31    "read_protect_efuse",
32    "write_protect_efuse",
33    "burn_efuse",
34    "burn_block_data",
35    "burn_bit",
36    "burn_key",
37    "burn_key_digest",
38    "burn_custom_mac",
39    "set_flash_voltage",
40    "execute_scripts",
41]
42
43SUPPORTED_READ_COMMANDS = [
44    "summary",
45    "dump",
46    "get_custom_mac",
47    "adc_info",
48    "check_error",
49]
50
51SUPPORTED_COMMANDS = SUPPORTED_READ_COMMANDS + SUPPORTED_BURN_COMMANDS
52
53SUPPORTED_CHIPS = {
54    "esp32": DefChip("ESP32", esp32_efuse, esptool.targets.ESP32ROM),
55    "esp32c2": DefChip("ESP32-C2", esp32c2_efuse, esptool.targets.ESP32C2ROM),
56    "esp32c3": DefChip("ESP32-C3", esp32c3_efuse, esptool.targets.ESP32C3ROM),
57    "esp32c6": DefChip("ESP32-C6", esp32c6_efuse, esptool.targets.ESP32C6ROM),
58    "esp32c61": DefChip("ESP32-C61", esp32c61_efuse, esptool.targets.ESP32C61ROM),
59    "esp32c5": DefChip("ESP32-C5", esp32c5_efuse, esptool.targets.ESP32C5ROM),
60    "esp32c5beta3": DefChip(
61        "ESP32-C5(beta3)", esp32c5beta3_efuse, esptool.targets.ESP32C5BETA3ROM
62    ),
63    "esp32h2": DefChip("ESP32-H2", esp32h2_efuse, esptool.targets.ESP32H2ROM),
64    "esp32p4": DefChip("ESP32-P4", esp32p4_efuse, esptool.targets.ESP32P4ROM),
65    "esp32h2beta1": DefChip(
66        "ESP32-H2(beta1)", esp32h2beta1_efuse, esptool.targets.ESP32H2BETA1ROM
67    ),
68    "esp32s2": DefChip("ESP32-S2", esp32s2_efuse, esptool.targets.ESP32S2ROM),
69    "esp32s3": DefChip("ESP32-S3", esp32s3_efuse, esptool.targets.ESP32S3ROM),
70    "esp32s3beta2": DefChip(
71        "ESP32-S3(beta2)", esp32s3beta2_efuse, esptool.targets.ESP32S3BETA2ROM
72    ),
73}
74
75
76def get_esp(
77    port,
78    baud,
79    connect_mode,
80    chip="auto",
81    skip_connect=False,
82    virt=False,
83    debug=False,
84    virt_efuse_file=None,
85):
86    if chip not in ["auto"] + list(SUPPORTED_CHIPS.keys()):
87        raise esptool.FatalError("get_esp: Unsupported chip (%s)" % chip)
88    if virt:
89        efuse = SUPPORTED_CHIPS.get(chip, SUPPORTED_CHIPS["esp32"]).efuse_lib
90        esp = efuse.EmulateEfuseController(virt_efuse_file, debug)
91    else:
92        if chip == "auto" and not skip_connect:
93            esp = esptool.cmds.detect_chip(port, baud, connect_mode)
94        else:
95            esp = SUPPORTED_CHIPS.get(chip, SUPPORTED_CHIPS["esp32"]).chip_class(
96                port if not skip_connect else StringIO(), baud
97            )
98            if not skip_connect:
99                esp.connect(connect_mode)
100                if esp.sync_stub_detected:
101                    esp = esp.STUB_CLASS(esp)
102    return esp
103
104
105def get_efuses(
106    esp,
107    skip_connect=False,
108    debug_mode=False,
109    do_not_confirm=False,
110    extend_efuse_table=None,
111):
112    for name in SUPPORTED_CHIPS:
113        if SUPPORTED_CHIPS[name].chip_name == esp.CHIP_NAME:
114            efuse = SUPPORTED_CHIPS[name].efuse_lib
115            return (
116                efuse.EspEfuses(
117                    esp, skip_connect, debug_mode, do_not_confirm, extend_efuse_table
118                ),
119                efuse.operations,
120            )
121    else:
122        raise esptool.FatalError("get_efuses: Unsupported chip (%s)" % esp.CHIP_NAME)
123
124
125def split_on_groups(all_args):
126    """
127    This function splits the all_args list into groups,
128    where each item is a cmd with all its args.
129
130    Example:
131    all_args:
132    ['burn_key_digest', 'secure_images/ecdsa256_secure_boot_signing_key_v2.pem',
133     'burn_key', 'BLOCK_KEY0', 'images/efuse/128bit_key',
134     'XTS_AES_128_KEY_DERIVED_FROM_128_EFUSE_BITS']
135
136    used_cmds: ['burn_key_digest', 'burn_key']
137    groups:
138    [['burn_key_digest', 'secure_images/ecdsa256_secure_boot_signing_key_v2.pem'],
139     ['burn_key', 'BLOCK_KEY0', 'images/efuse/128bit_key',
140      'XTS_AES_128_KEY_DERIVED_FROM_128_EFUSE_BITS']]
141    """
142
143    groups = []
144    cmd = []
145    used_cmds = []
146    for item in all_args:
147        if item in SUPPORTED_COMMANDS:
148            used_cmds.append(item)
149            if cmd != []:
150                groups.append(cmd)
151            cmd = []
152        cmd.append(item)
153    if cmd:
154        groups.append(cmd)
155    return groups, used_cmds
156
157
158def main(custom_commandline=None, esp=None):
159    """
160    Main function for espefuse
161
162    custom_commandline - Optional override for default arguments parsing
163    (that uses sys.argv), can be a list of custom arguments as strings.
164    Arguments and their values need to be added as individual items to the list
165    e.g. "--port /dev/ttyUSB1" thus becomes ['--port', '/dev/ttyUSB1'].
166
167    esp - Optional override of the connected device previously
168    returned by esptool.get_default_connected_device()
169    """
170
171    external_esp = esp is not None
172
173    init_parser = argparse.ArgumentParser(
174        description="espefuse.py v%s - [ESP32xx] efuse get/set tool"
175        % esptool.__version__,
176        prog="espefuse",
177        add_help=False,
178    )
179
180    init_parser.add_argument(
181        "--chip",
182        "-c",
183        help="Target chip type",
184        choices=["auto"] + list(SUPPORTED_CHIPS.keys()),
185        default=os.environ.get("ESPTOOL_CHIP", "auto"),
186    )
187
188    init_parser.add_argument(
189        "--baud",
190        "-b",
191        help="Serial port baud rate used when flashing/reading",
192        type=esptool.arg_auto_int,
193        default=os.environ.get("ESPTOOL_BAUD", esptool.loader.ESPLoader.ESP_ROM_BAUD),
194    )
195
196    init_parser.add_argument(
197        "--port",
198        "-p",
199        help="Serial port device",
200        default=os.environ.get("ESPTOOL_PORT", esptool.loader.ESPLoader.DEFAULT_PORT),
201    )
202
203    init_parser.add_argument(
204        "--before",
205        help="What to do before connecting to the chip",
206        choices=["default_reset", "usb_reset", "no_reset", "no_reset_no_sync"],
207        default="default_reset",
208    )
209
210    init_parser.add_argument(
211        "--debug",
212        "-d",
213        help="Show debugging information (loglevel=DEBUG)",
214        action="store_true",
215    )
216    init_parser.add_argument(
217        "--virt",
218        help="For host tests, the tool will work in the virtual mode "
219        "(without connecting to a chip).",
220        action="store_true",
221    )
222    init_parser.add_argument(
223        "--path-efuse-file",
224        help="For host tests, saves efuse memory to file.",
225        type=str,
226        default=None,
227    )
228    init_parser.add_argument(
229        "--do-not-confirm",
230        help="Do not pause for confirmation before permanently writing efuses. "
231        "Use with caution.",
232        action="store_true",
233    )
234    init_parser.add_argument(
235        "--postpone",
236        help="Postpone burning some efuses from BLOCK0 at the end, "
237        "(efuses which disable access to blocks or chip).",
238        action="store_true",
239    )
240    init_parser.add_argument(
241        "--extend-efuse-table",
242        help="CSV file from ESP-IDF (esp_efuse_custom_table.csv)",
243        type=argparse.FileType("r"),
244        default=None,
245    )
246
247    common_args, remaining_args = init_parser.parse_known_args(custom_commandline)
248    debug_mode = common_args.debug
249    just_print_help = [
250        True for arg in remaining_args if arg in ["--help", "-h"]
251    ] or remaining_args == []
252
253    print("espefuse.py v{}".format(esptool.__version__))
254
255    if not external_esp:
256        try:
257            esp = get_esp(
258                common_args.port,
259                common_args.baud,
260                common_args.before,
261                common_args.chip,
262                just_print_help,
263                common_args.virt,
264                common_args.debug,
265                common_args.path_efuse_file,
266            )
267        except esptool.FatalError as e:
268            raise esptool.FatalError(
269                f"{e}\nPlease make sure that you have specified "
270                "the right port with the --port argument"
271            )
272            # TODO: Require the --port argument in the next major release, ESPTOOL-490
273
274    efuses, efuse_operations = get_efuses(
275        esp,
276        just_print_help,
277        debug_mode,
278        common_args.do_not_confirm,
279        common_args.extend_efuse_table,
280    )
281
282    parser = argparse.ArgumentParser(parents=[init_parser])
283    subparsers = parser.add_subparsers(
284        dest="operation", help="Run espefuse.py {command} -h for additional help"
285    )
286
287    efuse_operations.add_commands(subparsers, efuses)
288
289    # Enable argcomplete only on Unix-like systems
290    if sys.platform != "win32":
291        try:
292            import argcomplete
293
294            argcomplete.autocomplete(parser)
295        except ImportError:
296            pass
297
298    grouped_remaining_args, used_cmds = split_on_groups(remaining_args)
299    if len(grouped_remaining_args) == 0:
300        parser.print_help()
301        parser.exit(1)
302    there_are_multiple_burn_commands_in_args = (
303        sum(cmd in SUPPORTED_BURN_COMMANDS for cmd in used_cmds) > 1
304    )
305    if there_are_multiple_burn_commands_in_args:
306        efuses.batch_mode_cnt += 1
307
308    efuses.postpone = common_args.postpone
309
310    try:
311        for rem_args in grouped_remaining_args:
312            args, unused_args = parser.parse_known_args(rem_args, namespace=common_args)
313            if args.operation is None:
314                parser.print_help()
315                parser.exit(1)
316            assert (
317                len(unused_args) == 0
318            ), 'Not all commands were recognized "{}"'.format(unused_args)
319
320            operation_func = vars(efuse_operations)[args.operation]
321            # each 'operation' is a module-level function of the same name
322            print('\n=== Run "{}" command ==='.format(args.operation))
323
324            if hasattr(args, "show_sensitive_info"):
325                if args.show_sensitive_info or args.debug:
326                    args.show_sensitive_info = True
327                else:
328                    print("Sensitive data will be hidden (see --show-sensitive-info)")
329
330            operation_func(esp, efuses, args)
331
332        if there_are_multiple_burn_commands_in_args:
333            efuses.batch_mode_cnt -= 1
334            if not efuses.burn_all(check_batch_mode=True):
335                raise esptool.FatalError("BURN was not done")
336            print("Successful")
337
338        if (
339            sum(cmd in SUPPORTED_BURN_COMMANDS for cmd in used_cmds) > 0
340            and sum(cmd in SUPPORTED_READ_COMMANDS for cmd in used_cmds) > 0
341        ):
342            # [burn_cmd1] [burn_cmd2] [read_cmd1] [burn_cmd3] [read_cmd2]
343            print("\n=== Run read commands after burn commands ===")
344            for rem_args in grouped_remaining_args:
345                args, unused_args = parser.parse_known_args(
346                    rem_args, namespace=common_args
347                )
348                current_cmd = args.operation
349                if current_cmd in SUPPORTED_READ_COMMANDS:
350                    print(f"\n=== Run {args.operation} command ===")
351                    operation_func = vars(efuse_operations)[current_cmd]
352                    operation_func(esp, efuses, args)
353    finally:
354        if not external_esp and not common_args.virt and esp._port:
355            esp._port.close()
356
357
358def _main():
359    try:
360        main()
361    except esptool.FatalError as e:
362        print("\nA fatal error occurred: %s" % e)
363        sys.exit(2)
364
365
366if __name__ == "__main__":
367    _main()
368