#!/bin/bash

# Copyright 2021  Stuart Winter, Donostia, Spain.
# All rights reserved.
#
# Redistribution and use of this script, with or without modification, is
# permitted provided that the following conditions are met:
#
# 1. Redistributions of this script must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
#  THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED
#  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO
#  EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
#  OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
#  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
#  OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
#  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###################################################################################
# Program: /usr/sbin/os-initrd-mgr
# Purpose: Manage the Slackware ARM / AArch64 OS InitRD.
# Author : Stuart Winter <mozes@slackware.com>
# Date...: 20-Sep-2021
# Version: 1.02
###################################################################################
# Changelog:
# v1.02 - 20-Sep-2021
# * Support /etc/mdadm.conf
# v1.01 - 18-Aug-2021
# * Added -F/--force-rebuild to enable easy removal of local customisations.
#   If the Slackware ARM Kernel Module Loader local customisation files aren't
#   present within /boot/local, remove them from the OS InitRD.
# v1.00 - 20-May-2021
###################################################################################
# Todo:
# Automatically detect whether the mdadm.conf and other customisations should be
# updated/removed using a file hash.
# Perhaps /boot/.customisation_hashes
# For mdadm.conf it require extra handling though - don't include it if the file
# is 100% comments.
# Logic:
# If mdadm.conf is within the OS InitRD but is absent within /etc, or copy in
# /etc/ is pure comments:
# remove from OS InitRD, remove entry from hash file.
#
# If mdadm.conf present within OS InitRD and /etc/, but has different hash:
# re-include version from /etc, update entry in hash file.
#
# If mdadm.conf absent within OS InitRD, but present in /etc and isn't pure comments:
# include within OS InitRD and add entry to hash file.
#
# At present /etc/mdadm.conf is only re-included or removed if local customisations
# are found, or if -F was supplied.
# This means that any /etc/mdadm.conf may be skipped during the Kernel package
# upgrade process.
# However, this is low priority since init tries to set up Software RAID by itself.
###################################################################################
# Version
PROGNAME=os-initrd-mgr
VERSION="${PROGNAME} v1.02 by Stuart Winter <mozes@slackware.com>"
# Settings:
BOOTDIR=/boot
INITRDCUSTOMISATIONSDIR=$BOOTDIR/local
INITRDDETAILS=$BOOTDIR/.boot_details
QUIETMODE=No
FORCEREBUILD=No
# Kernel Module Loader customisation files:
KMODLOADCUSTS="load_kernel_modules.pre load_kernel_modules.pre-modload \
               load_kernel_modules.post"
# Possible customisations (configs and scripts) that the user may
# place within /boot/local for reincorporation into the OS InitRD:
CUSTOMISATIONS="rootdev rootfs wait-for-root resumedev luksdev lukstrim \
                lukskey keymap initrd-name \
                $KMODLOADCUSTS"

# Test if getopt is available.  It's not available within the Slackware Installer
# which is fine, since there's no need to run this tool during the initial
# installation of the Kernel package.
[ ! -x /bin/getopt ] && exit 0

# Temporary location into which the OS InitRD will be unpacked:
TMPDIR="$( mktemp -d /tmp/os-initrd-mgr.XXXXXX )"

# Exit status:
# 0  if OK
# 1  any error.  There is no requirement to surface different codes
#                as os-initrd-mgr is not designed to be called by any
#                tool that will handle the effects of this tool in any
#                manner.

############################## Functions###################################

function display_usage () {
   printf "Usage: ${PROGNAME} [options]\n"
   printf "By default, ${PROGNAME} will scan for any local additions\n"
   printf "within $INITRDCUSTOMISATIONSDIR and will reincorporate them within\n"
   printf "the OS Initial RAM Disk."

   if [ ! -z "$1" ]; then
      echo "Use $( basename $0 ) --help for a list of options"
   fi
}

function display_help() {
   printf "${VERSION}\n\n$( display_usage )

Main options:
 -q,   --quiet   Surpress output messages if there are no user
                 customisations within /boot/local.
                 This option is used by the kernel package's
                 post installation script to surpress any messages
                 attached to inaction on the part of this tool, with
                 the intention of only supplying informative messages
                 should any actions need to be taken.
 -F,   --force-rebuild
                 Force rebuilding of the OS InitRD even when no local
                 customisations are found.
 -h,   --help    Output this help text.
 -v,   --version Display version information.
"
}

function cleanup() {
   echo "Cleaning up ..."
   rm -rf $TMPDIR
}

function unpackinitrd() {
   local initrdfile="$1"
   if [ -s $initrdfile ]; then
      cd $TMPDIR
      echo "Unpacking $initrdfile ..."
      zcat $initrdfile | cpio --quiet -di
      return $?
    else
      echo "Error: $initrdfile does not exist or is of zero bytes."
      return 1
   fi
}

# In theory we shouldn't have any issues repacking the initrd
# (at least from a storage capacity perspective) because we're
# unpacking outside of /boot, and we're only adding scripts to the
# initrd, and we're repacking in place.
function packinitrd() {
   local found_customisations=$1
   local initrdfile="$2"
   local bootdetailsfile="$3"
   local bootcustomisationsdir="$4"
   local initrdsize
   local customisation

   cd $TMPDIR
   #printf "Adding customisation script:"
   # There may be no customisations if --force-rebuild was set.
   # Handle all customisations - Kernel Module Loader and initrd config files:
   if [ $found_customisations -eq 1 ]; then
      printf "Adding customisations ...\n"
      for customisation in $CUSTOMISATIONS ; do
         if [ -f $bootcustomisationsdir/$customisation ]; then
            #printf " $customisation"
            install -pm644 $bootcustomisationsdir/$customisation . # config lives in the OS InitRD's root directory.
         fi
      done
   fi

   # Handle the case where there are Kernel Module Loader customisation
   # scripts within the unpacked initrd, but not present within the OS' /boot/local
   # This indicates that the user wants to remove them, and will have supplied
   # the -F option to os-initrd-mgr
   # Our PWD is presently the root of the OS InitRD:
   # TOFIX: Note: we cannot remove the initrd customisations because /sbin/init
   #        performs no sanity checking on the presence of those files.
   #        This requires a patch to be merged upstream, but should be simple.
   for customisation in $KMODLOADCUSTS ; do
      if [ -f $customisation -a ! -f $bootcustomisationsdir/$customisation ]; then
         printf "Removing customisation script from the OS InitRD: $customisation\n"
         rm -f $customisation
      fi
   done

   # Include RAID support in initrd
   # Note: Slackware ARM/AArch64 OS InitRD ships /etc/mdadm.conf by default
   #       because the mkinitrd tool packages it plus the binaries and udev rule.
   if [ -s /etc/mdadm.conf ] ; then
      # By default Slackware ships a fully commented mdadm.conf, so
      # before copying it to the OS InitRD, we'll check if it's a functional
      # config:
      if egrep -qv '^#' /etc/mdadm.conf ; then
         printf "Installing /etc/mdadm.conf\n"
         cp -fa /etc/mdadm.conf etc/
      fi
     else
      # /etc/mdadm.conf is managed within the OS. If that file is absent yet there
      # is a copy within the OS InitRD, this indicates RAID support is no longer
      # required.
      if [ -s etc/mdadm.conf ]; then
         printf "/etc/mdadm.conf no longer present, removing from the OS InitRD\n"
         rm -f etc/mdadm.conf
      fi
   fi

   printf "Repacking $initrdfile ...\n"
   initrdsize=$( du -sb . | awk '{print $1}' )
   # Update the size information:
   sed -i 's?^initrdsize=.*?initrdsize='"$initrdsize"'?g' $bootdetailsfile
   # Pack:
   find . | cpio --quiet -o -H newc | gzip -9f > $initrdfile
   return $?
}

function preallocate() {
  local tmpdir=$1
  local size=$2
  local exitval
  mkdir -p $tmpdir
  fallocate -l $size $tmpdir/testfile-$size > /dev/null 2>&1
  exitval=$?
  rm -f $tmpdir/testfile-$size
  return $exitval
}

###############################################################################

# Trap errors and clean up:
trap cleanup 1 2 14 15 # trap CTRL+C and kill

PARAMS="$( getopt -qn "$( basename $0 )" -o h,q,v,F -l help,quiet,version,force-rebuild -- "$@" )"
# If params are incorrect then
if [ $? -gt 0 ]; then display_help ; exit 2 ; fi
eval set -- "${PARAMS}"
for param in $*; do
  case "$param" in
     -q|--quiet)   QUIETMODE=Yes
                   shift 1;;
     -F|--force-rebuild) FORCEREBUILD=Yes
                   shift 1;;
     -h|--help)    display_help ; exit ;;

     -v|--version) printf "${VERSION}\n" ; exit;;

     --) shift; break;;
  esac
done

# Check if we have the boot details.  We need these to
# determine whether we can safely unpack the OS Init RD:
# I may add this to the scope of the --force-rebuild option here at some point.
if [ ! -s $INITRDDETAILS ]; then
   echo "Error: unable to locate details file $INITRDDETAILS"
   echo "       You may need to reinstall the kernel package."
   exit 1
fi

# Load in the details file:
. $INITRDDETAILS || exit 1

# We'll detect the platform now, but it should also be included within
# /boot/.boot_details
if [ -z "$platform" ]; then
   case "$( uname -m )" in
      arm*)    platform="armv7";;
      aarch64) platform="armv8";;
      i?86)    platform="x86";;
      *)       platform=$( uname -m );;
   esac
fi

# Versionless symlink to the OS InitRD:
# On ARM / AArch64 the file name is /boot/initrd-armv{7,8}
OSINITRDFILE=$BOOTDIR/initrd-$platform

# Check if there are any local customisations:
found_customisations=0

[ "$QUIETMODE" = "No" ] && printf "Searching for local customisations ... "
for customisation in $CUSTOMISATIONS ; do
   if [ -f $INITRDCUSTOMISATIONSDIR/$customisation ]; then
      [ $found_customisations = 0 ] && printf "\n"
      #[ "$QUIETMODE" = "No" ] && printf "Found: $INITRDCUSTOMISATIONSDIR/$customisation\n"
      printf "Found: $INITRDCUSTOMISATIONSDIR/$customisation\n"
      found_customisations=1
   fi
done
if [ $found_customisations -eq 0 -a "$FORCEREBUILD" = "No" ]; then
   [ "$QUIETMODE" = "No" ] && printf " none found, quitting.\n"
   cleanup > /dev/null
   # It's not an error not to have local customisations, it just means
   # that this script has no actions to take by default.
   exit 0
 else
   printf "\n"
fi

# initrdsize is set within the boot details file (the file name is
# set within the $INITRDDETAILS variable).
#
preallocate $TMPDIR $initrdsize
if [ $? -gt 0 ]; then
   echo "Error: Insufficient temporary storage capacity to extract"
   echo "       the OS Init RD ($initrdsize bytes required)"
   cleanup > /dev/null
   exit 1
fi

unpackinitrd $OSINITRDFILE $INITRDDETAILS
if [ $? -gt 0 ]; then
   echo "Error: Unable to successfully unpack the OS InitRD to"
   echo "       temporary storage"
   echo "       See the error messages above and take corrective"
   echo "       action."
   cleanup > /dev/null
   exit 1
fi

packinitrd $found_customisations $OSINITRDFILE $INITRDDETAILS $INITRDCUSTOMISATIONSDIR
if [ $? -gt 0 ]; then
   echo "Warning: The OS InitRD may be corrupt."
   echo "         You may need to ensure that there is adequate storage"
   echo "         ($initrdsize bytes) within $BOOTDIR and reinstall the"
   echo "         kernel package."
   cleanup
   exit 1
 else
   echo "Successfully rebuilt the OS InitRD."
   cleanup > /dev/null
   exit 0
fi