#!/bin/sh

set -e

. /usr/share/proxmox-kernel-helper/scripts/functions

_add_entry_to_list_file() {
	file="$1"
	entry="$2"

	if [ -e "$file" ]; then
		cp "$file" "$file.new"
	fi
	echo "$entry" >> "$file.new"
	sort -uo "$file.new" "$file.new"
	mv "$file.new" "$file"
}

_remove_entry_from_list_file() {
	file="$1"
	entry="$2"

	# guard against removing whole file by accident!
	if [ -z "$entry" ]; then
		echo "cannot remove empty entry from '$file'."
		return
	fi

	if [ -e "$file" ]; then
		grep -vFx -- "$entry" "$file" > "$file.new" || true
		mv "$file.new" "$file"
	else
		echo "'$file' does not exist.."
	fi
}

_get_partition_info() {
	if [ ! -e "$1" ]; then
		warn "E: '$1' does not exist!"
		exit 1
	fi
	bdev=$(realpath "$1")
	if [ ! -b "$bdev" ]; then
		warn "E: '$bdev' is not a block device!"
		exit 1
	fi

	bdev_info=$( \
	          lsblk \
	          --bytes \
	          --pairs \
	          -o 'UUID,SIZE,FSTYPE,PARTTYPE,PKNAME,MOUNTPOINT' \
	          "$bdev" \
	          )
	if [ -z "$bdev_info" ]; then
		warn "E: unable to get information about block device '$1'!"
		exit 1
	fi

	count=$(echo "$bdev_info" | grep -c '^')
	if [ "$count" -ne '1' ]; then
		echo "$bdev_info"
		warn "E: block device '$1' has children!"
		exit 1
	fi

	echo "$bdev_info"
	eval "$bdev_info"

	if [ -z "$PKNAME" ]; then
		warn "E: cannot determine parent device of '$1' - please provide a partition, not a full disk."
		exit 1
	fi

	if [ -n "$SIZE" ] && [ "$SIZE" -lt 268435456 ]; then
		warn "E: '$1' is too small (<256M)."
		exit 1
	fi

	if [ -n "$MOUNTPOINT" ]; then
		warn "E: '$1' is mounted on '$MOUNTPOINT' - exiting."
		exit 1
	fi
}

format() {
	part="$1"
	force="$2"

	_get_partition_info "$part"

	if [ -n "$FSTYPE" ]; then
		if [ -z "$force" ] || [ "$force" != '--force' ]; then
			warn "E: '$part' contains a filesystem ('$FSTYPE') - exiting (use --force to override)"
			exit 1
		fi
	fi

	part_basename=$(basename "$bdev")
	if [ -z "$part_basename" ]; then
		if [ "$part" != "$bdev" ]; then
		    symlinkmsg=" -> '$bdev'"
		fi
		warn "E: unable to determine basename of '$part'$symlinkmsg"
		exit 1
	fi

	part_num=$(cat /sys/block/"$PKNAME"/"$part_basename"/partition)
	if [ -z "$part_num" ]; then
		warn "E: unable to determine partition number of '$part'"
		exit 1
	fi

	if [ -z "$PARTTYPE" ] || [ "$PARTTYPE" != "$ESPTYPE" ]; then
		echo "Setting partition type of '$part' to '$ESPTYPE'.."
		sgdisk "-t$part_num:$ESPTYPE" "/dev/$PKNAME"
		echo "Calling 'udevadm settle'.."
		udevadm settle --timeout=5
	fi

	echo "Formatting '$part' as vfat.."
	mkfs.vfat -F 32 "$part"
	echo "Done."
	exit 0
}

init_bootloader() {
	part="$1"
	mode="$2"

	_get_partition_info "$part"

	if [ -z "$PARTTYPE" ] || [ "$PARTTYPE" != "$ESPTYPE" ]; then
		warn "E: '$part' has wrong partition type (!= $ESPTYPE)."
		exit 1
	fi

	if [ -z "$FSTYPE" ] || [ "$FSTYPE" != 'vfat' ]; then
		warn "E: '$part' has wrong filesystem (!= vfat)."
		exit 1
	fi

	if [ -z "$UUID" ]; then
		warn "E: '$part' has no UUID set, required for mounting."
		exit 1
	fi

	esp_mp="/var/tmp/espmounts/$UUID"

	mkdir -p "$esp_mp"
	echo "Mounting '$part' on '$esp_mp'."
	mount -t vfat -o umask=0077 "$part" "$esp_mp"

	if [ -z "$mode" ]; then
		if [ -e "$esp_mp/EFI/proxmox/grubx64.efi" ]; then
			# for reinit support
			mode=grub
		fi
	fi

	if [ -d /sys/firmware/efi ]; then
		if [ -n "$mode" ] && [ "$mode" = "grub" ]; then
			echo "Installing grub x86_64 target.."
			grub-install.real \
				--boot-directory "${esp_mp}" \
				--target x86_64-efi \
				--no-floppy \
				--efi-directory "${esp_mp}" \
				--bootloader-id 'proxmox' \
				"/dev/$PKNAME"
			echo "Installing grub x86_64 target (removable).."
			grub-install.real \
				--boot-directory "${esp_mp}" \
				--target x86_64-efi \
				--no-floppy \
				--efi-directory "${esp_mp}" \
				--removable \
				--bootloader-id 'proxmox' \
				"/dev/$PKNAME"
		else
			echo "Installing systemd-boot.."
			mkdir -p "$esp_mp/$PMX_ESP_DIR"
			if ! command -V bootctl >/dev/null 2>&1 ;
			then
				warn "E: bootctl is not available - make sure systemd-boot is installed"
				exit 1
			fi

			bootctl --graceful --path "$esp_mp" install

			echo "Configuring systemd-boot.."
			echo "timeout 3" > "$esp_mp/$PMX_LOADER_CONF.tmp"
			echo "default proxmox-*" >> "$esp_mp/$PMX_LOADER_CONF.tmp"
			mv "$esp_mp/$PMX_LOADER_CONF.tmp" "$esp_mp/$PMX_LOADER_CONF"
		fi
	else
		echo "Installing grub i386-pc target.."
		grub-install.real \
			--boot-directory "$esp_mp" \
			--target i386-pc \
			--no-floppy \
			--bootloader-id='proxmox' \
			"/dev/$PKNAME"
	fi
	echo "Unmounting '$part'."
	umount "$part"

	echo "Adding '$part' to list of synced ESPs.."
	_add_entry_to_list_file "$ESP_LIST" "$UUID"

}

reinit() {
	if ! (echo "${curr_uuid}" | grep -qE '[0-9a-fA-F]{4}-[0-9a-fA-F]{4}'); then
		warn "WARN: ${curr_uuid} read from ${ESP_LIST} does not look like a VFAT-UUID - skipping"
		return
	fi

	path="/dev/disk/by-uuid/$curr_uuid"
	if [ ! -e "${path}" ]; then
		warn "WARN: ${path} does not exist - clean '${ESP_LIST}'! - skipping"
		return
	fi
	init_bootloader "$path" "$mode"
}

_clean_impl() {
	if [ ! -e "/dev/disk/by-uuid/" ]; then
		warn 'E: /dev/disk/by-uuid does not exist, aborting!'
		exit 1
	fi
	printf "Checking whether ESP '%s' exists.. " "$curr_uuid" # avoid newline
	if [ -e "/dev/disk/by-uuid/$curr_uuid" ]; then
		echo "Found!"
	else
		echo "Not found!"
		if [ -z "$dry_run" ] || [ "$dry_run" != '--dry-run' ]; then
			_remove_entry_from_list_file "$ESP_LIST" "$curr_uuid"
		fi
	fi
}

clean() {
	dry_run="$1"
	rm -f "$ESP_LIST".tmp
	loop_esp_list _clean_impl
	if [ "$?" -eq 2 ]; then
		warn "E: $ESP_LIST does not exist."
		exit 1
	fi
	if [ -e "$ESP_LIST".tmp ]; then
		mv "$ESP_LIST".tmp "$ESP_LIST"
	fi

	echo "Sorting and removing duplicate ESPs.."
	sort -uo "$ESP_LIST".tmp "$ESP_LIST"
	mv "$ESP_LIST".tmp "$ESP_LIST"
}

refresh() {
	hook=$1
	hookscripts='proxmox-auto-removal zz-proxmox-boot'

	if [ -n "$hook" ]; then
	    if echo "$hookscripts" | grep -sqE "(^|[[:space:]]+)$hook([[:space:]]+|$)"; then
		hookscripts="$hook"
	    else
		warn "E: '$hook' is not a valid hook script name.";
		exit 1;
	    fi
	fi

	for script in $hookscripts; do
		scriptpath="/etc/kernel/postinst.d/$script"
		if [ -f "$scriptpath" ] && [ -x "$scriptpath" ]; then
			echo "Running hook script '$script'.."
			$scriptpath
		else
			warn "Hook script '$script' not found or not executable, skipping."
		fi
	done
}

add_kernel() {
	ver="$1"

	if [ -z "$ver" ]; then
		warn "E: <kernel-version> is mandatory"
		warn ""
		exit 1
	fi

	if [ ! -e "/boot/vmlinuz-$ver" ]; then
		warn "E: no kernel image found in /boot for '$ver', not adding."
		exit 1
	fi
	_add_entry_to_list_file "$MANUAL_KERNEL_LIST" "$ver"
	echo "Added kernel '$ver' to manual kernel list. Use the 'refresh' command to update the ESPs."
}

remove_kernel() {
	ver="$1"

	if [ -z "$ver" ]; then
		warn "E: <kernel-version> is mandatory"
		warn ""
		exit 1
	fi

	if grep -sqFx -- "$ver" "$MANUAL_KERNEL_LIST"; then
		_remove_entry_from_list_file "$MANUAL_KERNEL_LIST" "$ver"
		echo "Removed kernel '$ver' from manual kernel list. Use the 'refresh' command to update the ESPs."
	else
		echo "Kernel '$ver' not found in manual kernel list."
	fi
}

list_kernels() {
	boot_kernels="$(boot_kernel_list)"

	if [ -e "$MANUAL_KERNEL_LIST" ]; then
		manual_kernels="$(cat "$MANUAL_KERNEL_LIST" || true)"
		boot_kernels="$(echo "$boot_kernels" | grep -Fxv -f "$MANUAL_KERNEL_LIST" || true)"
	fi

	if [ -z "$manual_kernels" ]; then
		manual_kernels="None."
	fi

	echo "Manually selected kernels:"
	echo "$manual_kernels"
	echo ""
	echo "Automatically selected kernels:"
	echo "$boot_kernels"

	pinned_kernel="$(get_first_line "$PINNED_KERNEL_CONF")"
	nextboot_kernel="$(get_first_line "$NEXT_BOOT_PIN")"
	if [ -n "$pinned_kernel" ]; then
		echo ""
		echo "Pinned kernel:"
		echo "${pinned_kernel}"
	fi
	if [ -n "$nextboot_kernel" ]; then
		echo ""
		echo "Kernel pinned on next-boot:"
		echo "${nextboot_kernel}"
	fi
}

usage() {
	subcmd="$1"
	if [ -z "$subcmd" ]; then
		warn "USAGE: $0 <commands> [ARGS]"
		warn ""
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "format" ]; then
		warn "  $0 format <partition> [--force]"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "init" ]; then
		warn "  $0 init <partition> [grub]"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "reinit" ]; then
		warn "  $0 reinit"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "clean" ]; then
		warn "  $0 clean [--dry-run]"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "refresh" ]; then
		warn "  $0 refresh [--hook <name>]"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "kernel" ]; then
		warn "  $0 kernel <add|remove> <kernel-version>"
		warn "  $0 kernel pin <kernel-version> [--next-boot]"
		warn "  $0 kernel unpin [--next-boot]"
		warn "  $0 kernel list"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "status" ]; then
		warn "  $0 status [--quiet]"
	fi
	if [ -z "$subcmd" ] || [ "$subcmd" = "help" ]; then
		warn "  $0 help"
	fi
}

help() {
	echo "USAGE: $0 format <partition> [--force]"
	echo ""
	echo "    format <partition> as EFI system partition. Use --force to format even if <partition> is currently in use."
	echo ""
	echo "USAGE: $0 init <partition>"
	echo ""
	echo "    initialize EFI system partition at <partition> for automatic synchronization of Proxmox kernels and their associated initrds."
	echo ""
	echo "USAGE: $0 reinit"
	echo ""
	echo "    reinitialize all configured EFI system partitions from $ESP_LIST."
	echo ""
	echo "USAGE: $0 clean [--dry-run]"
	echo ""
	echo "    remove no longer existing EFI system partition UUIDs from $ESP_LIST. Use --dry-run to only print outdated entries instead of removing them."
	echo ""
	echo "USAGE: $0 refresh [--hook <name>]"
	echo ""
	echo "    refresh all configured EFI system partitions. Use --hook to only run the specified hook, omit to run all."
	echo ""
	echo "USAGE: $0 kernel <add|remove> <kernel-version>"
	echo ""
	echo "    add/remove proxmox-kernel with ABI <kernel-version> to list of synced kernels, in addition to automatically selected ones."
	echo "    NOTE: you need to manually run 'refresh' once you're finished with adding/removing kernels from the list"
	echo ""
	echo "USAGE: $0 kernel pin <kernel-version> [--next-boot]"
	echo ""
	echo "    pin proxmox-kernel with ABI <kernel-version> as the default entry to be booted."
	echo "    with --next-boot sets <kernel-version> only for the next boot."
	echo "    NOTE: you need to manually run 'refresh' once you're finished with pinning kernels"
	echo ""
	echo "USAGE: $0 kernel unpin [--next-boot]"
	echo ""
	echo "    unpin removes pinned and next-boot kernel settings."
	echo "    with --next-boot only removes the pin for the next boot."
	echo ""
	echo "USAGE: $0 kernel list"
	echo ""
	echo "    list kernel versions currently selected for inclusion on ESPs."
	echo ""
	echo "USAGE: $0 status [--quiet]"
	echo ""
	echo "    Print details about the ESPs configuration. Exits with 0 if any ESP is configured, else with 2."
	echo ""
}

_status_detail() {
	if ! (echo "${curr_uuid}" | grep -qE '[0-9a-fA-F]{4}-[0-9a-fA-F]{4}'); then
		warn "WARN: ${curr_uuid} read from ${ESP_LIST} does not look like a VFAT-UUID - skipping"
		return
	fi

	path="/dev/disk/by-uuid/$curr_uuid"
	if [ ! -e "${path}" ]; then
		warn "WARN: ${path} does not exist - clean '${ESP_LIST}'! - skipping"
		return
	fi

	mountpoint="${MOUNTROOT}/${curr_uuid}"
	mkdir -p "${mountpoint}" || \
		{ warn "creation of mountpoint ${mountpoint} failed - skipping"; return; }
	mount "${path}" "${mountpoint}" || \
		{ warn "mount of ${path} failed - skipping"; return; }

	result=""
	if [ -f "${mountpoint}/$PMX_LOADER_CONF" ]; then
		if [ ! -d "${mountpoint}/$PMX_ESP_DIR" ]; then
			warn "${path}/$PMX_ESP_DIR does not exist"
		fi
		versions_uefi=$(ls -1 ${mountpoint}/$PMX_ESP_DIR | awk '{printf (NR>1?", ":"") $0}')
		result="uefi (versions: ${versions_uefi})"
	fi
	if [ -d "${mountpoint}/grub" ]; then
		versions_grub=$(ls -1 ${mountpoint}/vmlinuz-* | awk '{ gsub(/.*\/vmlinuz-/, ""); printf (NR>1?", ":"") $0 }')
		if [ -n "$result" ]; then
		    result="${result}, grub (versions: ${versions_grub})"
		else
		    result="grub (versions: ${versions_grub})"
		fi
	fi
	echo "$curr_uuid is configured with: $result"
	umount "${mountpoint}" || \
		{ warn "umount of ${path} failed - failure"; exit 0; }

	rmdir "${mountpoint}" || true
}

status() {
	quiet="$1"
	if [ ! -e "${ESP_LIST}" ]; then
		if [ -z "$quiet" ]; then
		    warn "E: $ESP_LIST does not exist."
		fi
		exit 2
	fi
	if [ -z "$quiet" ]; then
		if [ -d /sys/firmware/efi ]; then
		    echo "System currently booted with uefi"
		else
		    echo "System currently booted with legacy bios"
		fi
		loop_esp_list _status_detail
	fi
}

_ask_interactive_refresh() {
	msg="$1"

	if [ -t 0 ] && [ -t 1 ]; then # check if interactive
		echo "$msg."
		printf "Refresh the actual boot ESPs now? [yN] "
		read -r do_refresh
		if [ "$do_refresh" != "${do_refresh#[Yy]}" ] ;then
			refresh
		else
			echo "Skip auto-refresh, you can call it any time to enact boot changes."
		fi
	else
		echo "$msg. Use the 'refresh' command to update the ESPs."
	fi
}

pin_kernel() {
	ver="$1"
	pin_file="$2"

	if [ -z "$ver" ]; then
		boot_kernels="$(boot_kernel_list)"
		warn "E: <kernel-version> is mandatory"
		warn ""
		warn "Possible Proxmox kernel versions are:"
		warn "$boot_kernels"
		exit 1
	fi

	if [ -z "$pin_file" ]; then
	    pin_file="$PINNED_KERNEL_CONF"
	fi

	if [ ! -e "/boot/vmlinuz-$ver" ]; then
		boot_kernels="$(boot_kernel_list)"
		warn "E: no kernel image found in /boot for '$ver', not setting default."
		warn ""
		warn "Possible Proxmox kernel versions are:"
		warn "$boot_kernels"
		exit 1
	fi

	if [ -e "$pin_file" ]; then
		old_pin=$(get_first_line "${pin_file}")
		if [ "$ver" != "$old_pin" ]; then
			echo "Overriding previously pinned version '$old_pin' with '$ver'"
		fi
	fi
	echo "$ver" > "$pin_file"

	if [ -f "${ESP_LIST}" ]; then
		_ask_interactive_refresh "Set kernel '$ver' in $pin_file"
	else
		next_boot_ver=$(get_first_line "${NEXT_BOOT_PIN}")
		pin_ver="${next_boot_ver:-$ver}"
		echo "Setting '$pin_ver' as grub default entry and running update-grub."
		set_grub_default "$pin_ver"
		update-grub
	fi
}

unpin_kernel() {
	last_pin=$(get_first_line "${NEXT_BOOT_PIN}")
	rm -f "$NEXT_BOOT_PIN"
	echo "Removed $NEXT_BOOT_PIN."
	if [ -z "$1" ]; then
		old_pin=$(get_first_line "${PINNED_KERNEL_CONF}")
		last_pin=${old_pin:-$last_pin}
		rm -f "$PINNED_KERNEL_CONF"
		echo "Removed $PINNED_KERNEL_CONF."
	fi

	if [ -f "${ESP_LIST}" ]; then
		if [ -n "$last_pin" ]; then
			_ask_interactive_refresh "Unpinned kernel '$last_pin'"
		fi
	else
		echo "Reset default grub entry and running update-grub."
		pinned_kernel=$(get_first_line "${PINNED_KERNEL_CONF}")
		set_grub_default "$pinned_kernel"
		update-grub
	fi
}

if [ -z "$1" ]; then
    usage
    exit 0
fi

case "$1" in
	'format')
		shift
		if [ -z "$1"  ]; then
			warn "E: <partition> is mandatory."
			warn ""
			usage "format"
			exit 1
		fi
		format "$@"
		exit 0
	;;
	'init')
		reexec_in_mountns "$@"
		shift
		if [ -z "$1"  ]; then
			warn "E: <partition> is mandatory."
			warn ""
			usage "init"
			exit 1
		fi
		init_bootloader "$@"
		echo "Refreshing kernels and initrds.."
		refresh
		exit 0
	;;
	'reinit')
		reexec_in_mountns "$@"
		shift
		if [ "$#" -eq 1 ]; then
			warn "E: no arguments allowed."
			warn ""
			usage "reinit"
			exit 1
		fi
                loop_esp_list reinit "$@"
		exit 0
	;;
	'clean')
		shift
		clean "$@"
		exit 0
	;;
	'refresh')
		shift
		if [ "$#" -eq 0 ]; then
			refresh
		elif [ "$#" -eq 2 ] && [ "$1" = "--hook" ]; then
			refresh "$2"
		else
			usage "refresh"
			exit 1
		fi
		exit 0
	;;
	'kernel'|'kernels')
		shift
		if [ -z "$1" ]; then
			warn "E: subcommand is mandatory for 'kernel'."
			warn ""
			usage "kernel"
			exit 1
		fi
		cmd="$1"
		case "$cmd" in
			'add')
				add_kernel "$2"
				exit 0
			;;
			'remove')
				remove_kernel "$2"
				exit 0
			;;
			'list')
				list_kernels
				exit 0
			;;
			'pin')
				if [ "$#" -eq 3 ] && [ "$3" = '--next-boot' ]; then
					pin_kernel "$2" "${NEXT_BOOT_PIN}"
					echo "Pinned for next boot only."
				elif [ "$#" -eq 2 ]; then
					pin_kernel "$2"
				else
					usage "kernel"
					exit 1
				fi
				exit 0
			;;
			'unpin')
				if [ "$#" -eq 2 ] && [ "$2" = '--next-boot' ]; then
					unpin_kernel "$2"
				elif [ "$#" -eq 1 ]; then
					unpin_kernel
				else
					usage "kernel"
					exit 1
				fi
				exit 0
			;;
			*)
				warn "E: invalid 'kernel' subcommand '$cmd'."
				warn ""
				usage "kernel"
				exit 1
			;;
		esac
	;;
	'status')
		if [ "$#" -eq 2 ] && [ "$2" = '--quiet' ]; then
			shift
			status "$1"
		elif [ "$#" -eq 1 ]; then
			reexec_in_mountns "$@"
			shift
			status
		else
			usage "status"
			exit 1
		fi
		exit 0
	;;
	'help')
		shift
		help
		exit 0
	;;
	*)
		warn "Invalid/unknown command '$1'."
		warn ""
		usage
		exit 1
	;;
esac

exit 1
