From 7f242e6c69c5807872192f763f9b6fb12a2dca17 Mon Sep 17 00:00:00 2001 From: tadeubas Date: Mon, 9 Dec 2024 01:06:53 -0300 Subject: [PATCH 1/3] Fix firmware.py opening files and not closing --- src/krux/firmware.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/krux/firmware.py b/src/krux/firmware.py index 3d99ca09..b45fa573 100644 --- a/src/krux/firmware.py +++ b/src/krux/firmware.py @@ -227,7 +227,8 @@ def upgrade(): sig = None try: - sig = open(firmware_path + ".sig", "rb").read() + with open(firmware_path + ".sig", "rb") as sig_file: + sig = sig_file.read() except: display.flash_text(t("Missing signature file")) return False @@ -258,17 +259,19 @@ def status_text(text): display.clear() display.draw_centered_text(text) + firmware_file = open(firmware_path, "rb", buffering=0) write_data( lambda pct: status_text( t("Processing..") + "1/3" + "\n\n%d%%" % int(pct * 100) ), new_address, - open(firmware_path, "rb", buffering=0), + firmware_file, new_size, 65536, True, firmware_with_header_hash, ) + firmware_file.close() write_data( lambda pct: status_text( From 7f6a0807946757f1af0360f5ee265001b44e8db1 Mon Sep 17 00:00:00 2001 From: tadeubas Date: Mon, 9 Dec 2024 14:22:46 -0300 Subject: [PATCH 2/3] open file using with + color for error --- src/krux/firmware.py | 42 ++++++++++++++++++++++-------------------- 1 file changed, 22 insertions(+), 20 deletions(-) diff --git a/src/krux/firmware.py b/src/krux/firmware.py index af610391..95a175ef 100644 --- a/src/krux/firmware.py +++ b/src/krux/firmware.py @@ -31,6 +31,7 @@ from .display import display from .krux_settings import t from .wdt import wdt +from .themes import theme FLASH_SIZE = 2**24 MAX_FIRMWARE_SIZE = 0x300000 @@ -217,14 +218,16 @@ def status_text(text): return False if new_size > MAX_FIRMWARE_SIZE: - display.flash_text("Firmware exceeds max size: %d" % MAX_FIRMWARE_SIZE) + display.flash_text( + "Firmware exceeds max size: %d" % MAX_FIRMWARE_SIZE, theme.error_color + ) return False pubkey = None try: pubkey = ec.PublicKey.from_string(SIGNER_PUBKEY) except: - display.flash_text("Invalid public key") + display.flash_text("Invalid public key", theme.error_color) return False sig = None @@ -232,17 +235,17 @@ def status_text(text): with open(firmware_path + ".sig", "rb") as sig_file: sig = sig_file.read() except: - display.flash_text(t("Missing signature file")) + display.flash_text(t("Missing signature file"), theme.error_color) return False try: # Parse, serialize, and reparse to ensure signature is compact prior to verification sig = ec.Signature.parse(ec.Signature.parse(sig).serialize()) if not pubkey.verify(sig, firmware_hash): - display.flash_text(t("Bad signature")) + display.flash_text(t("Bad signature"), theme.error_color) return False except: - display.flash_text(t("Bad signature")) + display.flash_text(t("Bad signature"), theme.error_color) return False boot_config_sector = flash.read(MAIN_BOOT_CONFIG_SECTOR_ADDRESS, 4096) @@ -251,25 +254,25 @@ def status_text(text): boot_config_sector = flash.read(BACKUP_BOOT_CONFIG_SECTOR_ADDRESS, 4096) address, _, entry_index = find_active_firmware(boot_config_sector) if address is None: - display.flash_text("Invalid bootloader") + display.flash_text("Invalid bootloader", theme.error_color) return False # Write new firmware to the opposite slot new_address = FIRMWARE_SLOT_2 if address == FIRMWARE_SLOT_1 else FIRMWARE_SLOT_1 try: - firmware_file = open(firmware_path, "rb", buffering=0) - write_data( - lambda pct: status_text( - t("Processing..") + "1/3" + "\n\n%d%%" % int(pct * 100) - ), - new_address, - firmware_file, - new_size, - 65536, - True, - firmware_with_header_hash, - ) + with open(firmware_path, "rb", buffering=0) as firmware_file: + write_data( + lambda pct: status_text( + t("Processing..") + "1/3" + "\n\n%d%%" % int(pct * 100) + ), + new_address, + firmware_file, + new_size, + 65536, + True, + firmware_with_header_hash, + ) write_data( lambda pct: status_text( @@ -280,7 +283,6 @@ def status_text(text): len(boot_config_sector), 4096, ) - firmware_file.close() new_boot_config_sector = update_boot_config_sector( boot_config_sector, entry_index, new_address, new_size @@ -295,7 +297,7 @@ def status_text(text): 4096, ) except: - display.flash_text("Error read/write data") + display.flash_text("Error read/write data", theme.error_color) return False status_text( From e7a90d875c330edc43277b018e5dc4a556a69ea4 Mon Sep 17 00:00:00 2001 From: tadeubas Date: Tue, 10 Dec 2024 19:36:22 -0300 Subject: [PATCH 3/3] added test case for new code branch --- tests/test_firmware.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/test_firmware.py b/tests/test_firmware.py index da1b8762..644b8880 100644 --- a/tests/test_firmware.py +++ b/tests/test_firmware.py @@ -12901,6 +12901,48 @@ def test_upgrade(mocker, m5stickv, mock_success_input_cls, tdata): ) +def test_upgrade_fails_write_data(mocker, m5stickv, mock_success_input_cls, tdata): + from krux import firmware + + mocker.patch( + "builtins.open", + new=get_mock_open( + { + "/sd/firmware-v0.0.0.bin": tdata.TEST_FIRMWARE, + "/sd/firmware-v0.0.0.bin.sig": tdata.TEST_FIRMWARE_SIG, + } + ), + ) + mocker.patch( + "os.listdir", + new=mocker.MagicMock( + return_value=["firmware-v0.0.0.bin", "firmware-v0.0.0.bin.sig"] + ), + ) + mocker.patch("krux.firmware.Input", new=mock_success_input_cls) + mocker.patch( + "krux.firmware.flash.read", + new=mocker.MagicMock( + return_value=bytes(tdata.SECTOR_WITH_ACTIVE_FIRMWARE_AT_INDEX_1_SLOT_1) + ), + ) + mocker.patch( + "krux.firmware.ec.PublicKey.from_string", + new=mocker.MagicMock(return_value=tdata.TEST_SIGNER_PUBLIC_KEY), + ) + + mocker.patch.object(firmware, "write_data", side_effect=ValueError) + mocker.spy(firmware, "display") + + assert not firmware.upgrade() + + from krux.themes import theme + + firmware.display.flash_text.assert_called_with( + "Error read/write data", theme.error_color + ) + + def test_upgrade_uses_backup_sector_when_main_sector_is_missing_active_firmware( mocker, m5stickv, mock_success_input_cls, tdata ):