feat: Add support for k, M suffix for flash size

Author: peterdragun

docs/en/esptool/advanced-commands.rst

@@ -41,7 +41,8 @@ The ``dump-mem`` command will dump a region from the chip's memory space to a fi
 
 ::
 
-    esptool.py dump-mem 0x40000000 65536 iram0.bin
+    esptool.py dump-mem 0x40000000 64k iram0.bin
+
 
 .. _load-ram:
 

docs/en/esptool/basic-commands.rst

@@ -108,6 +108,8 @@ The read-flash command allows reading back the contents of flash. The arguments
     esptool.py -p PORT -b 460800 read-flash 0 0x200000 flash_contents.bin
 
 
+Size can be specified in bytes, or with suffixes like ``k`` and ``M``. So ``0x200000`` in example can be replaced with ``2M``.
+
 It is also possible to autodetect flash size by using ``ALL`` as size. The above example with autodetection would look like this:
 
 ::
@@ -135,11 +137,11 @@ To erase the entire flash chip (all data replaced with 0xFF bytes):
 
     esptool.py erase-flash
 
-To erase a region of the flash, starting at address 0x20000 with length 0x4000 bytes (16KB):
+To erase a region of the flash, starting at address 0x20000 with length 16 kB (0x4000 bytes):
 
 ::
 
-    esptool.py erase-region 0x20000 0x4000
+    esptool.py erase-region 0x20000 16k
 
 The address and length must both be multiples of the SPI flash erase sector size. This is 0x1000 (4096) bytes for supported flash chips.
 

esptool/__init__.py

@@ -44,6 +44,7 @@
 from esptool.cmds import (
     chip_id,
     detect_chip,
+    detect_flash_size,
     dump_mem,
     elf2image,
     erase_flash,
@@ -81,6 +82,7 @@
 from esptool.util import (
     FatalError,
     NotImplementedInROMError,
+    flash_size_bytes,
 )
 from itertools import chain, cycle, repeat
 
@@ -271,6 +273,18 @@ def check_flash_size(esp: ESPLoader, address: int, size: int) -> None:
             f"(set size {size:#x} from address {address:#010x} goes past 16MB "
             f"by {address + size - 0x1000000:#x} bytes)."
         )
+    # Check if we are writing/reading past detected flash size
+    if not esp.secure_download_mode:
+        detected_size_str = detect_flash_size(esp)
+        if not detected_size_str:
+            return
+        detected_size = flash_size_bytes(detected_size_str)
+        if address + size > detected_size:
+            raise FatalError(
+                f"Can't access flash regions larger than detected flash size "
+                f"(set size {size:#x} from address {address:#010x} goes past "
+                f"{detected_size_str} by {address + size - detected_size:#x} bytes)."
+            )
 
 
 ############################### GLOBAL OPTIONS AND MAIN ###############################
@@ -544,7 +558,7 @@ def load_ram_cli(ctx, filename):
 
 @cli.command("dump-mem")
 @click.argument("address", type=AnyIntType())
-@click.argument("size", type=AnyIntType())
+@click.argument("size", type=AutoSizeType(allow_all=False))
 @click.argument("output", type=click.Path())
 @click.pass_context
 def dump_mem_cli(ctx, address, size, output):

esptool/cli_util.py

@@ -58,13 +58,29 @@ def convert(
 
 
 class AutoSizeType(AnyIntType):
-    """Similar to AnyIntType but allows 'all' as a value to e.g. read whole flash"""
+    """Similar to AnyIntType but allows 'k', 'M' suffixes for kilo(1024), Mega(1024^2)
+    and 'all' as a value to e.g. read whole flash"""
+
+    def __init__(self, allow_all: bool = True):
+        self.allow_all = allow_all
+        super().__init__()
 
     def convert(
         self, value: str, param: click.Parameter | None, ctx: click.Context
     ) -> Any:
-        if value.lower() == "all":
+        if self.allow_all and value.lower() == "all":
             return value
+        # Handle suffixes like 'k', 'M' for kilo, mega
+        if value[-1] in ("k", "M"):
+            try:
+                num = arg_auto_int(value[:-1])
+            except ValueError:
+                raise click.BadParameter(f"{value!r} is not a valid integer")
+            if value[-1] == "k":
+                num *= 1024
+            elif value[-1] == "M":
+                num *= 1024 * 1024
+            return num
         return super().convert(value, param, ctx)
 
 
@@ -388,6 +404,13 @@ def parse_port_filters(
 def parse_size_arg(esp: ESPLoader, size: int | str) -> int:
     """Parse the flash size argument and return the size in bytes"""
     if isinstance(size, int):
+        if not esp.secure_download_mode:
+            detected_size = flash_size_bytes(detect_flash_size(esp))
+            if detected_size and size > detected_size:
+                raise FatalError(
+                    f"Specified size {size:#x} is greater than detected flash size "
+                    f"{detected_size:#x}.",
+                )
         return size
     if size.lower() != "all":
         raise FatalError(f"Invalid size value: {size}. Use an integer or 'all'.")

test/test_esptool.py

@@ -439,6 +439,7 @@ def test_short_flash(self):
         self.run_esptool("write-flash 0x0 images/one_kb.bin")
         self.verify_readback(0, 1024, "images/one_kb.bin")
 
+    @pytest.mark.skipif(arg_chip != "esp32", reason="Don't need to test multiple times")
     def test_short_flash_deprecated(self):
         out = self.run_esptool(
             "--before default_reset write_flash 0x0 images/one_kb.bin --flash_size keep"
@@ -877,6 +878,13 @@ def test_write_past_end_fails(self):
         assert "File 'images/one_kb.bin'" in output
         assert "will not fit" in output
 
+    @pytest.mark.skipif(arg_chip != "esp32", reason="Don't need to test multiple times")
+    def test_read_past_end_fails(self):
+        output = self.run_esptool_error(
+            "read-flash 0xffffff 1 out.bin"
+        )  # 0xffffff is well past the end of the flash in most cases (16MB)
+        assert "Can't access flash regions larger than detected flash size" in output
+
     def test_write_no_compression_past_end_fails(self):
         output = self.run_esptool_error(
             "write-flash -u -fs 1MB 0x280000 images/one_kb.bin"
@@ -913,7 +921,7 @@ def test_read_nostub_high_offset(self):
         # readback with no-stub and flash-size set
         try:
             self.run_esptool(
-                f"--no-stub read-flash -fs detect {offset} 1024 {dump_file.name}"
+                f"--no-stub read-flash -fs detect {offset} 1k {dump_file.name}"
             )
             with open(dump_file.name, "rb") as f:
                 rb = f.read()