Merge pull request #491 from tadeubas/firmware-fix

Fix firmware.py opening files and not closing
selfcustody · Dec 11, 2024 · 8572c62 · 8572c62
2 parents 97f50ba + e7a90d8
commit 8572c62
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 19 deletions.
diff --git 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,31 +218,34 @@ 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
     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"))
+        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)
@@ -250,24 +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:
-        write_data(
-            lambda pct: status_text(
-                t("Processing..") + "1/3" + "\n\n%d%%" % int(pct * 100)
-            ),
-            new_address,
-            open(firmware_path, "rb", buffering=0),
-            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(
@@ -292,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(

diff --git 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
 ):