#!/bin/sh

STATEDIR=/var/lib/dracut
BOOTDIR=/boot
CONF=/etc/dracut.conf
KPKGCONF=/etc/kernel-img.conf
USETRIGGERS=true
mode=""
version=""

set -e

if        $USETRIGGERS                                         \
       && [ x"$DPKG_MAINTSCRIPT_PACKAGE" != x ]                \
       && [ $# = 1 ]                                           \
       && [ x"$1" = x-u ]                                      \
       && dpkg-trigger --check-supported 2>/dev/null
then
       if dpkg-trigger --no-await update-initramfs; then
               echo "dracut-update-initramfs: deferring update (trigger activated)"
               exit 0
       fi
fi

usage()
{
       if [ -n "${1}" ]; then
               printf "${@}\n\n" >&2
       fi
       cat >&2 << EOF
Usage: ${0} [OPTION]...

Options:
 -k [version]  Specify kernel version or 'all'
 -c            Create a new initramfs
 -u            Update an existing initramfs
 -d            Remove an existing initramfs
 -t            Take over a custom initramfs with this one
 -b            Set alternate boot directory
 -v            Be verbose
 -h            This message

EOF
       exit 1
}

# chroot check
chrooted()
{
        # borrowed from udev's postinst
        if [ "$(stat -c %d/%i /)" = "$(stat -Lc %d/%i /proc/1/root 2>/dev/null)" ]; then
                # the devicenumber/inode pair of / is the same as that of
                # /sbin/init's root, so we're *not* in a chroot and hence
                # return false.
                return 1
        fi
        return 0
}

mild_panic()
{
       if [ -n "${1}" ]; then
               printf "${@}\n" >&2
       fi
       exit 0
}

panic()
{
       if [ -n "${1}" ]; then
               printf "${@}\n" >&2
       fi
       exit 1
}

verbose()
{
       if [ "${verbose}" = 1 ]; then
               printf "${@}\n"
       fi
}

version_exists()
{
       [ -e "${STATEDIR}/${1}" ] && [ -e "${initramfs}" ]
       return $?
}

set_initramfs()
{
       initramfs="${BOOTDIR}/dracut.img-${version}"
}


# backup initramfs while running
backup_initramfs()
{
       [ ! -r "${initramfs}" ] && return 0
       initramfs_bak="${initramfs}.dpkg-bak"
       [ -r "${initramfs_bak}" ] && rm -f "${initramfs_bak}"
       ln -f "${initramfs}" "${initramfs_bak}" \
               || cp -a "${initramfs}" "${initramfs_bak}"
       verbose "Keeping ${initramfs_bak}"
}

# keep booted initramfs
backup_booted_initramfs()
{
       initramfs_bak="${initramfs}.dpkg-bak"

       # first time run thus no backup
       [ ! -r "${initramfs_bak}" ] && return 0

       # chroot with no /proc
       [ ! -r /proc/uptime ] && rm -f "${initramfs_bak}" && return 0

       # no kept backup wanted
       [ "${backup_initramfs}" = "no" ] && rm -f "${initramfs_bak}" && return 0

       # no backup yet
       if [ ! -r "${initramfs}.bak" ]; then
               mv -f ${initramfs_bak} "${initramfs}.bak"
               verbose "Backup ${initramfs}.bak"
               return 0
       fi

       # keep booted initramfs
       uptime_days=$(awk '{printf "%d", $1 / 3600 / 24}' /proc/uptime)
       if [ -n "$uptime_days" ]; then
               boot_initramfs=$(find "${initramfs}.bak" -mtime +${uptime_days})
       fi
       if [ -n "${boot_initramfs}" ]; then
               mv -f "${initramfs_bak}" "${initramfs}.bak"
               verbose "Backup ${initramfs}.bak"
               return 0
       fi
       verbose "Removing current backup ${initramfs_bak}"
       rm -f ${initramfs_bak}
}

# nuke generated copy
remove_initramfs()
{
       [ -z "${initramfs_bak}" ] && return 0
       rm -f "${initramfs_bak}"
       verbose "Removing ${initramfs_bak}"
}


generate_initramfs()
{
       echo "dracut-update-initramfs: Generating ${initramfs}"
       OPTS=""
       if [ "${verbose}" = 1 ]; then
               OPTS="-v ${OPTS}"
       fi
       ##WORK HERE!
       if dracut ${OPTS} "${initramfs}.new" "${version}"; then
               mv -f "${initramfs}.new" "${initramfs}"
               set_sha1
       else
               mkinitramfs_return="$?"
               remove_initramfs
               rm -f "${initramfs}.new"
               if [ "$mkinitramfs_return" = "2" ]; then
                       # minversion wasn't met, exit 0
                       exit 0
               fi
               echo "update-initramfs: failed for ${initramfs}"
               exit $mkinitramfs_return
       fi
}

# lilo call
run_lilo()
{
       # show lilo errors on failure
       if ! lilo -t  > /dev/null 2>&1 ; then
               echo "ERROR lilo fails for new ${initramfs}:"
               echo
               lilo -t
       fi
       lilo
}

# check if lilo is on mbr
mbr_check()
{
       # try to discover grub|grub2 and be happy
       [ -r /boot/grub/grub.cfg ] \
               && groot=$(awk '/^set root=/{print substr($2, 7, 3); exit}' \
                       /boot/grub/grub.cfg)
       [ -r /boot/grub/menu.lst ] \
               && groot=$(awk '/^root/{print substr($2, 2, 3); exit}' \
                       /boot/grub/menu.lst)
       [ -e /boot/grub/device.map ] && [ -n "${groot}" ] \
               && dev=$(awk "/${groot}/{ print \$NF}" /boot/grub/device.map)
       [ -n "${dev}" ] && [ -r ${dev} ] \
               && dd if="${dev}" bs=512 skip=0 count=1 2> /dev/null \
               | grep -q GRUB && return 0

       # check out lilo.conf for validity
       boot=$(awk -F = '/^boot=/{ print $2}' /etc/lilo.conf)
       [ -z "${boot}" ] && return 0
       case ${boot} in
       /dev/md/*)
               if [ -r /proc/mdstat ]; then
                       MD=${boot#/dev/md/}
                       boot="/dev/$(awk "/^md${MD}/{print substr(\$5, 1, 3)}" \
                       /proc/mdstat)"
               fi
               ;;
       /dev/md*)
               if [ -r /proc/mdstat ]; then
                       MD=${boot#/dev/}
                       boot="/dev/$(awk "/^${MD}/{print substr(\$5, 1, 3)}" \
                       /proc/mdstat)"
               fi
               ;;
       esac
       [ ! -r "${boot}" ] && return 0
       dd if="${boot}" bs=512 skip=0 count=1 2> /dev/null | grep -q LILO \
               && run_lilo && return 0

       # no idea which bootloader is used
       echo
       echo "WARNING: grub and lilo installed."
       echo "If you use grub as bootloader everything is fine."
       echo "If you use lilo as bootloader you must run lilo!"
       echo
}

# Invoke bootloader
run_bootloader()
{
       # if both lilo and grub around, figure out if lilo needs to be run
       if ( command -v update-grub >/dev/null 2>&1 \
               || [ -e /boot/grub/menu.lst ] || [ -e /boot/grub/grub.cfg ] ) \
       && ( [ -e /etc/lilo.conf ] && command -v lilo >/dev/null 2>&1 ); then
               [ -r "${KPKGCONF}" ] && \
               do_b=$(awk  '/^do_bootloader/{print $3}' "${KPKGCONF}")
               if [ "${do_b}" = "yes" ] || [ "${do_b}" = "Yes" ] \
               || [ "${do_b}" = "YES" ]; then
                       run_lilo
                       return 0
               elif [ "${do_b}" = "no" ] || [ "${do_b}" = "No" ] \
               || [ "${do_b}" = "NO" ]; then
                       return 0
               fi

               # do_bootloader unconfigured
               mbr_check
               return 0
       fi
       if [ -r /etc/lilo.conf ] && command -v lilo >/dev/null 2>&1; then
               run_lilo
               return 0
       fi
       if command -v elilo >/dev/null 2>&1; then
               elilo
               return 0
       fi
       if [ -r /etc/zipl.conf ]; then
               zipl
       fi
       if flash-kernel --supported >/dev/null 2>&1; then
               flash-kernel ${version}
       fi
}

compare_sha1()
{
       sha1sum "${initramfs}" | diff "${STATEDIR}/${version}" - >/dev/null 2>&1
       return $?
}

# Note that this must overwrite so that updates work.
set_sha1()
{
       sha1sum "${initramfs}" > "${STATEDIR}/${version}"
}

delete_sha1()
{
       rm -f "${STATEDIR}/${version}"
}

# ro /boot is not modified
ro_boot_check()
{
       # check irrelevant inside of a chroot
       if [ ! -r /proc/mounts ] || chrooted; then
               return 0
       fi

       boot_opts=$(awk '/boot/{if ((match($4, /^ro/) || match($4, /,ro/)) \
               && $2 == "/boot") print "ro"}' /proc/mounts)
       if [ -n "${boot_opts}" ]; then
               echo "WARNING: /boot is ro mounted."
               echo "update-initramfs: Not updating ${initramfs}"
               exit 0
       fi
}

get_sorted_versions()
{
       version_list=""

       for gsv_x in "${STATEDIR}"/*; do
               gsv_x="$(basename "${gsv_x}")"
               if [ "${gsv_x}" = '*' ]; then
                       return 0
               fi
               worklist=""
               for gsv_i in $version_list; do
                       if dpkg --compare-versions "${gsv_x}" '>' "${gsv_i}"; then
                               worklist="${worklist} ${gsv_x} ${gsv_i}"
                               gsv_x=""
                       else
                               worklist="${worklist} ${gsv_i}"
                       fi
               done
               if [ "${gsv_x}" != "" ]; then
                       worklist="${worklist} ${gsv_x}"
               fi
               version_list="${worklist}"
       done

       verbose "Available versions: ${version_list}"
}

set_current_version()
{
       if [ -f /boot/dracut.img-`uname -r` ]; then
               version=`uname -r`
       fi
}

set_linked_version()
{
       if [ -e /initrd.img ] && [ -L /initrd.img ]; then
               linktarget="$(basename "$(readlink /initrd.img)")"
       fi

       if [ -e /boot/initrd.img ] && [ -L /boot/initrd.img ]; then
               linktarget="$(basename "$(readlink /boot/initrd.img)")"
       fi

       if [ -z "${linktarget}" ]; then
               return
       fi

       version="${linktarget##*img-}"
}

set_highest_version()
{
       get_sorted_versions
       set -- ${version_list}
       version=${1}
}

create()
{
       if [ -z "${version}" ]; then
               usage "Create mode requires a version argument"
       fi

       set_initramfs

       if [ "${takeover}" = 0 ]; then
               if version_exists "${version}"; then
                       panic "Cannot create version ${version}: already exists"
               fi

               if [ -e "${initramfs}" ]; then
                       panic "${initramfs} already exists, cannot create."
               fi
       fi

       generate_initramfs
}

update()
{
       if [ "${update_initramfs}" = "no" ]; then
               echo "update-initramfs: Not updating initramfs."
               exit 0
       fi

       if [ -z "${version}" ]; then
               set_highest_version
       fi

       if [ -z "${version}" ]; then
               set_linked_version
       fi

       if [ -z "${version}" ]; then
               set_current_version
       fi

       if [ -z "${version}" ]; then
               verbose "Nothing to do, exiting."
               exit 0
       fi

       set_initramfs

       ro_boot_check

       altered_check

       backup_initramfs

       generate_initramfs

       run_bootloader

       backup_booted_initramfs
}

delete()
{
       if [ -z "${version}" ]; then
               usage "Delete mode requires a version argument"
       fi

       set_initramfs

       if [ ! -e "${initramfs}" ]; then
               panic "Cannot delete ${initramfs}, doesn't exist."
       fi

       if ! version_exists "${version}"; then
               panic "Cannot delete version ${version}: Not created by this utility."
       fi

       altered_check

       echo "update-initramfs: Deleting ${initramfs}"

       delete_sha1

       rm -f "${initramfs}"
}

# Check for update mode on existing and modified initramfs
altered_check()
{
       # No check on takeover
       [ "${takeover}" = 1 ] && return 0
       if [ ! -e "${initramfs}" ]; then
               mild_panic "${initramfs} does not exist. Cannot update."
       fi
       if ! compare_sha1; then
               echo "update-initramfs: ${initramfs} has been altered." >&2
               mild_panic "update-initramfs: Cannot update. Override with -t option."
       fi
}

# Defaults
verbose=0
yes=0
# We default to takeover=1 in Ubuntu, but not Debian
takeover=0

##

while getopts "k:cudyvtb:h?" flag; do
       case "${flag}" in
       k)
               version="${OPTARG}"
               ;;
       c)
               mode="c"
               ;;
       d)
               mode="d"
               ;;
       u)
               mode="u"
               ;;
       v)
               verbose="1"
               ;;
       y)
               yes="1"
               ;;
       t)
               takeover="1"
               ;;
       b)
               BOOTDIR="${OPTARG}"
               if [ ! -d "${BOOTDIR}" ]; then
                       echo "Error: ${BOOTDIR} is not a directory."
                       exit 1
               fi
               ;;
       h|?)
               usage
               ;;
       esac
done

shift $((${OPTIND} - 1))

if [ $# -ne 0 ]; then
       echo "Invalid argument for option -k."
       usage
fi

# Validate arguments
if [ -z "${mode}" ]; then
       usage "You must specify at least one of -c, -u, or -d."
fi

if [ "${version}" = "all" ] \
       || ( [ "${update_initramfs}" = "all" ] && [ -z "${version}" ] ); then
       : FIXME check for --yes, and if not ask are you sure
       get_sorted_versions
       if [ -z "${version_list}" ]; then
               verbose "Nothing to do, exiting."
               exit 0
       fi

       OPTS="-b ${BOOTDIR}"
       if [ "${verbose}" = "1" ]; then
               OPTS="${OPTS} -v"
       fi
       if [ "${takeover}" = "1" ]; then
               OPTS="${OPTS} -t"
       fi
       if [ "${yes}" = "1" ]; then
               OPTS="${OPTS} -y"
       fi
       for u_version in ${version_list}; do
               # Don't stop if one version doesn't work.
               set +e
               verbose "Execute: ${0} -${mode} -k \"${u_version}\" ${OPTS}"
               "${0}" -${mode} -k "${u_version}" ${OPTS}
               set -e
       done
       exit 0
fi


case "${mode}" in
       c)
               create
               ;;
       d)
               delete
               ;;
       u)
               update
               ;;
esac