#!/bin/bash
############################################################################
#                                                                          #
# This file is part of the IPFire Firewall.                                #
#                                                                          #
# IPFire 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 2 of the License, or        #
# (at your option) any later version.                                      #
#                                                                          #
# IPFire 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 IPFire; if not, write to the Free Software                    #
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307 USA #
#                                                                          #
# Copyright (C) 2014 IPFire Team <info@ipfire.org>.                        #
#                                                                          #
############################################################################

GRUB_INSTALL_ARGS="--no-floppy --recheck --force"

function find_bootloader_device() {
	local mp
	for mp in /boot /; do
		if find_device "${mp}"; then
			return 0
		fi
	done

	return 1
}

function find_device() {
	local mountpoint="${1}"

	local root
	local dev mp fs flags rest
	while read -r dev mp fs flags rest; do
		# Skip unwanted entries
		[ "${dev}" = "rootfs" ] && continue

		if [ "${mp}" = "${mountpoint}" ] && [ -b "${dev}" ]; then
			root="$(basename "${dev}")"
			break
		fi
	done < /proc/mounts

	# Get the actual device from the partition that holds /
	while [ -n "${root}" ]; do
		if [ -e "/sys/block/${root}" ]; then
			echo "/dev/${root}"
			return 0
		fi

		# Remove last character
		root="${root::-1}"
	done

	return 1
}

function device_is_mdraid() {
	local device="${1}"

	[ -d "/sys/block/${device/\/dev/}/md" ]
}

function mdraid_get_slaves() {
	local device="${1}"

	local slave
	for slave in /sys/block/${device/\/dev/}/slaves/*; do
		echo "/dev/$(basename "${slave}")"
	done 2>/dev/null
}

function grub_update_config() {
	echo "Updating configuration..."

	if ! grub-mkconfig -o /boot/grub/grub.cfg &>/dev/null; then
		echo "Could not update configuration. Aborting." >&2
		return 1
	fi

	return 0
}

function grub_install() {
	local device="${1}"

	echo "Installing GRUB on ${device}..."

	if [ ! -b "${device}" ]; then
		echo "${device} does not exist or is not a block device" >&2
		return 1
	fi

	local arches
	case "$(uname -m)" in
		aarch64)
			arches="arm64-efi"
			;;
		x86_64)
			arches="i386-pc x86_64-efi"
			;;
	esac

	local arch
	for arch in ${arches}; do
		local args="--target=${arch}"

		case "${arch}" in
			*-efi)
				# Skip all EFI architectures if no EFI partition exists
				if [ ! -d "/boot/efi" ]; then
					continue
				fi

				args="${args} --efi-directory=/boot/efi"

				# Don't try to modify the BIOS when we are
				# not running on EFI right now
				if [ ! -d "/sys/firmware/efi" ]; then
					args="${args} --no-nvram"
				fi
				;;
		esac

		local removable
		for removable in "" "--removable"; do
			if ! grub-install ${GRUB_INSTALL_ARGS} ${args} \
					${removable} "${device}" &>/dev/null; then
				echo "Could not install GRUB on ${device}" >&2
				return 1
			fi

			# Do not try to install with --removable for non-efi architectures
			[[ "${arch}" =~ \-efi$ ]] || break
		done
	done

	return 0
}

function main() {
	local device="${1}"

	# Find the root device
	if [ -z "${device}" ]; then
		device="$(find_bootloader_device)"
		if [ -z "${device}" ]; then
			echo "Could not find root device. Aborting." >&2
			return 1
		fi

		echo "Found bootloader device: ${device}"
	fi

	if [ ! -b "${device}" ]; then
		echo "${device} does not exist" >&2
		return 2
	fi

	# Update configuration files
	grub_update_config || return $?

	# Handle mdraid devices
	if device_is_mdraid "${device}"; then
		local slave
		for slave in $(mdraid_get_slaves "${device}"); do
			grub_install "${slave}" || return $?
		done

	# Handle normal block devices
	else
		grub_install "${device}" || return $?
	fi

	return 0
}

# Run main function
main "$@" || exit $?