#!/bin/sh

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

[ -r ${CONF} ] && . ${CONF}

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}/initrd.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 "update-initramfs: Generating ${initramfs}"
	OPTS="-o"
	if [ "${verbose}" = 1 ]; then
		OPTS="-v ${OPTS}"
	fi
	if mkinitramfs ${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 ( [ -x "$(command -v update-grub)" ] || [ -e /boot/grub/menu.lst ] \
	     || [ -e /boot/grub/grub.cfg ] ) \
	&& ( [ -e /etc/lilo.conf ] && [ -x /sbin/lilo ] ); 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 ] && [ -x /sbin/lilo ]; then
		run_lilo
		return 0
	fi
	if [ -x /sbin/elilo ]; then
		elilo
		return 0
	fi
	if [ -r /etc/zipl.conf ]; then
		zipl
	fi
	if flash-kernel --supported >/dev/null 2>&1; then
		flash-kernel
	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
			xbps-uhelper cmpver "${gsv_x}" "${gsv_i}"
			if [ $? -eq 0 ] || [ $? -eq 1 ]; 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/initrd.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##initrd.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