Harden write_bootable_mmc with sudo check and safer device selection

write_bootable_mmc.py

@@ -36,6 +36,12 @@ import subprocess
 import sys
 import os
 import time
+import json
+import shlex
+import stat
+
+MIN_DEVICE_BYTES = 7 * 1024 * 1024 * 1024
+MAX_DEVICE_BYTES = 17 * 1024 * 1024 * 1024
 
 # functions
 # useful link 1: http://superuser.com/questions/868117/layouting-a-disk-image-and-copying-files-into-it
@@ -49,7 +55,7 @@ def shout(cmd):
 def print_help():
     """Print help information"""
     print("\nDescription:\n")
-    print("  * Required programs: kpartx, parted")
+    print("  * Required programs: lsblk, kpartx, parted")
     print("  * Run under superuser. Make sure the correct device is provided.")
     print("  * Erases partition table on the provided device")
     print("  * If given someimage.img file - burns the sd card from it")
@@ -58,12 +64,148 @@ def print_help():
     print("  * Creates EXT4 partition labeled 'root' and extracts rootfs.tar.gz")
     print("\nExamples:\n")
     print("  * Use files (names are hardcoded) from the current dir ('build/tmp/deploy/images/elphel393/mmc/'):")
-    print("      ~$ sudo python3 make_sdcard.py /dev/sdz")
+    print("      ~$ sudo write_bootable_mmc.py /dev/sdz")
     print("  * Use someimage.img file:")
-    print("      ~$ sudo python3 make_sdcard.py /dev/sdz someimage.img")
+    print("      ~$ sudo write_bootable_mmc.py /dev/sdz someimage.img")
+    print("  * Auto-detect likely removable 8/16GB devices and pick interactively:")
+    print("      ~$ sudo write_bootable_mmc.py")
     print("  * To write *.iso use a standard OS tool that burns bootable USB drives")
     print("")
 
+def list_block_devices():
+    """Return top-level block devices from lsblk."""
+    out = subprocess.check_output(
+        ["lsblk", "-b", "-J", "-o", "NAME,PATH,TYPE,SIZE,RM,TRAN,MODEL,SERIAL,MOUNTPOINTS"],
+        text=True
+    )
+    data = json.loads(out)
+    return data.get("blockdevices", [])
+
+def has_mounted_children(dev):
+    """Return True if any child partition is mounted."""
+    for child in dev.get("children", []) or []:
+        mountpoints = child.get("mountpoints")
+        if isinstance(mountpoints, list):
+            if any(mp for mp in mountpoints if mp):
+                return True
+        elif mountpoints:
+            return True
+        if has_mounted_children(child):
+            return True
+    return False
+
+def as_disk_entry(dev):
+    """Normalize lsblk disk entry."""
+    path = dev.get("path")
+    if not path:
+        name = dev.get("name", "")
+        path = f"/dev/{name}" if name else ""
+    size = int(dev.get("size") or 0)
+    model = (dev.get("model") or "").strip()
+    serial = (dev.get("serial") or "").strip()
+    tran = (dev.get("tran") or "").strip().lower()
+    rm = int(dev.get("rm") or 0)
+    return {
+        "path": path,
+        "size": size,
+        "model": model,
+        "serial": serial,
+        "tran": tran,
+        "rm": rm,
+        "mounted": has_mounted_children(dev),
+    }
+
+def get_disks():
+    """Return normalized disk list."""
+    disks = []
+    for dev in list_block_devices():
+        if dev.get("type") == "disk":
+            disks.append(as_disk_entry(dev))
+    return disks
+
+def format_gib(size_bytes):
+    """Format size in GiB."""
+    return f"{size_bytes / (1024.0 ** 3):.1f} GiB"
+
+def choose_device_interactive(disks):
+    """Choose target device from likely removable 8/16GB disks."""
+    candidates = []
+    for d in disks:
+        if d["size"] < MIN_DEVICE_BYTES or d["size"] > MAX_DEVICE_BYTES:
+            continue
+        if d["mounted"]:
+            continue
+        if d["rm"] == 1 or d["tran"] in ("usb", "mmc", "sdio"):
+            candidates.append(d)
+
+    if not candidates:
+        print("ERROR: no safe removable 8/16GB candidate device found.")
+        print("Disks discovered:")
+        for d in disks:
+            mounted = "mounted" if d["mounted"] else "not-mounted"
+            print(f"  {d['path']:<14} {format_gib(d['size']):>8} rm={d['rm']} tran={d['tran'] or '-'} {mounted}")
+        print("Provide target device explicitly, for example:")
+        print(f"  sudo {os.path.basename(sys.argv[0])} /dev/sdX")
+        sys.exit(1)
+
+    print("Select target block device:")
+    for i, d in enumerate(candidates, start=1):
+        label = " ".join(x for x in (d["model"], d["serial"]) if x).strip()
+        print(
+            f"  [{i}] {d['path']:<14} {format_gib(d['size']):>8} "
+            f"rm={d['rm']} tran={d['tran'] or '-'} {label}"
+        )
+
+    while True:
+        val = input(f"Enter number [1-{len(candidates)}] or 'q' to quit: ").strip().lower()
+        if val in ("q", "quit", "exit"):
+            sys.exit(1)
+        if val.isdigit():
+            idx = int(val)
+            if 1 <= idx <= len(candidates):
+                return candidates[idx - 1]["path"]
+        print("Invalid selection.")
+
+def get_disk_meta(disks, device):
+    """Return metadata for a selected disk path."""
+    for d in disks:
+        if d["path"] == device:
+            return d
+    return None
+
+def ensure_running_as_root():
+    """Require sudo/root execution."""
+    if os.geteuid() != 0:
+        cmd = " ".join(shlex.quote(a) for a in sys.argv)
+        print("ERROR: this program must be launched with sudo.")
+        print(f"Try:\n  sudo {cmd}")
+        sys.exit(1)
+
+def ensure_block_device(device):
+    """Validate that provided path is a block device."""
+    if not os.path.exists(device):
+        print(f"No such device: {device}")
+        sys.exit(1)
+    mode = os.stat(device).st_mode
+    if not stat.S_ISBLK(mode):
+        print(f"Not a block device: {device}")
+        sys.exit(1)
+
+def confirm_erase(device, meta):
+    """Ask for explicit confirmation before erasing."""
+    if meta:
+        size = format_gib(meta["size"])
+        tran = meta["tran"] or "-"
+        model = " ".join(x for x in (meta["model"], meta["serial"]) if x).strip()
+        print(f"Selected: {device} ({size}, rm={meta['rm']}, tran={tran}) {model}".rstrip())
+        if meta["size"] < MIN_DEVICE_BYTES or meta["size"] > MAX_DEVICE_BYTES:
+            print("WARNING: selected device size is outside expected 8/16GB range.")
+    print(f"WARNING: this will erase all data on {device}.")
+    answer = input("Type YES to continue: ").strip()
+    if answer != "YES":
+        print("Aborted.")
+        sys.exit(1)
+
 def check_program_installed(program):
     """Check if a program is installed"""
     try:
@@ -75,6 +217,7 @@ def check_program_installed(program):
 
 # Check required programs
 required_programs = (
+    "lsblk",
     "parted",
     "kpartx"
 )
@@ -89,21 +232,41 @@ if something_is_missing:
     sys.exit(1)
 
 # Parse command line arguments
-if len(sys.argv) > 1:
-    DEVICE = sys.argv[1]
-else:
-    DEVICE = ""
+args = []
+for arg in sys.argv[1:]:
+    if arg in ("-h", "--help"):
         print_help()
         sys.exit(0)
+    args.append(arg)
+
+if len(args) > 2:
+    print("ERROR: wrong number of arguments.")
+    print_help()
+    sys.exit(1)
+
+ensure_running_as_root()
 
-if len(sys.argv) > 2:
-    IMAGE_FILE = sys.argv[2]
+disks = get_disks()
+if len(args) > 0:
+    DEVICE = args[0]
+else:
+    if not sys.stdin.isatty():
+        print("ERROR: no target device provided and no interactive terminal to choose one.")
+        print(f"Use: sudo {os.path.basename(sys.argv[0])} /dev/sdX [image.img]")
+        sys.exit(1)
+    DEVICE = choose_device_interactive(disks)
+
+if len(args) > 1:
+    IMAGE_FILE = args[1]
     if not IMAGE_FILE.endswith(".img"):
         print("ERROR: Please provide *.img file or leave argument empty to use certain image files in the current dir")
         sys.exit(1)
 else:
     IMAGE_FILE = ""
 
+ensure_block_device(DEVICE)
+confirm_erase(DEVICE, get_disk_meta(disks, DEVICE))
+
 print("NOTE: If plasma crashes, do not worry")
 
 # Parameters
@@ -140,10 +303,6 @@ else:
         print("No such file")
         something_is_missing = True
 
-if not os.path.exists(DEVICE):
-    print(f"No such device: {DEVICE}")
-    something_is_missing = True
-
 if something_is_missing:
     sys.exit(1)
 
@@ -167,6 +326,9 @@ shout(f"parted -s {DEVICE} align-check optimal 2")
 devs_created = False
 partition1 = f"{DEVICE}1" if not "mmcblk" in DEVICE else f"{DEVICE}p1"
 partition2 = f"{DEVICE}2" if not "mmcblk" in DEVICE else f"{DEVICE}p2"
+if os.path.basename(DEVICE)[-1].isdigit():
+    partition1 = f"{DEVICE}p1"
+    partition2 = f"{DEVICE}p2"
 
 print("= Waiting for device nodes...")
 while not devs_created: