#!/usr/bin/env bash
# Mkinitramfs creates a basic init ram fs with encryption support.
# Copyright (C)  2023  Aaron Ball <nullspoon@oper.io>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


# Hash of fully-qualified binaries
declare -a fqbins

#
# Searches PATH for all binaries listed as function arguments. Each found
# binary is then put into the global "bins" array.
#
# This function also serves to ensure all specified binaries are present and
# in PATH. It will exit code 1 with an error message if any binaries are not
# found.
#
# Note: To use this function, insert "declare -a fqbins" at top of script).
#
# @param bins All function arguments are the names of binaries to locate.
#
resolve_bins() {
  local args=${@}
  local path=''

  for i in ${args[@]}; do
    if ! path="$(type -p ${i} 2>/dev/null)"; then
      printf 'Could not find binary %s. Is it installed?\n' "${i}"
      exit 1
    fi

    printf 'Located %s at %s\n' "${i}" "${path}"
    fqbins+=("${path}")
  done
}


#
# Gets first path locatable in a string
#
get_first_path() {
  local str="${1:-}"
  local out=''

  # Find the first occurrence of `/`
  for ((i=0; i<${#str}; i++)); do
    [ "${str:$i:1}" = '/' ] && break
  done

  # Exit if no path found
  [ "${i}" -eq "${#str}" ] && return 2

  # Find end of string by looking for space
  while [ "${str:$i:1}" != ' ' ] && [ "${i}" -lt "${#str}" ]; do
    out="${out}${str:$i:1}"
    (( i += 1 ))
  done
  printf '%s\n' "${out}"
}


get_deps() {
  local bin="${1:-}"
  for dep in $(ldd ${bin}); do
     get_first_path "${dep}"
  done
}


#
# Creates cpio file for specified initrd cache directory.
# Outputs to /boot/initrd-${version}.
#
# @param cache   Path to the uncompressed source directory to be archived
# @param version Version of initrd archive
#
mkcpio() {
  [ -z "${1:-}" ] && printf 'Initramfs cache path required.\n' && exit 1
  [ -z "${2:-}" ] && printf 'Kernel version required.\n' && exit 1

  local cache="${1}"
  local version="${2}"

  cd "${cache}"
  fspath="/boot/initrd-${version}"

  # Notify user of initrd overwriting
  [ -f "${fspath}" ] && printf "\n\n%s exists. Overwriting.\n" "${fspath}"

  # Create the initrd
  printf 'Building initrd to %s\n' "${fspath}"
  find . -print0 | cpio --null -o --format=newc | xz -C crc32 -9 -c > "${fspath}"
}


#
# Sets up cache dir for later writing to an archive (via cpio).
# Performs simple operations like creating a basic directory structure,
#
# @param cache Path to the cache directory. Will be created if not exists
#
cache_dir_setup() {
  [ -z "${1}" ] && printf 'Please specify a cache dir.\n' && exit 1
  local cache="${1:-}"

  # Clean the cache so we start fresh
  rm -r "${cache:?}" && mkdir "${cache}"

  # Create the temporary directory structure
  for i in bin dev etc lib lib32 lib64 mnt/root proc root run sbin sys usr/share; do
    mkdir -p "${cache}/${i}"
  done

  # Copy in terminals database for `l` (linux - stock tty)
  mkdir -p "${cache}/usr/share/terminfo"
  cp -r /usr/share/terminfo/l "${cache}/usr/share/terminfo/l"
}


#
# Checks installed kernel modules for the specified version to see if the
# corresponding kernel has encryption support statically compiled in.
#
# @param version Version of kernel modules to check
#
check_crypto_support() {
  [ -z "${1:-}" ] && printf 'Kernel version required.\n' && exit 1
  local version="${1}"
  local buf=''
  local -a copy=(
    arch/x86/crypto
    crypto
    drivers/input/{keyboard,serio,*.ko}
    drivers/{crypto,usb,platform,md}
    lib
    fs/{btrfs,xfs,ext4}
  )

  mkdir -p "${cache}/lib/modules/${version}"
  cp -v /lib/modules/${version}/modules*.bin "${cache}/lib/modules/${version}/"

  # Copy in all the rewuired kernel modules
  for i in ${copy[@]}; do
    buf=(/lib/modules/${version}/kernel/${i})
    # For globbing in the array
    for ent in ${buf[@]}; do
      printf 'Copying in %s\n' "${ent}"
      mkdir -p "$(dirname ${cache}/${ent})"
      cp -r "${ent}" "${cache}/${ent}"
    done
  done
}


#
# Checks to see if kernel modules for the specified version have been
# installed.
#
# @param version Version to check for installation status
#
check_kernel_version() {
  [ -z ${1} ] && printf 'Kernel version required.\n' && exit 1
  local version="${1}"
  local modulespath="/lib/modules/${version}"

  if [ ! -d ${modulespath} ]; then
    printf 'Error: Could not find %s.\n' "${modulespath}"
    printf 'Has kernel version %s been comiled?' "${version}"
    exit 1
  fi
}


main() {
  [ -z "${1:-}" ] && printf 'Please specify a kernel version\n' && exit 1

  local version=${1}

  check_kernel_version ${version}

  local cache='/tmp/initrd'
  local res_path=/usr/share/mkinitramfs/

  # Ensure init script has no syntax errors (linting process)
  if ! bash -n ${res_path}/init; then
    printf '\n\nERROR: Syntax errors were found in the init script.\n'
    printf 'Please resolve the above issues and try again.\n'
    printf 'Exiting to avoid creating an unbootable initramfs.\n\n'
    return 1
  fi

  # List of binaries to exist in the new initramfs
  # Standard shell binaries
  resolve_bins bash pidof ps ls grep less switch_root kill
  # Module management
  resolve_bins udevadm udevd modprobe rmmod
  # For encrypted and block device support
  resolve_bins partprobe dmsetup mount umount cryptsetup

  # For debugging
  #resolve_bins sleep mkdir chmod chown tr clear cat

  # Set up the archive source dir.
  cache_dir_setup ${cache}

  # Ensure static crypto support exists in kernel
  check_crypto_support ${version}

  # Copy in libgcc_s for newer cryptsetup
  install -v -D /usr/lib/libgcc_s.so   "${cache}/usr/lib/libgcc_s.so"
  install -v -D /usr/lib/libgcc_s.so.1 "${cache}/usr/lib/libgcc_s.so.1"

  # Install udev rules for block devices
  mkdir -p "${cache}/lib/udev"
  cp -rv /lib/udev/rules.d "${cache}/lib/udev/rules.d"

  # Copy binary and dependencies to cache dir
  for bin in ${fqbins[@]}; do
    # Copy the binary of interest
    install -v -D "${bin}" "${cache}/${bin}"

    local deps=$(get_deps ${bin})
    # Copy each of the binary's deps
    for dep in ${deps[@]}; do
      install -v -D ${dep} "${cache}/${dep}"
    done
  done

  # Copy in message file
  [ -f "/etc/msg" ] && install -v -D "/etc/msg" "${cache}/etc/msg"
  # Copy in init script and add execute
  install -v -D -m 755 "${res_path}/init" "${cache}/init"
  # Copy in the system shell profile
  install -v -D "/etc/profile" "${cache}/etc/profile"

  # Create archive image
  mkcpio ${cache} ${version}
}

main ${@}