Building a canned filesystem dfu for micropython

fsbuilder.py

"""
MicroPython filesystem builder.

Example usage (build micropython-dev using make VARIANT=dev in the unix port):

    micropython-dev -X heapsize=10m fsbuilder.py 4096 512 lfs2 directory/

This will create an image "directory.img" which can then be deployed to a device.
To turn it into a DFU file use (dfu.py can be found in MicroPython tools/):

    dfu.py -b 0x80000000:directory.img directory.dfu

This can be deployed independently to the main firmware.  To bundle the firmware
and filesystem image into one DFU use, for PYBD-SF2:

    dfu.py -b 0x08008000:firmware0.bin 0x90000000:firmware1.bin -b 0x80000000:directory.img full.dfu

And for PYBD-SF6:

    dfu.py -b 0x08008000:firmware.bin -b 0x80000000:directory.img full.dfu

The firmware[01].bin files can be found in the board's build directory.
"""

import sys, ubinascii, uos
import ustruct as struct


class RAMBlockDevice:
    def __init__(self, block_size, num_blocks):
        self.block_size = block_size
        self.data = bytearray(block_size * num_blocks)
        for i in range(len(self.data)):
            self.data[i] = 0xFF

    def readblocks(self, block_num, buf, offset=0):
        addr = block_num * self.block_size + offset
        for i in range(len(buf)):
            buf[i] = self.data[addr + i]

    def writeblocks(self, block_num, buf, offset=None):
        if offset is None:
            # do erase, then write
            for i in range(len(buf) // self.block_size):
                self.ioctl(6, block_num + i)
            offset = 0
        addr = block_num * self.block_size + offset
        for i in range(len(buf)):
            self.data[addr + i] = buf[i]

    def ioctl(self, op, arg):
        if op == 4:  # block count
            return len(self.data) // self.block_size
        if op == 5:  # block size
            return self.block_size
        if op == 6:  # block erase
            return 0


def copy_recursively(vfs, src_dir, dest_dir):
    assert src_dir.endswith("/")
    assert dest_dir.endswith("/")
    DIR = 1 << 14
    for name, mode, _ in uos.ilistdir(src_dir):
        src_name = src_dir + name
        dest_name = dest_dir + name
        if mode & DIR:
            print(dest_name + "/")
            vfs.mkdir(dest_name)
            copy_recursively(vfs, src_name + "/", dest_name + "/")
        else:
            print(dest_name)
            with open(src_name, "rb") as src, vfs.open(dest_name, "wb") as dest:
                dest.write(src.read())


def build_dfu(elem_addr, elem_data):
    # Build single element.
    out_elem = struct.pack("<II", elem_addr, len(elem_data)) + elem_data

    # Build single target.
    sig = b"Target"
    alt = 0
    has_name = 1
    name = b"VFS"
    t_size = len(out_elem)
    num_elem = 1
    out_targ = struct.pack("<6sBi255sII", sig, alt, has_name, name, t_size, num_elem)
    out_targ += out_elem

    # Build DFU.
    sig = b"DfuSe"
    ver = 1
    size = len(out_targ) + 11
    num_targ = 1
    out = struct.pack("<5sBIB", sig, ver, size, num_targ) + out_targ

    # Add footer.
    dev = 0
    vid = 0x0483
    pid = 0xDF11
    out += struct.pack("<HHHH3sB", dev, pid, vid, 0x011A, b"UFD", 16)

    # Add CRC.
    crc = ~ubinascii.crc32(out)
    out += struct.pack("<I", crc)

    return out


def main():
    if len(sys.argv) != 5:
        print(
            "usage: {} <block_size> <num_blocks> <vfs_type> <dir>".format(sys.argv[0]),
            file=sys.stderr,
        )
        sys.exit(1)

    # Parse arguments.
    block_size = int(sys.argv[1])
    num_blocks = int(sys.argv[2])
    vfs_type = sys.argv[3]
    dir = sys.argv[4]
    if not dir.endswith("/"):
        dir += "/"

    # Get VFS class
    vfs_class_name = "Vfs{}{}".format(vfs_type[0].upper(), vfs_type[1:])
    try:
        vfs_class = getattr(uos, vfs_class_name)
    except AttributeError:
        print("unsupported VFS type: {}".format(vfs_type), file=sys.stderr)
        sys.exit(1)

    # Create the RAM device.
    ramdev = RAMBlockDevice(block_size, num_blocks)

    # Format the filesystem and create the VFS.
    vfs_class.mkfs(ramdev)
    vfs = vfs_class(ramdev)

    # Build the filesystem recursively.
    print(
        "Building filesystem: block_size={} num_blocks={} type={}".format(
            block_size, num_blocks, vfs_class_name
        )
    )
    print("Source directory: {}".format(dir))
    try:
        copy_recursively(vfs, dir, "/")
    except OSError as er:
        if er.args[0] == 28:  # ENOSPC
            print("Error: not enough space on filesystem", file=sys.stderr)
            sys.exit(1)
        else:
            print("Error: OSError {}".format(er.args[0]), file=sys.stderr)
            sys.exit(1)

    # Save the block device data.
    output = dir.rstrip("/") + ".img"
    print("Writing block device to {}".format(output))
    with open(output, "wb") as f:
        f.write(ramdev.data)

    # Save the block device data as a DFU file.
    if 0:
        output = dir.rstrip("/") + ".dfu"
        print("Writing block device to {}".format(output))
        with open(output, "wb") as f:
            f.write(build_dfu(0x80000000, ramdev.data))


if __name__ == "__main__":
    main()