From 4208707088b5b5475a98774e12f0b297b2a46487 Mon Sep 17 00:00:00 2001 From: matejcik Date: Wed, 31 Jul 2024 17:37:29 +0200 Subject: [PATCH] refactor(core/tools): make combine_firmware nicer --- .../trezor_core_tools/combine_firmware.py | 72 +++++++++++-------- core/tools/trezor_core_tools/common.py | 21 ++++++ core/tools/trezor_core_tools/layout_parser.py | 68 ++++++++++-------- 3 files changed, 104 insertions(+), 57 deletions(-) create mode 100644 core/tools/trezor_core_tools/common.py diff --git a/core/tools/trezor_core_tools/combine_firmware.py b/core/tools/trezor_core_tools/combine_firmware.py index d52d4743e3..7477a60685 100755 --- a/core/tools/trezor_core_tools/combine_firmware.py +++ b/core/tools/trezor_core_tools/combine_firmware.py @@ -3,24 +3,16 @@ from __future__ import annotations import datetime import io +import sys from pathlib import Path -import os -import subprocess import click -def get_layout_params(layout: Path, name: str) -> int: - directory = os.path.dirname(os.path.realpath(__file__)) - with subprocess.Popen(args=["python", Path(directory, "layout_parser.py"), str(layout), name], - stdout=subprocess.PIPE) as script: - return int(script.stdout.read().decode().strip()) +from .layout_parser import find_all_values + @click.command() -@click.argument( - "layout", - type=click.Path(dir_okay=False, writable=True, path_type=Path), - required=True, -) +@click.argument("model") @click.argument( "outfile", type=click.Path(dir_okay=False, writable=True, path_type=Path), @@ -32,37 +24,59 @@ def get_layout_params(layout: Path, name: str) -> int: type=(str, click.Path(exists=True, dir_okay=False, readable=True, path_type=Path)), multiple=True, ) -def main( - layout: Path, - bin: List[Tuple[Path, str]], - outfile: Path | None, +def main(model: str, bin: list[tuple[str, Path]], outfile: Path | None) -> None: + """Create a combined.bin file from components. -) -> None: + MODEL is the internal model name (e.g. T1B1, T2T1, T2B1). + BIN is a list of tuples with the type and path to the binary file. The usual types + are: + + \b + * BOARDLOADER + * BOOTLOADER + * FIRMWARE + + For example: + + \b + $ combine_firmware T3T1 -b boardloader build/boardloader.bin -b bootloader build/bootloader.bin -b firmware build/firmware.bin + """ if outfile is None: today = datetime.date.today().strftime(r"%Y-%m-%d") outfile = Path(f"combined-{today}.bin") - first_bin = bin[0] - (name, bin_path) = first_bin + offset = 0 + out_buf = io.BytesIO() - start_offset = get_layout_params(layout, name+ "_START") + all_values = find_all_values(model) - offset = start_offset - out_bytes = io.BytesIO() + def find_value(name: str) -> int: + value_name = f"{name.upper()}_START" + if value_name not in all_values: + click.echo(f"ERROR: component {name} not found in layout for model {model}") + click.echo("Try one of: boardloader, bootloader, firmware") + sys.exit(1) + return all_values[value_name] - for (name, bin_path) in bin: - bin_start = get_layout_params(layout, name + "_START") - # zero-pad until next section: - offset += out_bytes.write(b"\x00" * (bin_start - offset)) + for name, bin_path in bin: + bin_start = find_value(name) + + if not offset: + # initialize offset + offset = bin_start + else: + # pad until next section + offset += out_buf.write(b"\x00" * (bin_start - offset)) assert offset == bin_start # write binary - offset += out_bytes.write(bin_path.read_bytes()) + offset += out_buf.write(bin_path.read_bytes()) # write out contents - click.echo(f"Writing {outfile} ({offset - start_offset} bytes)") - outfile.write_bytes(out_bytes.getvalue()) + out_bytes = out_buf.getvalue() + click.echo(f"Writing {outfile} ({len(out_bytes)} bytes)") + outfile.write_bytes(out_bytes) if __name__ == "__main__": diff --git a/core/tools/trezor_core_tools/common.py b/core/tools/trezor_core_tools/common.py new file mode 100644 index 0000000000..7a4f2f5be7 --- /dev/null +++ b/core/tools/trezor_core_tools/common.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from pathlib import Path + +HERE = Path(__file__).parent.resolve() +CORE = HERE.parent.parent + +MODELS_DIR = CORE / "embed" / "models" + +MODELS_DICT = { + "1": "T1B1", + "T": "T2T1", + "R": "T2B1", + "DISC1": "D001", + "DISC2": "D002", +} + + +def get_layout_for_model(model: str) -> Path: + model = MODELS_DICT.get(model, model) + return MODELS_DIR / model / f"model_{model}.h" diff --git a/core/tools/trezor_core_tools/layout_parser.py b/core/tools/trezor_core_tools/layout_parser.py index 0b5adfd1fe..13b2e43303 100644 --- a/core/tools/trezor_core_tools/layout_parser.py +++ b/core/tools/trezor_core_tools/layout_parser.py @@ -1,40 +1,52 @@ #!/usr/bin/env python3 from __future__ import annotations - -from pathlib import Path +import re import click import ast +from .common import get_layout_for_model + + +# the following regular expression matches a thing looking like those examples: +# #define HEADER_START 0x123456 +# #define HEADER_START (1 * 2 * 3) // comment +# and returns two groups: the name and the value. Comment is ignored. +SEARCH_PATTERN = r"#define (\w+) (.+?)(?:\s*//.*)?$" + + +def find_all_values(model: str) -> dict[str, int]: + layout = get_layout_for_model(model) + values = {} + for line in open(layout): + match = re.match(SEARCH_PATTERN, line) + if match is not None: + name, value = match.groups() + try: + node = ast.parse(value, mode="eval") + parsed_value = eval(compile(node, "", "eval")) + values[name] = int(parsed_value) + except Exception: + pass + return values + + +def find_value(model: str, name: str) -> int: + all_values = find_all_values(model) + if name not in all_values: + raise ValueError(f"Value {name} not found in layout for model {model}") + return all_values[name] + @click.command() -@click.argument( - "layout", - type=click.Path(exists=True, dir_okay=False, readable=True, path_type=Path), -) -@click.argument( - "name", - type=click.STRING, -) -def main( - layout: Path, - name: str, -) -> None: - - search = f'#define {name} ' - for line in layout.read_text().splitlines(): - - if line.startswith(search): - line = line.split(search)[1] - if line.startswith("("): - line = line.split("(")[1].split(")")[-2] - if "//" in line: - line = line.split("//")[0] - line = line.strip() - node = ast.parse(line, mode='eval') - result = eval(compile(node, '', 'eval')) - print(result) +@click.argument("model") +@click.argument("name") +def main(model: str, name: str) -> None: + try: + print(find_value(model, name)) + except ValueError as e: + raise click.ClickException(str(e)) if __name__ == "__main__":