#!/bin/bash
##########################################################################
# Script....: /tools/edit-initrd (OS InitRD and Slackware Installer)
# Purpose...: Unpack an initrd, drop into an interactive shell in the
#            extracted tree, then repack it when you exit the shell.
#            This is a development tool for testing changes to the Initial
#            RAM disk.
# Designer..: Stuart Winter <mozes@slackware.com>
# Programmer: Vibe coded with ChatGPT, tweaked by MoZes.
#
# Supports: uncompressed, gzip, xz.
# Repack:   gzip (max) or xz (max)
# Typical use:
#   edit-initrd -i /boot/initrd.gz        # defaults to repack using gzip (regardless of input compression)
#   edit-initrd -i /boot/initrd.gz -c xz  # repack as xz
#   edit-initrd -i initrd.xz -s /bin/bash # use a specific shell
#
# Notes:
# - Repacking overwrites the original initrd unless -o is provided.
##########################################################################

set -euo pipefail

usage() {
  cat >&2 <<'EOF'
Usage:
  edit-initrd -i <initrd> [-o <output>] [-c gzip|xz|auto] [-s <shell>] [-k]

Options:
  -i <initrd>   Path to initrd image (required)
  -o <output>   Output path (default: overwrite input initrd)
  -c <mode>     Repack compression: gzip, xz, or auto (default: gzip)
                auto = reuse detected compression; if none, defaults to gzip
  -s <shell>    Shell to run (default: $SHELL or /bin/bash)
  -k            Keep temporary extraction directory (for debugging)

Examples:
  edit-initrd -i /boot/initrd-armv8
  edit-initrd -i /boot/initrd-armv8 -c gzip -o /tmp/initrd-armv8
EOF
  exit 2
}

INITRD=""
OUT=""
REPACK_MODE="gzip"
# Auto is nice, but the installers are delivered xz compressed. It's faster to repack
# with gzip, which is the best for use during development.
#REPACK_MODE="auto"
SHELL_BIN="${SHELL:-/bin/bash}"
KEEP_TMP=0

export XZ_OPTS="--threads 0 -ez9f -C crc32"

while getopts ":i:o:c:s:kh" opt; do
  case "$opt" in
    i) INITRD="$OPTARG" ;;
    o) OUT="$OPTARG" ;;
    c) REPACK_MODE="$OPTARG" ;;
    s) SHELL_BIN="$OPTARG" ;;
    k) KEEP_TMP=1 ;;
    h) usage ;;
    \?) echo "Unknown option: -$OPTARG" >&2; usage ;;
    :)  echo "Option -$OPTARG requires an argument." >&2; usage ;;
  esac
done

[[ -n "$INITRD" ]] || { echo "ERROR: -i <initrd> is required." >&2; usage; }
[[ -r "$INITRD" ]] || { echo "ERROR: Cannot read initrd: $INITRD" >&2; exit 1; }

if [[ -z "$OUT" ]]; then
  OUT="$INITRD"
fi

if [[ ! -x "$SHELL_BIN" ]]; then
  echo "ERROR: Shell not executable: $SHELL_BIN" >&2
  exit 1
fi

need_cmd() { command -v "$1" >/dev/null 2>&1 || { echo "ERROR: Missing required command: $1" >&2; exit 1; }; }
need_cmd cpio
need_cmd gzip
need_cmd xz
need_cmd mktemp

tmpdir="$(mktemp -d -t edit-initrd.XXXXXX)"
cleanup() {
  if (( KEEP_TMP )); then
    echo "Keeping temp directory: $tmpdir" >&2
  else
    rm -rf "$tmpdir"
  fi
}
trap cleanup EXIT

# Compression detection is done by magic bytes only (no 'file' utility required):
#  - gzip: 1f 8b
#  - xz:   fd 37 7a 58 5a 00
detect_compression() {
  # Returns: gzip | xz | none
  local f="$1"
  if head -c 2 "$f" | LC_ALL=C od -An -tx1 | tr -d ' \n' | grep -qi '^1f8b'; then
    echo "gzip"
  elif head -c 6 "$f" | LC_ALL=C od -An -tx1 | tr -d ' \n' | grep -qi '^fd377a585a00'; then
    echo "xz"
  else
    echo "none"
  fi
}

unpack_initrd() {
  local in="$1" comp="$2" outdir="$3"
  echo "Unpacking: $in"
  mkdir -p "$outdir"
  (
    cd "$outdir"
    case "$comp" in
      gzip) gzip -dc "$in" | cpio -idm ;;
      xz)   xz -dc "$in"   | cpio -idm ;;
      none) cpio -idm  < "$in" ;;
      *)    echo "ERROR: Unknown compression: $comp" >&2; exit 1 ;;
    esac
  )
}

choose_repack_mode() {
  local detected="$1" requested="$2"
  case "$requested" in
    auto)
      if [[ "$detected" == "gzip" || "$detected" == "xz" ]]; then
        echo "$detected"
      else
        echo "gzip"
      fi
      ;;
    gzip|xz) echo "$requested" ;;
    *)
      echo "ERROR: Invalid -c mode: $requested (use gzip, xz, or auto)" >&2
      exit 1
      ;;
  esac
}

repack_initrd() {
  local indir="$1" mode="$2" out="$3"
  local outtmp="${out}.tmp.$$"

  echo "Repacking to: $out (mode: $mode)"

  # Create cpio archive in a stable way (sorted paths), then compress.
  # Verbose compression requested; cpio verbosity kept quiet to avoid huge spew.
  (
    cd "$indir"
    # Important: do not include '.' itself; feed relative paths.
    find . -mindepth 1 -print0 \
      | LC_ALL=C sort -z \
      | cpio --null -o -H newc 2>/dev/null \
      | case "$mode" in
          gzip) gzip -9 -v > "$outtmp" ;;
          xz)   xz $XZ_OPTS > "$outtmp" ;;
        esac
  )

  # Preserve permissions on overwrite where possible; atomic replace.
  chmod --reference="$out" "$outtmp" 2>/dev/null || true
  mv -f "$outtmp" "$out"
}

COMP_DETECTED="$(detect_compression "$INITRD")"
REPACK_CHOSEN="$(choose_repack_mode "$COMP_DETECTED" "$REPACK_MODE")"

echo "Initrd:        $INITRD"
echo "Detected:      $COMP_DETECTED"
echo "Repack mode:   $REPACK_CHOSEN"
echo "Output:        $OUT"
echo "Work dir:      $tmpdir"
echo

unpack_initrd "$INITRD" "$COMP_DETECTED" "$tmpdir"

echo
echo "Entering shell in: $tmpdir"
echo "Tip: edit files, add/remove modules/firmware, etc. Exit the shell to repack."
echo

(
  cd "$tmpdir"
  export PS1="initrd-root# "
  exec "$SHELL_BIN" -i
)

echo
repack_initrd "$tmpdir" "$REPACK_CHOSEN" "$OUT"
echo "Done."

#read -p "Press ENTER to reboot, or CTRL-C to abort"
#reboot