deviceTree.applyOverlays: rewrite in Python/libfdt

from argparse import ArgumentParser

from dataclasses import dataclass

from functools import cached_property

import json

from pathlib import Path

from libfdt import Fdt, FdtException, FDT_ERR_NOSPACE, fdt_overlay_apply

@dataclass

class Overlay:

    name: str

    filter: str

    dtbo_file: Path

    @cached_property

    def fdt(self):

        with self.dtbo_file.open("rb") as fd:

            return Fdt(fd.read())

    @cached_property

    def compatible(self):

        return get_compatible(self.fdt)

def get_compatible(fdt):

    root_offset = fdt.path_offset("/")

    return set(fdt.getprop(root_offset, "compatible").as_stringlist())

def apply_overlay(dt: Fdt, dto: Fdt) -> Fdt:

    while True:

        # we need to copy the buffers because they can be left in an inconsistent state

        # if the operation fails (ref: fdtoverlay source)

        result = dt.as_bytearray().copy()

        err = fdt_overlay_apply(result, dto.as_bytearray().copy())

        if err == 0:

            new_dt = Fdt(result)

            # trim the extra space from the final tree

            new_dt.pack()

            return new_dt

        if err == -FDT_ERR_NOSPACE:

            # not enough space, add some blank space and try again

            # magic number of more space taken from fdtoverlay

            dt.resize(dt.totalsize() + 65536)

            continue

        raise FdtException(err)

def main():

    parser = ArgumentParser(description='Apply a list of overlays to a directory of device trees')

    parser.add_argument("--source", type=Path, help="Source directory")

    parser.add_argument("--destination", type=Path, help="Destination directory")

    parser.add_argument("--overlays", type=Path, help="JSON file with overlay descriptions")

    args = parser.parse_args()

    source: Path = args.source

    destination: Path = args.destination

    overlays: Path = args.overlays

    with overlays.open() as fd:

        overlays_data = [

            Overlay(

                name=item["name"],

                filter=item["filter"],

                dtbo_file=Path(item["dtboFile"]),

            )

            for item in json.load(fd)

        ]

    for source_dt in source.glob("**/*.dtb"):

        rel_path = source_dt.relative_to(source)

        print(f"Processing source device tree {rel_path}...")

        with source_dt.open("rb") as fd:

            dt = Fdt(fd.read())

        dt_compatible = get_compatible(dt)

        for overlay in overlays_data:

            if overlay.filter and overlay.filter not in str(rel_path):

                print(f"  Skipping overlay {overlay.name}: filter does not match")

                continue

            if not overlay.compatible.intersection(dt_compatible):

                print(f"  Skipping overlay {overlay.name}: {overlay.compatible} is incompatible with {dt_compatible}")

                continue

            print(f"  Applying overlay {overlay.name}")

            dt = apply_overlay(dt, overlay.fdt)

        print(f"Saving final device tree {rel_path}...")

        dest_path = destination / rel_path

        dest_path.parent.mkdir(parents=True, exist_ok=True)

        with dest_path.open("wb") as fd:

            fd.write(dt.as_bytearray())

if __name__ == '__main__':

    main()

{ lib, stdenv, stdenvNoCC, dtc }:

{ lib, stdenv, stdenvNoCC, dtc, writers, python3 }:

{

  # Compile single Device Tree overlay source

  applyOverlays = (base: overlays': stdenvNoCC.mkDerivation {

    name = "device-tree-overlays";

    nativeBuildInputs = [ dtc ];

    buildCommand = let

      overlays = lib.toList overlays';

    in ''

      mkdir -p $out

      cd "${base}"

      find -L . -type f -name '*.dtb' -print0 \

        | xargs -0 cp -v --no-preserve=mode --target-directory "$out" --parents

      for dtb in $(find "$out" -type f -name '*.dtb'); do

        dtbCompat=$(fdtget -t s "$dtb" / compatible 2>/dev/null || true)

        # skip files without `compatible` string

        test -z "$dtbCompat" && continue

        ${lib.flip (lib.concatMapStringsSep "\n") overlays (o: ''

        overlayCompat="$(fdtget -t s "${o.dtboFile}" / compatible)"

        # skip incompatible and non-matching overlays

        if [[ ! "$dtbCompat" =~ "$overlayCompat" ]]; then

          echo "Skipping overlay ${o.name}: incompatible with $(basename "$dtb")"

        elif ${if (o.filter == null) then "false" else ''

          [[ "''${dtb//${o.filter}/}" ==  "$dtb" ]]

        ''}

        then

          echo "Skipping overlay ${o.name}: filter does not match $(basename "$dtb")"

        else

          echo -n "Applying overlay ${o.name} to $(basename "$dtb")... "

          mv "$dtb"{,.in}

          fdtoverlay -o "$dtb" -i "$dtb.in" "${o.dtboFile}"

          echo "ok"

          rm "$dtb.in"

        fi

        '')}

      done

    nativeBuildInputs = [

      (python3.pythonOnBuildForHost.withPackages(ps: [ps.libfdt]))

    ];

    buildCommand = ''

      python ${./apply_overlays.py} --source ${base} --destination $out --overlays ${writers.writeJSON "overlays.json" overlays'}

    '';

  });

}