#!/bin/bash
# -*- mode: shell-script; indent-tabs-mode: nil; sh-basic-offset: 4; -*-
# ex: ts=8 sw=4 sts=4 et filetype=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