#!/bin/bash --norc

kver="$(uname -r)"

boot_dir="/boot"
quiet=0
host_only=0
force=0
set_hostonly=0
no_hostonly=0
dracut_dbg=0

error() { echo "$@" >&2; }

usage () {
    local cmd=""
    [[ $1 = '-n' ]] && cmd=echo || cmd=error

    $cmd "usage: ${0##*/} [--version] [--help] [-v] [-f] [--preload <module>]"
    $cmd "       [--image-version] [--with=<module>]"
    $cmd "       [--nocompress]"
    $cmd "       [--set-hostonly]"
    $cmd "       [--no-hostonly]"
    $cmd "       [--debug]"
    $cmd "       <initrd-image> <kernel-version>"
    $cmd ""
    $cmd "       (ex: ${0##*/} /boot/initramfs-$kver.img $kver)"

    [[ $1 = '-n' ]] && exit 0
    exit 1
}

# Little helper function for reading args from the commandline.
# it automatically handles -a b and -a=b variants, and returns 1 if
# we need to shift $3.
read_arg() {
    # $1 = arg name
    # $2 = arg value
    # $3 = arg parameter
    local param="$1"
    local rematch='^[^=]*=(.*)$' result
    if [[ $2 =~ $rematch ]]; then
        read "$param" <<< "${BASH_REMATCH[1]}"
    else
	for ((i=3; $i <= $#; i++)); do
            # Only read next arg if it not an arg itself.
            if [[ ${@:$i:1} = -* ]];then
		break
            fi
            result="$result ${@:$i:1}"
            # There is no way to shift our callers args, so
            # return "no of args" to indicate they should do it instead.
	done
	read "$1" <<< "$result"
        return $(($i - 3))
    fi
}

# For photon
default_kernel_images() {
  local kernel_version=""

  local installed_kernels=

  installed_kernels="$(find /boot/ -mindepth 1 -maxdepth 1 -type f \
                  -iname "vmlinuz-*.ph[[:digit:]]*" -exec basename {} \; | \
                  cut -d '-' -f 2-)"

  for kernel_version in ${installed_kernels}; do
    local modules_dir="/lib/modules/${kernel_version}"
    if [ ! -d "${modules_dir}" ]; then
      echo -e "\nERROR: ${modules_dir} not found for vmlinuz-${kernel_version}\n" >&2
      continue
    fi
    # Take this directory as the source of truth
    kernels="$kernels $kernel_version"
    targets="$targets $boot_dir/initrd.img-$kernel_version"
  done
}

while (($# > 0)); do
    case ${1%%=*} in
        --with-usb) read_arg usbmodule "$@" || shift $?
            basicmodules="$basicmodules ${usbmodule:-usb-storage}"
            unset usbmodule;;
        --with-avail) read_arg modname "$@" || shift $?
            basicmodules="$basicmodules $modname";;
        --with) read_arg modname "$@" || shift $?
            basicmodules="$basicmodules $modname";;
        --version)
            echo "mkinitrd: dracut compatibility wrapper"
            exit 0;;
        -v|--verbose) dracut_args="${dracut_args} -v";;
        -f|--force) force=1;;
        --preload) read_arg modname "$@" || shift $?
            basicmodules="$basicmodules $modname";;
        --image-version) img_vers=yes;;
        --rootfs|-d) read_arg rootfs "$@" || shift $?
            dracut_args="${dracut_args} --filesystems $rootfs";;
        --nocompress) dracut_args="$dracut_args --no-compress";;
        --set-hostonly) set_hostonly=1;;
        --no-hostonly) no_hostonly=1;;
        --debug) dracut_dbg=1;;
        --help) usage -n;;
        --builtin) ;;
        --without*) ;;
        --without-usb) ;;
        --fstab*) ;;
        --ifneeded) ;;
        --omit-scsi-modules) ;;
        --omit-ide-modules) ;;
        --omit-raid-modules) ;;
        --omit-lvm-modules) ;;
        --omit-dmraid) ;;
        --allow-missing) ;;
        --net-dev*) ;;
        --noresume) ;;
        --rootdev*) ;;
        --thawdev*) ;;
        --rootopts*) ;;
        --root*) ;;
        --loopdev*) ;;
        --loopfs*) ;;
        --loopopts*) ;;
        --looppath*) ;;
        --dsdt*) ;;
        --bootchart) ;;
	-s) ;;
	--quiet|-q) quiet=1;;
	-b) read_arg boot_dir "$@" || shift $?
	    if [ ! -d $boot_dir ];then
		error "Boot directory $boot_dir does not exist"
		exit 1
	    fi
	    ;;
	-k) # Would be nice to get a list of images here
	    read_arg kernel_images "$@" || shift $?
	    for kernel_image in $kernel_images;do
		kernels="$kernels ${kernel_image#*-}"
	    done
	    host_only=1
	    force=1
	    ;;
	-i) read_arg initrd_images "$@" || shift $?
	    for initrd_image in $initrd_images;do
		targets="$targets $boot_dir/$initrd_image"
	    done
	    ;;
        *)  if [[ ! $targets ]]; then
            targets=$1
            elif [[ ! $kernels ]]; then
            kernels=$1
            else
            usage
            fi;;
    esac
    shift
done

[[ $targets && $kernels ]] || default_kernel_images
[[ $targets && $kernels ]] || error "No Kernel Registered"

# We can have several targets/kernels, transform the list to an array
targets=( $targets )
[[ $kernels ]] && kernels=( $kernels )

# don't set hostonly flag if running in docker env
# if set initrd will be incomplete
# https://fedoraproject.org/wiki/Features/DracutHostOnly#Detailed_Description
# https://man7.org/linux/man-pages/man5/dracut.conf.5.html
if [ ${set_hostonly} -eq 0 ]; then
  if grep -qc docker /proc/self/cgroup; then
    echo "--- Generating initrd under docker environment ---"
    host_only=0
  elif [ "$(stat -c %d:%i /)" != "$(stat -c %d:%i /proc/1/root/.)" ]; then
    # set host_only=0 if initrd is generated in chroot
    # otherwise cloud images generated in a Photon host will have bad initrd
    echo "--- Generating initrd under chroot environment ---"
    host_only=0
  else
    host_only=1
  fi
else
  host_only=1
fi
force=1

if [ ${no_hostonly} -eq 1 ]; then
  host_only=0
fi

[[ $host_only == 1 ]] && dracut_args="${dracut_args} -H"
[[ $force == 1 ]]     && dracut_args="${dracut_args} -f"

initrd_sanity_check() {
  local ret=0
  local status="$1"
  local log_fn="$2"
  local target="$3"
  local errmsg=""

  # initrd issues can be fatal, so if anything goes wrong during initrd creation
  # print errors to console
  if [ ${status} -ne 0 ]; then
    errmsg=$(cat << EOF
\n\n------------------------ ERROR NOTICE ------------------------------
          DRACUT RETURNED NON-ZERO EXIT STATUS(${status})

PROBABLY ${target} IS FAULTY
SYSTEM MAY BECOME UNUSABLE POST REBOOT
---------------------- PROCEED WITH CAUTION ------------------------
EOF
)
    echo -e "${errmsg}" 1>&2 |& tee -a "${log_fn}"
    ret=1
  fi

  if ! lsinitrd ${target} 1>/dev/null; then
    errmsg=$(cat << EOF
\n\n------------------------ ERROR NOTICE ------------------------------
                  lsinitrd ${target} FAILED

PROBABLY ${target} IS FAULTY
SYSTEM MAY BECOME UNUSABLE POST REBOOT
---------------------- PROCEED WITH CAUTION -----------------------\n\n
EOF
)
    echo -e "${errmsg}" 1>&2 |& tee -a "${log_fn}"
    ret=1
  fi

  return ${ret}
}

# If running in non tty, enable verbose
# It's probably happening in some kind of build env, we need verbose.
if ! test -t 1; then
  quiet=0
fi

final_ret=0

for ((i = 0; i < ${#targets[@]}; i++)); do
  ret=0
  if [[ $img_vers ]];then
    target="${targets[$i]}-${kernels[$i]}"
  else
    target="${targets[$i]}"
  fi
  kernel="${kernels[$i]}"

  log_fn="/var/log/mkinitrd-${kernel}.log"

  echo "Creating $target"

  if [ $dracut_dbg -ne 0 ]; then
    export DRACUT_INSTALL="/usr/lib/dracut/dracut-install --verbose"
    dracut_cmd=(dracut -L 6 $dracut_args)
  else
    dracut_cmd=(dracut $dracut_args)
  fi

  # this check is for combination of newer and older kernels
  if [ -d "/lib/modules/${kernel}/dracut.conf.d" ]; then
    dracut_cmd+=(--confdir \"/lib/modules/${kernel}/dracut.conf.d /etc/dracut.conf.d\")
  elif [ -d "/var/lib/initramfs/kernel" ]; then
    kernel_ver_dir="/var/lib/initramfs/kernel"
    if [ -s "$kernel_ver_dir/$kernel" ]; then
      dracut_cmd+=("$(cat $kernel_ver_dir/$kernel)")
    fi
    unset kernel_ver_dir
  else
    error "ERROR: either /lib/modules/${kernel}/dracut.conf.d and /var/lib/initramfs/kernel not present"
    exit 1
  fi

  if [[ $basicmodules ]]; then
    dracut_cmd+=(--add-drivers "$basicmodules")
  fi
  dracut_cmd+=("$target" "$kernel")

  if [ ${quiet} -eq 1 ]; then
    eval "${dracut_cmd[@]}" &> "${log_fn}"
  else
    eval "${dracut_cmd[@]}" |& tee "${log_fn}"
  fi

  ret=${PIPESTATUS[0]}
  if ! initrd_sanity_check "${ret}" "${log_fn}" "${target}"; then
    final_ret=1
  fi
done

[ $final_ret -ne 0 ] && error "--- ERROR: mkinitrd FAILED, SOMETHING WENT WRONG ---"

exit ${final_ret}
