#!/bin/bash
# Init script for mkinitramfs
# 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/>.

export DEBUG=0
export QUIET=0
export INTERACTIVE=0
export ROOTDEV=''
export MOUNTOPTS=()

screen_init() {
  # Clear screen
  [ "${QUIET}" -eq 1 ] && printf '\033c'

  # Output message file if it exists
  [ -f /etc/msg ] && printf '%s\n' "$(</etc/msg)"
}


log() {
  local lvl="${1}"
  local msg="${2}"
  [ "${QUIET}" -eq 1 ] && [ "${lvl,,}" = "info"  ] && return

  printf -- "%s\n" "${msg}"
  if [ "${DEBUG}" -ne 0 ]; then
    echo "Press enter to continue"
    read
  fi
}
 
 
#
# Resolves the path to the root device from the cmdline root= syntax.
# Currently handles UUID syntax and /dev/sdx syntax, btrfs subvols, debugging,
# and quiet options.
#
parse_cmdline() {
  local cmdline="$(<${1:?})"

  for i in ${cmdline[@]}; do
    case "${i}" in
      root=*)
        if [ "${i:5:4}" == 'UUID' ]; then
          # mount by uuid
          ROOTDEV="/dev/disk/by-uuid/${i##*=}"
        elif [ "${i:5:5}" == 'LABEL' ]; then
          # mount by label
          ROOTDEV=/dev/disk/by-partlabel/${i##*=}
        else
          # mount by dev
          ROOTDEV="${i##*=}"
        fi
      ;;
      subvol=*)
        MOUNTOPTS+=(-o subvol=${i#*=})
      ;;
      initdebug)
        # Enable debug mode (this is gonna be slow)
        DEBUG=1
      ;;
      interactive)
        # Enable interactive mode
        INTERACTIVE=1
      ;;
      quiet)
        #  Enable quiet boot (less output)
        QUIET=1
      ;;
    esac
  done
}


#
# Takes an encrypted device path and executes cryptsetup luksOpen. Returns path
# to the new decrypted block device. This path takes the rough form of
# /dev/mapper/_dev_sdx
#
setup_encrypted() {
  local path="${1}"
  local name="${path//\//_}"

  # Decrypted block dev path is /dev/mapper/${name}
  cryptsetup luksOpen ${path} ${name}

  # Notify user and wait for input on failure
  if [ "$?" -gt 0 ]; then
    # Ensure decryption was successful, it not, drop to debug shell and reboot once done.
    printf 'ERROR: Could not decrypt device\n' >&2
    printf 'Dropping to debug shell (will automatically reboot on exit).\n' >&2
    /bin/bash --norc -i
    printf 'Rebooting\n' >&2
    printf 'b' > /proc/sysrq-trigger
  fi

  # Probe the new crypto device for an additional partition table
  partprobe "/dev/mapper/${name}" 2>/dev/null 1>/dev/null

  # Success. Return the path of the decrypted root device
  echo "/dev/mapper/${name}"
}


udev_populate() {
  # Start udevd if not already running
  pidof -s udevd > /dev/null || /sbin/udevd --daemon

  /sbin/udevadm trigger --type=devices --action=add --subsystem-match=block
  /sbin/udevadm trigger --type=devices --action=add --subsystem-match=input
  /sbin/udevadm settle

  # Stop udevd once devices are populated so main udevd can finish booting
  kill $(pidof udevd)
}


#
# Main function to keep the main operations code nicely separated from the
# rest.
#
main() {
  # Mount the /proc and /sys filesystems.
  mount -t tmpfs -o mode=0755,nousid,nodev none   /run
  mount -t sysfs -o nodev,noexec,nosuid    sysfs  /sys
  mount -t proc  -o nodev,noexec,nosuid    proc   /proc
  mount -t tmpfs                           shm    /dev/shm
  mount -t devtmpfs -o exec,nosuid,noatime dev    /dev
  mount -t devpts                          devpts /dev/pts

  local fakeroot='/mnt/root'

  if [ ! -d "${fakeroot}" ]; then
    log INFO "Fake root location ${fakeroot} does not exist. Cannot boot."
  fi

  parse_cmdline /proc/cmdline

  # display fanciful boot image
  screen_init
  udev_populate

  log INFO "Root device: ${ROOTDEV}"

  # Drop to maintenance shell if root device does not exist.
  if [ "$?" -ne 0 ]; then
    log ERROR "ERROR: Root device ${ROOTDEV} does not exist or can not be accessed."
    log ERROR "       Dropping to maintenance shell for troubleshooting."
    log ERROR "       (You may just need to wait longer for the root device)"
    /bin/bash -i
  fi

  # Modprobe atkbd to ensure user can type password
  for i in i8042 atkbd; do
    printf 'Loading required module [%s]\n' "${i}"
    modprobe "${i}" || printf 'Failed loading %s\n' "${i}"
  done

  if cryptsetup isLuks "${ROOTDEV}"; then
    # Set new rootdev location (/dev/mapper/something). This will update it to
    # the decrypted block device path.
    log INFO "Root device ${ROOTDEV} is encrypted."
    ROOTDEV=$(setup_encrypted ${ROOTDEV})
    log INFO "New rootdev: ${ROOTDEV}"
  fi

  # Drop to interactive shell if requested
  if [ "${INTERACTIVE}" -eq 1 ]; then
    log INFO "Interractive shell requested. Type 'exit' to continue boot sequence."
    /bin/bash --norc
  fi

  # Mount the fakeroot.
  log INFO "Mounting fakeroot"
  mount -o ro ${MOUNTOPTS[@]} ${ROOTDEV} ${fakeroot}

  # Ensure switch_root will be possible for destination fakeroot
  if [ ! -f "${fakeroot}/sbin/init" ]; then
    log ERROR "ERROR: Destination fakeroot ${fakeroot} does not have an init script."
    log ERROR "       Cannot proceed. Dropping to shell for troubleshooting."
    /bin/bash -i
  fi

  # Boot the real McCoy
  log INFO "Switching root"
  exec switch_root ${fakeroot} /sbin/init

  # This will be reached if previous command failed
  log ERROR "There was an error performing switch_root to ${fakeroot}."
  log ERROR "Starting recovery shell."
  /bin/bash -i
}

main ${@}