feat(write_flash): Allow encrypted writes using key from the Key Manager

Author: Harshal5

esptool/targets/esp32c5.py

@@ -25,6 +25,10 @@ class ESP32C5ROM(ESP32C6ROM):
 
     EFUSE_RD_REG_BASE = EFUSE_BASE + 0x030  # BLOCK0 read base address
 
+    EFUSE_FORCE_USE_KEY_MANAGER_KEY_REG = EFUSE_BASE + 0x34
+    EFUSE_FORCE_USE_KEY_MANAGER_KEY_SHIFT = 10
+    FORCE_USE_KEY_MANAGER_VAL_XTS_AES_KEY = 2
+
     EFUSE_PURPOSE_KEY0_REG = EFUSE_BASE + 0x34
     EFUSE_PURPOSE_KEY0_SHIFT = 22
     EFUSE_PURPOSE_KEY1_REG = EFUSE_BASE + 0x34
@@ -192,7 +196,13 @@ def is_flash_encryption_key_valid(self):
             self.get_key_block_purpose(b) for b in range(self.EFUSE_MAX_KEY + 1)
         ]
 
-        return any(p == self.PURPOSE_VAL_XTS_AES128_KEY for p in purposes)
+        if any(p == self.PURPOSE_VAL_XTS_AES128_KEY for p in purposes):
+            return True
+
+        return (
+            self.read_reg(self.EFUSE_FORCE_USE_KEY_MANAGER_KEY_REG)
+            >> self.EFUSE_FORCE_USE_KEY_MANAGER_KEY_SHIFT
+        ) & self.FORCE_USE_KEY_MANAGER_VAL_XTS_AES_KEY
 
     def check_spi_connection(self, spi_connection):
         if not set(spi_connection).issubset(set(range(0, 29))):

esptool/targets/esp32p4.py

@@ -43,6 +43,10 @@ class ESP32P4ROM(ESP32ROM):
 
     EFUSE_RD_REG_BASE = EFUSE_BASE + 0x030  # BLOCK0 read base address
 
+    EFUSE_FORCE_USE_KEY_MANAGER_KEY_REG = EFUSE_BASE + 0x34
+    EFUSE_FORCE_USE_KEY_MANAGER_KEY_SHIFT = 9
+    FORCE_USE_KEY_MANAGER_VAL_XTS_AES_KEY = 2
+
     EFUSE_PURPOSE_KEY0_REG = EFUSE_BASE + 0x34
     EFUSE_PURPOSE_KEY0_SHIFT = 24
     EFUSE_PURPOSE_KEY1_REG = EFUSE_BASE + 0x34
@@ -216,9 +220,15 @@ def is_flash_encryption_key_valid(self):
         if any(p == self.PURPOSE_VAL_XTS_AES128_KEY for p in purposes):
             return True
 
-        return any(p == self.PURPOSE_VAL_XTS_AES256_KEY_1 for p in purposes) and any(
+        if any(p == self.PURPOSE_VAL_XTS_AES256_KEY_1 for p in purposes) and any(
             p == self.PURPOSE_VAL_XTS_AES256_KEY_2 for p in purposes
-        )
+        ):
+            return True
+
+        return (
+            self.read_reg(self.EFUSE_FORCE_USE_KEY_MANAGER_KEY_REG)
+            >> self.EFUSE_FORCE_USE_KEY_MANAGER_KEY_SHIFT
+        ) & self.FORCE_USE_KEY_MANAGER_VAL_XTS_AES_KEY
 
     def change_baud(self, baud):
         ESPLoader.change_baud(self, baud)