Command-line tool to write a disk image on a USB device with macOS.

flashimg.py

#!/usr/bin/env python3

"""
Command-line tool to write a disk image on a USB device with macOS.

Exactly like [etcher](https://etcher.balena.io), [rpi-imager](https://github.com/raspberrypi/rpi-imager),
[rufus.ie](https://rufus.ie), [UNetbootin](https://unetbootin.github.io), etc.
No frills, no freaks, much simple and so more efficient.

Unix has the right tool for over 50 years: [`dd`](https://en.wikipedia.org/wiki/Dd_(Unix)). Sorry folks and

Note: root access is mandatory to run the `dd` command.

Requirements: `pip3 install defusedxml pick humanfriendly`
"""

import argparse
import json
import logging
import subprocess
import typing as t
import xml.etree.ElementTree
from pathlib import Path

import defusedxml.ElementTree as ET
import humanfriendly as hf
from pick import pick


def plist_parse_tag(node: xml.etree.ElementTree.Element) -> dict[str, t.Any]:
    """Convert a property value node to its Python equivalent."""

    if node.tag == "true":
        return True

    elif node.tag == "false":
        return False

    elif node.tag == "string":
        return node.text

    elif node.tag == "integer":
        return int(node.text)

    elif node.tag == "array":
        return [plist_parse_tag(child) for child in node]

    elif node.tag == "dict":
        result = {}
        i = iter(node)
        while True:
            try:
                key = next(i).text
                value = next(i)

                result[key] = plist_parse_tag(value)

            except StopIteration:
                break

        return result

    else:
        raise ValueError(f"Unknown tag '{node.tag}'")


def from_plist(xml_plist) -> dict:
    """Convert a macOS plist in XML format to a Python `dict`."""
    doc = ET.fromstring(xml_plist)
    assert doc.tag == "plist"
    return plist_parse_tag(doc[0])


def main():

    try:
        parser = argparse.ArgumentParser()
        parser.add_argument("-v", "--verbose", action="store_true", help="Show the device details")
        parser.add_argument("-k", "--keep", action="store_true", help="Do not eject device")
        parser.add_argument("imagefile", type=Path)
        args = parser.parse_args()

        logging.basicConfig(format="\033[32m%(asctime)s\033[0m - %(levelname)s - %(message)s", level=logging.DEBUG)

        if not args.imagefile.is_file():
            parser.error("file must exist")

        # get the list of external disks
        diskutil_list = subprocess.check_output(["diskutil", "list", "-plist", "external", "physical"])
        diskutil_list = from_plist(diskutil_list)

        externals = []
        for identifier in diskutil_list["WholeDisks"]:

            # retrieve the details of the disk
            info = subprocess.check_output(["diskutil", "info", "-plist", identifier])
            info = from_plist(info)

            if info["Writable"] and info["Ejectable"] and info["BusProtocol"] == "USB":
                for disk in diskutil_list["AllDisksAndPartitions"]:
                    if disk["DeviceIdentifier"] == identifier:
                        externals.append((disk, info))
                        break

        if len(externals) == 0:
            logging.error("No eligible device available")
            exit()

        if len(externals) > 1:
            title = "Please choose a disk device:"
            options = []
            for disk, info in externals:
                s = ",".join(
                    p.get("MountPoint") or p.get("VolumeName") or p["DeviceIdentifier"]
                    for p in disk.get("Partitions", ())
                )
                options.append(f"{disk['DeviceIdentifier']} : {hf.format_size(info['TotalSize'])} - {s}")
            option, index = pick(options, title, indicator=">", clear_screen=False, quit_keys=[27, ord("q")])
            if option is None:
                exit()

            disk, info = externals[index]
        else:
            disk, info = externals[0]

        if args.verbose:
            print(json.dumps(disk, indent=2))
            print(json.dumps(info, indent=2))

        logging.info(
            f"Disk {disk['DeviceIdentifier']} size: {hf.format_size(disk['Size'])} name: {info.get('MediaName')}"
        )
        for p in disk.get("Partitions", ()):
            logging.info(
                f"Partition {p['DeviceIdentifier']} - {p.get('Content')}  {p.get('VolumeName')}  {p.get('MountPoint')}"
            )

        identifier = disk["DeviceIdentifier"]

        raw_device = f"/dev/r{identifier}"  # dd requires /dev/rdiskX device

        response = input(
            f"About to flash \033[32m{args.imagefile}\033[0m to \033[95m{raw_device}\033[0m."
            " Enter \033[97mC\033[0m to continue, anything else to cancel: "
        )
        if response.casefold() != "c":
            logging.info("Cancelled")
            exit()

        # 1. ensure the disk is unmounted
        logging.info(f"Unmounting {identifier}")
        subprocess.check_call(["diskutil", "unmountDisk", identifier])

        # 2. copy the raw image to the media
        logging.info(f"Copying image {args.imagefile} to {raw_device} ...")
        subprocess.check_call(["sudo", "dd", f"if={args.imagefile}", f"of={raw_device}", "status=progress", "bs=16m"])

        # 3. eject (i.e. prevent macOS to mount any partition)
        if args.keep:
            logging.info(f"Device {identifier} not ejected !")
        else:
            logging.info(f"Eject {identifier}")
            subprocess.check_call(["diskutil", "eject", identifier])

        logging.info("Done")

    except subprocess.CalledProcessError as e:
        print()
        logging.info(e)

    except KeyboardInterrupt:
        print()
        logging.info("Interrupted")


if __name__ == "__main__":
    main()