diff --git a/gp2040ce_bintools/builder.py b/gp2040ce_bintools/builder.py index 84fa44e..90cf8e1 100644 --- a/gp2040ce_bintools/builder.py +++ b/gp2040ce_bintools/builder.py @@ -1,9 +1,13 @@ """Build binary files for a GP2040-CE board.""" import argparse +import copy import logging +from google.protobuf.message import Message + from gp2040ce_bintools import core_parser -from gp2040ce_bintools.storage import STORAGE_LOCATION, pad_config_to_storage_size +from gp2040ce_bintools.storage import (STORAGE_LOCATION, STORAGE_SIZE, pad_config_to_storage_size, + serialize_config_with_footer) logger = logging.getLogger(__name__) @@ -62,6 +66,52 @@ def pad_firmware_up_to_storage(firmware: bytes) -> bytearray: return bytearray(firmware) + bytearray(b'\x00' * bytes_to_pad) +def replace_config_in_binary(board_binary: bytearray, config_binary: bytearray) -> bytearray: + """Given (presumed) whole board and config binaries, combine the two to one, with proper offsets for GP2040-CE. + + Whatever is in the board binary is not sanity checked, and is overwritten. If it is + too small to be a board dump, it is nonetheless expanded and overwritten to fit the + proper size. + + Args: + board_binary: binary data of a whole board dump from a GP2040-CE board + config_binary: binary data of board config + footer, possibly padded to be a full storage section + Returns: + the resulting correctly-offset binary suitable for a GP2040-CE board + """ + if len(board_binary) < STORAGE_LOCATION + STORAGE_SIZE: + # this is functionally the same, since this doesn't sanity check the firmware + return combine_firmware_and_config(board_binary, config_binary) + else: + new_binary = bytearray(copy.copy(board_binary)) + new_binary[STORAGE_LOCATION:(STORAGE_LOCATION + STORAGE_SIZE)] = pad_config_to_storage_size(config_binary) + return new_binary + + +def write_new_config_to_filename(config: Message, filename: str, inject: bool = False) -> None: + """Serialize the provided config to the specified file. + + The file may be replaced, creating a configuration section-only binary, or appended to + an existing file that is grown to place the config section in the proper location. + + Args: + config: the Protobuf configuration to write to disk + filename: the filename to write the serialized configuration to + inject: if True, the file is read in and has its storage section replaced; if False, + the whole file is replaced + """ + if inject: + config_binary = serialize_config_with_footer(config) + with open(filename, 'rb') as file: + existing_binary = file.read() + binary = replace_config_in_binary(bytearray(existing_binary), config_binary) + else: + binary = serialize_config_with_footer(config) + + with open(filename, 'wb') as file: + file.write(binary) + + ############ # COMMANDS # ############ diff --git a/tests/test_builder.py b/tests/test_builder.py index c50722f..f1f8484 100644 --- a/tests/test_builder.py +++ b/tests/test_builder.py @@ -1,8 +1,28 @@ """Tests for the image builder module.""" -import pytest +import os +import sys -from gp2040ce_bintools.builder import FirmwareLengthError, combine_firmware_and_config, pad_firmware_up_to_storage -from gp2040ce_bintools.storage import get_config_footer, get_storage_section +import pytest +from decorator import decorator + +from gp2040ce_bintools import get_config_pb2 +from gp2040ce_bintools.builder import (FirmwareLengthError, combine_firmware_and_config, pad_firmware_up_to_storage, + replace_config_in_binary, write_new_config_to_filename) +from gp2040ce_bintools.storage import get_config, get_config_footer, get_storage_section + +HERE = os.path.dirname(os.path.abspath(__file__)) + + +@decorator +def with_pb2s(test, *args, **kwargs): + """Wrap a test with precompiled pb2 files on the path.""" + proto_path = os.path.join(HERE, 'test-files', 'pb2-files') + sys.path.append(proto_path) + + test(*args, **kwargs) + + sys.path.pop() + del sys.modules['config_pb2'] def test_padding_firmware(firmware_binary): @@ -29,7 +49,88 @@ def test_firmware_plus_config_binary(firmware_binary, config_binary): assert footer_size == 2032 +def test_replace_config_in_binary(config_binary): + """Test that a config binary is placed in the storage location of a source binary to overwrite.""" + whole_board = replace_config_in_binary(bytearray(b'\x00' * 3 * 1024 * 1024), config_binary) + assert len(whole_board) == 3 * 1024 * 1024 + # if this is valid, we should be able to find the storage and footer again + storage = get_storage_section(whole_board) + footer_size, _, _ = get_config_footer(storage) + assert footer_size == 2032 + + +def test_replace_config_in_binary_not_big_enough(config_binary): + """Test that a config binary is placed in the storage location of a source binary to pad.""" + whole_board = replace_config_in_binary(bytearray(b'\x00' * 1 * 1024 * 1024), config_binary) + assert len(whole_board) == 2 * 1024 * 1024 + # if this is valid, we should be able to find the storage and footer again + storage = get_storage_section(whole_board) + footer_size, _, _ = get_config_footer(storage) + assert footer_size == 2032 + + def test_padding_firmware_too_big(firmware_binary): """Test that firmware is padded to the expected size.""" with pytest.raises(FirmwareLengthError): _ = pad_firmware_up_to_storage(firmware_binary + firmware_binary + firmware_binary) + + +@with_pb2s +def test_write_new_config_to_whole_board(whole_board_dump, tmp_path): + """Test that the config can be overwritten on a whole board dump.""" + tmp_file = os.path.join(tmp_path, 'whole-board-dump-copy.bin') + with open(tmp_file, 'wb') as file: + file.write(whole_board_dump) + # reread just in case + with open(tmp_file, 'rb') as file: + board_dump = file.read() + + config = get_config(get_storage_section(board_dump)) + assert config.boardVersion == 'v0.7.2' + config.boardVersion = 'v0.7.2-COOL' + write_new_config_to_filename(config, tmp_file, inject=True) + + # read new file + with open(tmp_file, 'rb') as file: + new_board_dump = file.read() + config = get_config(get_storage_section(new_board_dump)) + assert config.boardVersion == 'v0.7.2-COOL' + assert len(board_dump) == len(new_board_dump) + + +@with_pb2s +def test_write_new_config_to_firmware(firmware_binary, tmp_path): + """Test that the config can be added on a firmware.""" + tmp_file = os.path.join(tmp_path, 'firmware-copy.bin') + with open(tmp_file, 'wb') as file: + file.write(firmware_binary) + + config_pb2 = get_config_pb2() + config = config_pb2.Config() + config.boardVersion = 'v0.7.2-COOL' + write_new_config_to_filename(config, tmp_file, inject=True) + + # read new file + with open(tmp_file, 'rb') as file: + new_board_dump = file.read() + config = get_config(get_storage_section(new_board_dump)) + assert config.boardVersion == 'v0.7.2-COOL' + assert len(new_board_dump) == 2 * 1024 * 1024 + + +@with_pb2s +def test_write_new_config_to_config_bin(firmware_binary, tmp_path): + """Test that the config can be written to a file.""" + tmp_file = os.path.join(tmp_path, 'config.bin') + config_pb2 = get_config_pb2() + config = config_pb2.Config() + config.boardVersion = 'v0.7.2-COOL' + write_new_config_to_filename(config, tmp_file) + + # read new file + with open(tmp_file, 'rb') as file: + config_dump = file.read() + config = get_config(config_dump) + config_size, _, _ = get_config_footer(config_dump) + assert config.boardVersion == 'v0.7.2-COOL' + assert len(config_dump) == config_size + 12 diff --git a/tests/test_storage.py b/tests/test_storage.py index f3c2793..c8ac44c 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -59,7 +59,7 @@ def test_config_footer_bad_magic(storage_dump): def test_config_footer_bad_crc(storage_dump): """Test that a config footer isn't detected if the CRC checksums don't match.""" corrupt = bytearray(storage_dump) - corrupt[-50:-40] = bytearray(0*10) + corrupt[-50:-40] = bytearray(0 * 10) with pytest.raises(storage.ConfigCrcError): _, _, _ = storage.get_config_footer(corrupt)