tools/mpremote: Allow .img for ROMFS file and validate ROMFS image.

Currently the tool allows writing an invalid ROMFS image, with a bad header
or images smaller than minimum size, and only checks the image extension.

This commit allows deploying a ROMFS with either a ".img" or ".romfs"
extension (in the future support may be added for other extensions that
have different semantics, eg a manifest), and validates the image header
before writing.

Signed-off-by: iabdalkader <i.abdalkader@gmail.com>

1 files changed, 7 insertions, 2 deletions

diff --git a/tools/mpremote/mpremote/commands.py b/tools/mpremote/mpremote/commands.py
index 7dd448c8b..e385d0509 100644
--- a/tools/mpremote/mpremote/commands.py
+++ b/tools/mpremote/mpremote/commands.py
@@ -9,7 +9,7 @@ import serial.tools.list_ports
 
 from .transport import TransportError, TransportExecError, stdout_write_bytes
 from .transport_serial import SerialTransport
-from .romfs import make_romfs
+from .romfs import make_romfs, VfsRomWriter
 
 
 class CommandError(Exception):
@@ -555,7 +555,7 @@ def _do_romfs_deploy(state, args):
     romfs_filename = args.path
 
     # Read in or create the ROMFS filesystem image.
-    if romfs_filename.endswith(".romfs"):
+    if os.path.isfile(romfs_filename) and romfs_filename.endswith((".img", ".romfs")):
         with open(romfs_filename, "rb") as f:
             romfs = f.read()
     else:
@@ -581,6 +581,11 @@ def _do_romfs_deploy(state, args):
         rom_size = transport.eval("len(dev)")
         print(f"ROMFS{rom_id} partition has size {rom_size} bytes")
 
+    # Check if ROMFS image is valid
+    if not romfs.startswith(VfsRomWriter.ROMFS_HEADER):
+        print("Invalid ROMFS image")
+        sys.exit(1)
+
     # Check if ROMFS filesystem image will fit in the target partition.
     if len(romfs) > rom_size:
         print("ROMFS image is too big for the target partition")