add more methods for writing a config to file

these allow for adding a config to the proper section of an existing
binary file, or writing a config to a new config-only binary file

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 #
 ############

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

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)