nerves_system_openwrt_one

defmodule Mix.Tasks.OpenwrtOne.FirmwarePost do
  @shortdoc "Splice user-merged .itb + .ubi into the firmware archive"

  @moduledoc """
  Post-process a `mix firmware` .fw so it ships the correct
  user-merged initramfs/UBI rather than the system-only versions
  baked in at system build time.

  Designed to be chained after `firmware` via a mix alias:

      defp aliases do
        [
          firmware: ["firmware", "openwrt_one.firmware_post"]
        ]
      end

  With the alias in place:

    * `mix firmware` produces a `.fw` whose `data/openwrt-one-initramfs.itb`
      and `data/openwrt-one-nand.ubi` are derived from the .fw's own
      combined squashfs — i.e. they contain `/srv/erlang`.
    * `mix burn` works directly: fwup writes the spliced UBI to the
      recovery USB stick, so the first boot after the USB dance finds
      the Erlang release.
    * `mix upload` can extract `data/openwrt-one-initramfs.itb` straight
      out of the `.fw` (no on-the-fly rebuild needed).

  Without the alias, `mix firmware` ships the stale system-only
  resources and the device boots with `erlinit: No release found in
  /srv/erlang.`

  ## What it does

    1. Calls `scripts/wrap-firmware.sh` from the system, which extracts
       the .fw's squashfs, unsquashes it, repacks as cpio.gz, and runs
       `mkimage` + `ubinize` to produce a fresh `.itb` and `.ubi`.
    2. Replaces `data/openwrt-one-initramfs.itb` and
       `data/openwrt-one-nand.ubi` inside the .fw zip with those fresh
       artifacts.
    3. Rewrites the matching `length` and `blake2b-256` entries in
       `meta.conf` so fwup's integrity check still passes.
    4. Preserves the original zip entry ordering (fwup expects
       resources to appear in the order their `on-resource` clauses
       are declared).

  ## Requirements

    * `b2sum` (GNU coreutils ≥ 8.26) — `apt install coreutils` on
      Linux, `brew install coreutils` on macOS (binary is `gb2sum`).
    * Everything `scripts/wrap-firmware.sh` needs: `unzip`,
      `unsquashfs`, `cpio`, `gzip`, `mkimage`, `dtc`, `ubinize`.
  """
  use Mix.Task

  @itb_resource "openwrt-one-initramfs.itb"
  @itb_member "data/openwrt-one-initramfs.itb"
  @ubi_resource "openwrt-one-nand.ubi"
  @ubi_member "data/openwrt-one-nand.ubi"

  @impl Mix.Task
  def run(_args) do
    unless Mix.target() == :openwrt_one do
      Mix.raise(
        "mix openwrt_one.firmware_post requires MIX_TARGET=openwrt_one (got #{inspect(Mix.target())})"
      )
    end

    fw = firmware_path()
    File.exists?(fw) || Mix.raise("Firmware not found at #{fw} — `mix firmware` did not produce it")

    wrap_script = Path.join([system_source_dir!(), "scripts", "wrap-firmware.sh"])

    File.exists?(wrap_script) ||
      Mix.raise("wrap-firmware.sh not found at #{wrap_script}")

    b2sum = find_b2sum!()

    work = Path.join(System.tmp_dir!(), "openwrt_one_post-#{System.unique_integer([:positive])}")
    File.mkdir_p!(work)

    try do
      itb = Path.join(work, "openwrt-one-initramfs.itb")
      ubi = Path.join(work, "openwrt-one-nand.ubi")

      Mix.shell().info(
        "==> Building per-app initramfs.itb + nand.ubi from #{Path.relative_to_cwd(fw)}"
      )

      run_cmd!(wrap_script, [fw, itb, ubi])

      Mix.shell().info("==> Splicing per-app resources into #{Path.relative_to_cwd(fw)}")

      itb_length = File.stat!(itb).size
      itb_hash = b2sum_256!(b2sum, itb)
      ubi_length = File.stat!(ubi).size
      ubi_hash = b2sum_256!(b2sum, ubi)

      rewrite_fw!(fw, work,
        itb: {itb, itb_length, itb_hash},
        ubi: {ubi, ubi_length, ubi_hash}
      )
    after
      File.rm_rf!(work)
    end
  end

  defp rewrite_fw!(fw, work, replacements) do
    extract = Path.join(work, "extract")
    File.mkdir_p!(extract)
    fw_charlist = String.to_charlist(fw)
    extract_charlist = String.to_charlist(extract)

    {:ok, listing} = :zip.list_dir(fw_charlist)

    order =
      for entry <- listing, match?({:zip_file, _, _, _, _, _}, entry) do
        {:zip_file, name, _info, _comment, _offset, _comp_size} = entry
        List.to_string(name)
      end

    {:ok, _} = :zip.extract(fw_charlist, [{:cwd, extract_charlist}])

    {itb_file, itb_length, itb_hash} = Keyword.fetch!(replacements, :itb)
    {ubi_file, ubi_length, ubi_hash} = Keyword.fetch!(replacements, :ubi)

    File.cp!(itb_file, Path.join(extract, @itb_member))
    File.cp!(ubi_file, Path.join(extract, @ubi_member))

    meta_path = Path.join(extract, "meta.conf")

    meta_path
    |> File.read!()
    |> update_meta(@itb_resource, itb_length, itb_hash)
    |> update_meta(@ubi_resource, ubi_length, ubi_hash)
    |> then(&File.write!(meta_path, &1))

    File.rm!(fw)
    files = for name <- order, do: String.to_charlist(name)
    {:ok, _} = :zip.create(fw_charlist, files, [{:cwd, extract_charlist}])
  end

  defp update_meta(meta, resource, length, hash) do
    pattern =
      ~r/(file-resource\s+"#{Regex.escape(resource)}"\s*\{[^}]*?length=)\d+([^}]*?blake2b-256=)[0-9a-fA-F]+/

    case Regex.run(pattern, meta) do
      nil ->
        Mix.raise("Could not find file-resource \"#{resource}\" block in meta.conf to update")

      _ ->
        Regex.replace(pattern, meta, fn _, prefix, mid -> "#{prefix}#{length}#{mid}#{hash}" end)
    end
  end

  defp firmware_path do
    app = Mix.Project.config()[:app]
    Path.join([Mix.Project.build_path(), "nerves", "images", "#{app}.fw"])
  end

  defp system_source_dir! do
    case Mix.Project.deps_paths()[:nerves_system_openwrt_one] do
      nil ->
        Mix.raise("nerves_system_openwrt_one is not in your project's deps")

      path ->
        Path.expand(path)
    end
  end

  defp find_b2sum! do
    cond do
      System.find_executable("b2sum") ->
        "b2sum"

      System.find_executable("gb2sum") ->
        "gb2sum"

      true ->
        Mix.raise("""
        b2sum (GNU coreutils) not found in PATH.

          Linux:  apt install coreutils
          macOS:  brew install coreutils  (binary is gb2sum)
        """)
    end
  end

  defp b2sum_256!(b2sum, file) do
    case System.cmd(b2sum, ["-l", "256", file]) do
      {out, 0} ->
        out |> String.split() |> List.first()

      {out, code} ->
        Mix.raise("#{b2sum} -l 256 #{file} failed (exit #{code}): #{out}")
    end
  end

  defp run_cmd!(cmd, args) do
    case System.cmd(cmd, args, into: IO.stream(:stdio, :line), stderr_to_stdout: true) do
      {_, 0} -> :ok
      {_, code} -> Mix.raise("#{Path.basename(cmd)} failed with exit code #{code}")
    end
  end
end