#!/bin/bash


# This file based in part on the mkinitramfs script for the LFS LiveCD (And later Tucana :))
# written by Alexander E. Patrakov and Jeremy Huntwork.

# Heavily modified by Rahul Chandra
# - Added forceful early module loading
# - Added live mode to include all modules
# - Fixed udev by adding cgroup support, adding /etc/group, and libkmod files

LIVE_MODE=false

# Check arguments
for arg in "$@"; do
  case "$arg" in
    --live)
      LIVE_MODE=true
      ;;
    *)
      if [[ "$arg" != "--"* ]]; then
        KERNEL_VERSION=$arg
        INITRAMFS_FILE=initrd.img-$KERNEL_VERSION
      fi
      ;;
  esac
done

if [ -z "$INITRAMFS_FILE" ]; then
  INITRAMFS_FILE=initrd.img-no-kmods
fi

copy()
{
  local file

  if [ "$2" = "lib" ]; then
    file=$(PATH=/usr/lib type -p $1)
  else
    file=$(type -p $1)
  fi

  if [ -n "$file" ] ; then
    cp $file $WDIR/usr/$2
  else
    echo "Missing required file: $1 for directory $2"
    rm -rf $WDIR
    exit 1
  fi
}


if [ -n "$KERNEL_VERSION" ] && [ ! -d "/usr/lib/modules/$1" ] ; then
  echo "No modules directory named $1"
  exit 1
fi

printf "Creating $INITRAMFS_FILE... "
depmod $KERNEL_VERSION
binfiles="sh cat cp dd killall ls mkdir mknod mount find"
binfiles="$binfiles umount sed sleep ln rm uname"
binfiles="$binfiles readlink basename"

# Systemd installs udevadm in /bin. Other udev implementations have it in /sbin
if [ -x /usr/bin/udevadm ] ; then binfiles="$binfiles udevadm"; fi

sbinfiles="modprobe blkid switch_root"

# Optional files and locations
for f in mdadm mdmon udevd udevadm; do
  if [ -x /usr/sbin/$f ] ; then sbinfiles="$sbinfiles $f"; fi
done

# Add lvm if present (cannot be done with the others because it
# also needs dmsetup
if [ -x /usr/sbin/lvm ] ; then sbinfiles="$sbinfiles lvm dmsetup"; fi

unsorted=$(mktemp /tmp/unsorted.XXXXXXXXXX)

DATADIR=/usr/share/mkinitramfs
INITIN=init.in

# Create a temporary working directory
WDIR=$(mktemp -d /tmp/initrd-work.XXXXXXXXXX)

# Create base directory structure
mkdir -p $WDIR/{dev,run,sys,proc,usr/{bin,lib/{firmware,modules},sbin}}
mkdir -p $WDIR/etc/{modprobe.d,udev/rules.d}
touch $WDIR/etc/modprobe.d/modprobe.conf

# Create /etc/group so udev will work
cat > $WDIR/etc/group << "FOF"
kmem:x:3:
tape:x:4:
tty:x:5:
disk:x:8:
lp:x:9:
dialout:x:10:
audio:x:11:
video:x:12:
cdrom:x:15:
input:x:24:
kvm:x:61:
FOF

# Copy modprobe.conf

if [ -d /etc/modprobe.d ]; then
   cp -rp /etc/modprobe.d/* $WDIR/etc/modprobe.d/
fi
ln -s usr/bin  $WDIR/bin
ln -s usr/lib  $WDIR/lib
ln -s usr/sbin $WDIR/sbin
ln -s lib      $WDIR/lib64

# Create necessary device nodes
mknod -m 640 $WDIR/dev/console c 5 1
mknod -m 664 $WDIR/dev/null    c 1 3

# Install the udev configuration files
if [ -f /etc/udev/udev.conf ]; then
  cp /etc/udev/udev.conf $WDIR/etc/udev/udev.conf
fi

for file in $(find /etc/udev/rules.d/ -type f) ; do
  cp $file $WDIR/etc/udev/rules.d
done





# Copy the RAID configuration file if present
if [ -f /etc/mdadm.conf ] ; then
  cp /etc/mdadm.conf $WDIR/etc
fi

# Install the init file
install -m0755 $DATADIR/$INITIN $WDIR/init

# Add the early boot modules
source /etc/early_load_modules.conf
for EARLY_LOAD_MODULE in ${EARLY_LOAD_MODULES[@]}; do
    sed -i "s/do_try_resume\ /modprobe\ $EARLY_LOAD_MODULE\ndo_try_resume\ /" $WDIR/init
done

if [  -n "$KERNEL_VERSION" ] ; then
  if [ -x /usr/bin/kmod ] ; then
    binfiles="$binfiles kmod"
  else
    binfiles="$binfiles lsmod"
    sbinfiles="$sbinfiles insmod"
  fi
fi

# Install basic binaries
for f in $binfiles ; do
  ldd /usr/bin/$f | sed "s/\t//" | cut -d " " -f1 >> $unsorted
  copy /usr/bin/$f bin
done

for f in $sbinfiles ; do
  ldd /usr/sbin/$f | sed "s/\t//" | cut -d " " -f1 >> $unsorted
  copy $f sbin
done

# Add udevd libraries if not in /usr/sbin
if [ -x /usr/lib/udev/udevd ] ; then
  ldd /usr/lib/udev/udevd | sed "s/\t//" | cut -d " " -f1 >> $unsorted
elif [ -x /usr/lib/systemd/systemd-udevd ] ; then
  ldd /usr/lib/systemd/systemd-udevd | sed "s/\t//" | cut -d " " -f1 >> $unsorted
fi

# Add libkmod
echo libkmod.so.2 >> $unsorted

# Add module symlinks if appropriate
if [ -n "$KERNEL_VERSION" ] && [ -x /usr/bin/kmod ] ; then
  ln -s kmod $WDIR/usr/bin/lsmod
  ln -s kmod $WDIR/usr/bin/insmod
fi



# Add lvm symlinks if appropriate
# Also copy the lvm.conf file
if  [ -x /usr/sbin/lvm ] ; then
  ln -s lvm $WDIR/usr/sbin/lvchange
  ln -s lvm $WDIR/usr/sbin/lvrename
  ln -s lvm $WDIR/usr/sbin/lvextend
  ln -s lvm $WDIR/usr/sbin/lvcreate
  ln -s lvm $WDIR/usr/sbin/lvdisplay
  ln -s lvm $WDIR/usr/sbin/lvscan

  ln -s lvm $WDIR/usr/sbin/pvchange
  ln -s lvm $WDIR/usr/sbin/pvck
  ln -s lvm $WDIR/usr/sbin/pvcreate
  ln -s lvm $WDIR/usr/sbin/pvdisplay
  ln -s lvm $WDIR/usr/sbin/pvscan

  ln -s lvm $WDIR/usr/sbin/vgchange
  ln -s lvm $WDIR/usr/sbin/vgcreate
  ln -s lvm $WDIR/usr/sbin/vgscan
  ln -s lvm $WDIR/usr/sbin/vgrename
  ln -s lvm $WDIR/usr/sbin/vgck
  # Conf file(s)
  cp -a /etc/lvm $WDIR/etc
fi

# Install libraries
sort $unsorted | uniq | while read library ; do
# linux-vdso and linux-gate are pseudo libraries and do not correspond to a file
# libsystemd-shared is in /lib/systemd, so it is not found by copy, and
# it is copied below anyway
  if [[ "$library" == linux-vdso.so.1 ]] ||
     [[ "$library" == linux-gate.so.1 ]] ||
     [[ "$library" == libsystemd-shared* ]]; then
    continue
  fi

  copy $library lib
done

if [ -d /usr/lib/udev ]; then
  cp -a /usr/lib/udev $WDIR/usr/lib
fi
if [ -d /usr/lib/systemd ]; then
  cp -a /usr/lib/systemd $WDIR/usr/lib
fi
if [ -d /usr/lib/elogind ]; then
  cp -a /usr/lib/elogind $WDIR/usr/lib
fi


# Install the kernel modules if requested
if [ -n "$KERNEL_VERSION" ]; then
  # This whoel situation leads to some serious module conflicts with things like the SOF firmware
  # if $LIVE_MODE; then
  #   echo "Including storage modules for live mode"
  #   find \
  #     /usr/lib/modules/$KERNEL_VERSION/kernel/{block,crypto,fs,lib,arch,virt} \
  #     /usr/lib/modules/$KERNEL_VERSION/kernel/drivers/{mmc,misc,block,ata,cdrom,nvme,md,firewire} \
  #     /usr/lib/modules/$KERNEL_VERSION/kernel/drivers/{scsi,message,pcmcia,virtio,char,media,video,vfio,i2c,input} \
  #     /usr/lib/modules/$KERNEL_VERSION/kernel/drivers/usb/{host,storage} \
  #     -type f 2> /dev/null | cpio --make-directories -p --quiet $WDIR
  # else
  #   echo "Including currently loaded modules only"
  #   MODULES=$(lsmod | awk 'NR>1 {print $1}')
  #   for MODULE in $MODULES; do
  #     MODFILE=$(modinfo -n $MODULE)
  #     REL_PATH="${MODFILE#*/kernel/}"
  #     MODFILE="/usr/lib/modules/$KERNEL_VERSION/kernel/$REL_PATH"
  #     if [ -f "$MODFILE" ]; then
  #       echo "Copying $MODULE"
  #       cp --parents "$MODFILE" "$WDIR"
  #     else 
  #       echo "Could not find module $MODULE, regenerate initramfs after reboot"
  #     fi
  #   done
  # fi

  find \
      /usr/lib/modules/$KERNEL_VERSION/kernel/{block,crypto,fs,lib,arch,virt} \
      /usr/lib/modules/$KERNEL_VERSION/kernel/drivers/{mmc,misc,block,ata,cdrom,nvme,md,firewire} \
      /usr/lib/modules/$KERNEL_VERSION/kernel/drivers/{scsi,message,pcmcia,virtio,char,media,video,vfio,i2c,input} \
      /usr/lib/modules/$KERNEL_VERSION/kernel/drivers/usb/{host,storage} \
      -type f 2> /dev/null | cpio --make-directories -p --quiet $WDIR

  # Copy required metadata files
  cp /usr/lib/modules/$KERNEL_VERSION/modules.{builtin,order} \
     $WDIR/usr/lib/modules/$KERNEL_VERSION 2>/dev/null

  if [ -f /usr/lib/modules/$KERNEL_VERSION/modules.builtin.modinfo ]; then
    cp /usr/lib/modules/$KERNEL_VERSION/modules.builtin.modinfo \
       $WDIR/usr/lib/modules/$KERNEL_VERSION
  fi

  # Generate new module dependency files inside WDIR
  depmod -b $WDIR $KERNEL_VERSION
fi

# Install any firmware present
if [ -n "$KERNEL_VERSION" ]; then
  MODULES=$(find $WDIR/usr/lib/modules/$KERNEL_VERSION/kernel -type f 2>/dev/null)
  for MODULE in $MODULES; do
    FIRMWARE=$(modinfo $MODULE | grep firmware: | sed 's/firmware: *//')
    for FILE in $FIRMWARE; do
      if [ -f "/usr/lib/firmware/$FILE" ]; then
        cp --parents /usr/lib/firmware/$FILE $WDIR
      elif [ -f "/usr/lib/firmware/$FILE.xz" ]; then
        cp --parents /usr/lib/firmware/$FILE.xz $WDIR
      fi
    done
  done
fi


( cd $WDIR ; find . | cpio -o -H newc --quiet | gzip -9 ) > $INITRAMFS_FILE

# Prepare early loading of microcode if available
if ls /usr/lib/firmware/intel-ucode/* >/dev/null 2>&1 ||
   ls /usr/lib/firmware/amd-ucode/*   >/dev/null 2>&1; then

# first empty WDIR to reuse it
  rm -r $WDIR/*

  DSTDIR=$WDIR/kernel/x86/microcode
  mkdir -p $DSTDIR

  if [ -d /usr/lib/firmware/amd-ucode ]; then
    cat /usr/lib/firmware/amd-ucode/microcode_amd*.bin.xz > $DSTDIR/AuthenticAMD.bin
  fi

  if [ -d /usr/lib/firmware/intel-ucode ]; then
    cat /usr/lib/firmware/intel-ucode/* > $DSTDIR/GenuineIntel.bin
  fi

  ( cd $WDIR; find . | cpio -o -H newc --quiet ) > microcode.img
  cat microcode.img $INITRAMFS_FILE > tmpfile
  mv tmpfile $INITRAMFS_FILE
  rm microcode.img
fi

# Remove the temporary directories and files
rm -rf $WDIR $unsorted
printf "done.\n"
