#! /bin/sh

# Only PC architecture
is_a_pc() {
	case "$(uname -m)" in
		x86_64|i[3456]86)
			echo "yes"
			return 0
			;;
		*)
			echo "no"
			return 1
			;;
	esac
}

# Which virtualization we use -- or "(none)"
detect_virt() {
	local detect="/usr/bin/systemd-detect-virt"
	if [ ! -x "$detect" ]; then
		echo "unknown"
		return 1
	fi
	"$detect"
}

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

# Is this "efi", "bios" or "other" ?
get_firmware_kind() {
	if is_a_pc >/dev/null; then
		if [ -d /sys/firmware/efi ]; then
			echo "efi"
		else
			echo "bios"
		fi
	else
		echo "other"
	fi
}

# Read the boot disk from debconf
debconf_get_boot_disks() {
	local grub_kind="$1"
	debconf-show "$grub_kind" | sed -rn \
		-e 's/\s*,\s*/ /g' \
		-e "s,^.*\<$grub_kind/install_devices:\s*,,p" \
		#
}

# Which GRUB package is used ("grub-pc" or "grub-efi")
get_grub_kind() {
	local firmware_kind grub_kind
	firmware_kind="$(get_firmware_kind)"
	case "$firmware_kind" in
		bios) echo 'grub-pc';;
		efi) echo 'grub-efi';;
	esac
}

# Validate that "debconf" installation disk actually exist
configured_disks_exist() {
	local grub_kind configured_disk
	grub_kind="$(get_grub_kind)"
	echo >&2 "I: Checking GRUB configured disks (for '$grub_kind'):"
	for configured_disk in $(debconf_get_boot_disks "$grub_kind"); do
		if test -b "$configured_disk"; then
			echo >&2 "I:    > Configured disk exist: '$configured_disk'"
		else
			echo >&2 "W:    > Missing configured disk: '$configured_disk')"
			return 1
		fi
	done
	return 0
}

# Find the disk that GRUB should use
find_grub_disk() {
	local firmware_kind
	firmware_kind="$(get_firmware_kind)"
	case "$firmware_kind" in
		bios)
			# Legacy BIOS -- the disk where "/boot" partition is mounted
			lsblk --list --paths --fs -o "NAME,MOUNTPOINT" --noheadings | \
				awk '$2 == "/boot" { print $1 }'
			;;
		efi)
			# EFI -- the disk where "ESP" partition exist
			lsblk --list --paths --output "PKNAME,PARTTYPE" | \
				awk '$2 == "c12a7328-f81f-11d2-ba4b-00a0c93ec93b" { print $1 }'
			;;
	esac
}

# Read lines from stdin and output a single long line of itmes
flatten_list() {
	tr -s '\n' ' ' | sed -e 's/ *$//'
}

# Output a list of devices for given "/dev/md*" name
find_md_partnames() {
	local md_disk="$1"
	local md_basename
	md_basename=$(basename "$md_disk")
	local md_prefix="/sys/class/block/$md_basename/md"
	find "$md_prefix" -maxdepth 1 -name 'dev-*' 2>/dev/null | \
		sed "s,$md_prefix/dev-,," | flatten_list
}

# Output basename of the disk that contain the given partition basename
disk_of_part() {
	local partname="$1"
	lsblk --list --nodeps -noheadings -o "PKNAME" "/dev/$partname"
}

# Given a single device name, output a comma-separated list of disks
# This is designed for "/dev/md*" devices
expand_disk_list() {
	local dev_name="$1"
	local disk
	case "$dev_name" in
		/dev/md*)
			for part in $(find_md_partnames "$dev_name"); do
				disk=$(disk_of_part "$part")
				#echo >&2 "D: '$disk' ($part)"
				find "/dev/disk/by-id" -lname "*${disk}" | sort | head -1
			done | sort -u | flatten_list | sed 's/ /, /g'
			;;
		*)
			local firmware_kind="$(get_firmware_kind)"
			if [ "$firmware_kind" = "bios" ]; then
					#Install GRUb on the parent disk (like /dev/sda)
					disk=$(lsblk -no PKNAME "$dev_name")
					#echo >&2 "D: '$disk'"
					find "/dev/disk/by-id" -lname "*${disk}" | sort | head -1
			else
					echo "$dev_name"
			fi
			;;
	esac
}

# Setting "debconf" when running from "dpkg" during pre-config stage
set_via_confmodule() {
	local template="$1"
	local value="$2"

	. /usr/share/debconf/confmodule

	db_get "$template" || true
	if [ "$RET" = "$value" ]; then
		echo >&2 "D: $template is already set. Skipping."
		exit 0
	fi
	echo >&2 "D: Running inside debconf session. Using db_set."
	db_set "$template" "$value"
	db_fset "$template" seen true
}

# Setting "debconf" when running as standalone script
set_via_communicate() {
	local template="$1"
	local value="$2"

	echo >&2 "D: Running standalone. Using debconf-communicate."
	echo "SET $template $value" | debconf-communicate > /dev/null
	echo "FSET $template seen true" | debconf-communicate > /dev/null
}

set_grub_installed_disks() {
	local template_question="grub-pc/install_devices"

	if [ -n "$DEBIAN_HAS_FRONTEND" ]; then
		# Running from "debconf"
		set_via_confmodule "$template_question" "$DISK_LIST"
	else
		# Running standalone
		set_via_communicate "$template_question" "$DISK_LIST"
	fi
}

check_fixing_grub() {
	if [ "$(id -u)" != 0 ]; then
		echo >&2 "E: this script should run as root and not '$(id -u)'"
		exit 0
	fi

	if [ ! is_a_pc > /dev/null ]; then
		echo >&2 "I: Skipping pre-grub (not a PC platform)"
		exit 0
	fi

	if chrooted >/dev/null; then
		echo >&2 "I: Skipping pre-grub (inside a chroot)"
		exit 0
	fi

	virt_type="$(detect_virt)"
	case "$virt_type" in
		lxc)
			echo >&2 "I: Skipping pre-grub (inside an '$virt_type' guest)"
			exit 0
			;;
		*)
			;;
	esac

	if [ ! -x /usr/bin/debconf-show ]; then
		echo >&2 "I: Skipping pre-grub (no 'debconf-show')"
		exit 0
	fi

	if [ "$(get_firmware_kind)" != "bios" ]; then
		echo >&2 "I: Skipping pre-grub (not BIOS firmware)"
		exit 0
	fi

	if configured_disks_exist; then
		echo >&2 "I: Skipping pre-grub (all configured disks exist)"
		exit 0
	fi
}

run_fixing() {
	local grub_disk
	check_fixing_grub
	# Need to find the correct disk ourself
	grub_disk=$(find_grub_disk)
	echo >&2 "I: /boot partition: '$grub_disk'"
	DISK_LIST=$(expand_disk_list "$grub_disk")
	echo >&2 "I: grub installed disks: '$DISK_LIST'"
	# Now fix it
	set_grub_installed_disks
}

usage() {
	echo >&2 "$@"
	echo >&2 "Usage: $0 {fix|check-only|manual <disk...>}"
	exit 1
}

############################################################################# MAIN

[ "$#" -gt 0 ] || usage "E: No action was specified"
action="$1"
shift

case "$action" in
	fix)
		#echo >&2 "DEBUG0: DEBIAN_HAS_FRONTEND='$DEBIAN_HAS_FRONTEND'"
		run_fixing
		;;
	check-only)
		check_fixing_grub
		;;
	manual)
		DISK_LIST="$@"
		set_grub_installed_disks
		;;
	*)
		usage "E: Unkown action '$action'"
		;;
esac
