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

@@ -708,10 +708,18 @@ def write_flash(esp, args):
                     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][2]
+
             # 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 args.compress and not last_file_encrypted:
+                esp.flash_defl_finish(reboot=False, timeout=timeout)
+            else:
+                esp.flash_finish(reboot=False, timeout=timeout)
 
         t = time.time() - t
         speed_msg = ""
@@ -750,22 +758,6 @@ def write_flash(esp, args):
 
     print("\nLeaving...")
 
-    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 triplets like:
-        # (address: Integer, filename: String, encrypted: Boolean)
-        last_file_encrypted = all_files[-1][2]
-
-        # Check whether the last file flashed was compressed or not
-        if args.compress and not last_file_encrypted:
-            esp.flash_defl_finish(False)
-        else:
-            esp.flash_finish(False)
-
 
 def _parse_app_info(app_info_segment):
     """

esptool/loader.py

@@ -1040,11 +1040,11 @@ 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 mode", self.ESP_FLASH_END, pkt)
+        self.check_command("leave Flash mode", self.ESP_FLASH_END, pkt, timeout=timeout)
 
     def run(self, reboot=False):
         """Run application code in flash"""
@@ -1329,14 +1329,16 @@ 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_FLASH_DEFL_END, pkt)
+        self.check_command(
+            "leave compressed flash mode", self.ESP_FLASH_DEFL_END, pkt, timeout=timeout
+        )
         self.in_bootloader = False
 
     @stub_and_esp32_function_only