fix(write_flash): Make write flash mem independent

* Remove reading of magic register to avoid issues with different memory layouts.

Author: jedi7

esptool/cmds.py

@@ -905,10 +905,21 @@ def write_flash(
                     raise  # Reconnect limit reached
 
         if esp.IS_STUB:
+            # Get the "encrypted" flag for the last file flashed
+            # Note: all_files list contains quadruplets like:
+            # (address: int, filename: str | None, data: bytes, encrypted: bool)
+            last_file_encrypted = all_files[-1][3]
+
             # Stub only writes each block to flash after 'ack'ing the receive,
-            # so do a final dummy operation which will not be 'ack'ed
+            # so do a final operation which will not be 'ack'ed
             # until the last block has actually been written out to flash
-            esp.read_reg(ESPLoader.CHIP_DETECT_MAGIC_REG_ADDR, timeout=timeout)
+            if compress and not last_file_encrypted:
+                esp.flash_defl_finish(reboot=False, timeout=timeout)
+            else:
+                esp.flash_finish(reboot=False, timeout=timeout)
+
+        # Skip sending flash_finish to ROM loader here,
+        # as it causes the loader to exit and run user code
 
         t = time.time() - t
         speed_msg = ""
@@ -948,22 +959,6 @@ def write_flash(
                 "Cannot verify written data if encrypted or in secure download mode."
             )
 
-    if esp.IS_STUB:
-        # skip sending flash_finish to ROM loader here,
-        # as it causes the loader to exit and run user code
-        esp.flash_begin(0, 0)
-
-        # Get the "encrypted" flag for the last file flashed
-        # Note: all_files list contains quadruplets like:
-        # (address: int, filename: str | None, data: bytes, encrypted: bool)
-        last_file_encrypted = all_files[-1][3]
-
-        # Check whether the last file flashed was compressed or not
-        if compress and not last_file_encrypted:
-            esp.flash_defl_finish(False)
-        else:
-            esp.flash_finish(False)
-
 
 def read_mac(esp: ESPLoader) -> None:
     """

esptool/loader.py

@@ -1066,11 +1066,16 @@ def flash_encrypt_block(self, data, seq, timeout=DEFAULT_TIMEOUT):
                 else:
                     raise
 
-    def flash_finish(self, reboot=False):
+    def flash_finish(self, reboot=False, timeout=DEFAULT_TIMEOUT):
         """Leave flash mode and run/reboot"""
         pkt = struct.pack("<I", int(not reboot))
         # stub sends a reply to this command
-        self.check_command("leave flash download mode", self.ESP_CMDS["FLASH_END"], pkt)
+        self.check_command(
+            "leave flash download mode",
+            self.ESP_CMDS["FLASH_END"],
+            pkt,
+            timeout=timeout,
+        )
 
     def run(self, reboot=False):
         """Run application code in flash"""
@@ -1400,15 +1405,18 @@ def flash_defl_block(self, data, seq, timeout=DEFAULT_TIMEOUT):
                     raise
 
     @stub_and_esp32_function_only
-    def flash_defl_finish(self, reboot=False):
+    def flash_defl_finish(self, reboot=False, timeout=DEFAULT_TIMEOUT):
         """Leave compressed flash mode and run/reboot"""
         if not reboot and not self.IS_STUB:
             # skip sending flash_finish to ROM loader, as this
             # exits the bootloader. Stub doesn't do this.
             return
         pkt = struct.pack("<I", int(not reboot))
         self.check_command(
-            "leave compressed flash mode", self.ESP_CMDS["FLASH_DEFL_END"], pkt
+            "leave compressed flash mode",
+            self.ESP_CMDS["FLASH_DEFL_END"],
+            pkt,
+            timeout=timeout,
         )
         self.in_bootloader = False