#!/bin/sh
# Copyright 2022-2023 Felix Freeman <libsys@hacktivista.org>
#
# This file is part of "Hackware boot" and licensed under the terms of the GNU
# General Public License version 2 or (at your option) any later version. You
# should have received a copy of this license along with the software. If not,
# see <https://www.gnu.org/licenses/>.

print_usage () {
	cat <<- EOF
	Usage: $0 <options>

	options:
	 --user                 sets hostname, changes root password and creates user
	 --luks <old-password>  sets a new LUKS password for / and /boot
	 --luks-root <old-pw>   sets a new LUKS password for /
	 --luks-boot <old-pw>   sets a new LUKS password for /boot
	 --grub-uuid            filter boot by UUID
	 --grub-password        set a GRUB password for \`root\`
	 --grub-secure-boot     prevent evil maid attacks with GPG signatures
	 --grub                 enable the three previous --grub options
	 --commit-grub          makes firmware grub.cfg changes permanent
	EOF
	exit 1
}

if [ -z "$1" ]; then
	print_usage
fi

if [ $(id -u) -ne 0 ]; then
	echo 'This program must be executed as root.'
	exit 2
fi

for param in "$@"; do
	case $param in
		--user)
			type useradd 1>/dev/null || requirements="$requirements useradd" ;;
		--luks*)
			type cryptsetup 1>/dev/null || requirements="$requirements cryptsetup" ;;
		--grub*|--commit-grub)
			for executable in flashrom cbfstool; do
				case "$requirements" in
					*$executable*) ;;
					*) type $executable 1>/dev/null || requirements="$requirements $executable"
				esac
			done
		    case $param in
				--grub-uuid)
					type findmnt 1>/dev/null || requirements="$requirements findmnt" ;;
				--grub-password)
					type grub-mkpasswd-pbkdf2 1>/dev/null || requirements="$requirements grub-mkpasswd-pbkdf2" ;;
				--grub-secure-boot)
					type gpg 1>/dev/null || requirements="$requirements gpg" ;;
			esac
			;;
	esac
done

if [ -n "$requirements" ]; then
	echo "This program requires the following missing executables: $requirements"
	exit 3;
fi

case "$*" in
	*--grub*|*--commit-grub*)
		if ! flashrom -p internal 2>/dev/null; then
			echo 'This program must be executed with kernel param iomem=relaxed'
			exit 4
		fi
		;;
	*--luks\ --*|*--luks)
		echo '--luks requires an <old-password>.'
		exit 5
		;;
esac

case "$*" in
	*--commit-grub*)
		if [ $# -gt 1 ]; then
			echo '--commit-grub cannot be executed with other options.'
			exit 6
		fi
		;;
esac

validate_hostname () {
	if [ -z "$1" ]; then
		echo 'please type an answer' >&2
		return 1
	fi
	case $1 in
		*[!-.a-z0-9]*)
			echo 'please enter lowercase letters, numbers, - and . only'
			return 1
	esac
}

set_hostname () {
	until
		printf '> Hostname: '
		read -r HOSTNAME
		validate_hostname "$HOSTNAME"
	do true; done
	echo "$HOSTNAME" > /etc/hostname
	printf "\n127.0.1.1\t$HOSTNAME\n" >> /etc/hosts
}

set_root_password () {
	echo '> Change root password'
	until
		passwd
	do true; done
}

create_user () {
	echo '> Create a user'
	while true; do printf 'Username: '
		read -r USERNAME
		useradd -m -s /bin/bash "$USERNAME" && break
	done

	until
		passwd $USERNAME
	do true; done
}

cryptboot_uuid () {
	sed -nEe 's/cryptboot[[:space:]]+UUID=([^[:space:]]+).*/\1/p' /etc/crypttab
}

request_new_luks_password () {
	stty -echo
	until
		printf "Password: "
		read -r NEW_LUKS_PASSWORD
	do true; done
	until
		printf "\nRepeat password: "
		read -r _NEW_LUKS_PASSWORD
	do true; done
	stty echo
	printf "\n"
	if [ "$NEW_LUKS_PASSWORD" != "$_NEW_LUKS_PASSWORD" ]; then
		echo "Password mismatch, try again."
		request_new_luks_password
	fi
}

set_luks_boot_password () {
	CRYPTBOOT_UUID=$(cryptboot_uuid)
	# Add user-supplied key for /boot.
	echo '> Set a /boot LUKS key'
	request_new_luks_password
	printf '%s\n%s' "$LUKS_PASSWORD" "$NEW_LUKS_PASSWORD" | cryptsetup -v luksChangeKey "/dev/disk/by-uuid/$CRYPTBOOT_UUID"
}

set_luks_root_password () {
	# Debian LVM
	CRYPTROOT_UUID=$(sed -nEe 's/.*_crypt[[:space:]]+UUID=([^[:space:]]+).*/\1/p' /etc/crypttab)
	if [ -z "$CRYPTROOT_UUID" ]; then
		# Parabola BTRFS
		CRYPTROOT_UUID=$(sed -nEe 's/.*cryptdevice=UUID=([^:]+):.*/\1/p' /etc/default/grub)
	fi

	# Add user-supplied key for /.
	echo '> Set a / LUKS key'
	request_new_luks_password
	printf '%s\n%s' "$LUKS_PASSWORD" "$NEW_LUKS_PASSWORD" | cryptsetup -v luksChangeKey "/dev/disk/by-uuid/$CRYPTROOT_UUID"
}

require_FIRMWARE_ROM () {
	if [ -z "$FIRMWARE_ROM" ]; then
		flashrom -p internal -r firmware.rom 1>/dev/null
		FIRMWARE_ROM='firmware.rom'
	fi
}

require_GRUB_CFG () {
	require_FIRMWARE_ROM
	if [ -z "$GRUB_CFG" ]; then
		cbfstool "$FIRMWARE_ROM" extract -n etc/grub.cfg -f grubtest.cfg 2>/dev/null
		GRUB_CFG="grubtest.cfg"
	fi
}

add_grub_cfg () {
	require_GRUB_CFG
	cbfstool "$FIRMWARE_ROM" remove -n grubtest.cfg 2>/dev/null
	cbfstool "$FIRMWARE_ROM" add -f "$GRUB_CFG" -n grubtest.cfg -t raw
}

append_boot_uuid () {
	require_GRUB_CFG
	if grep -q cryptboot /etc/crypttab 2>/dev/null; then
		UUID=$(cryptboot_uuid | sed 's/-//g')
	else
		UUID=$(findmnt /boot -no UUID)
		if [ -z "$UUID" ]; then
			UUID=$(findmnt / -no UUID)
		fi
	fi
	sed -E -e 's/# (set uuids=)/\1/' -e 's/# (export uuids)/\1/' -e "s/uuids='/\0 $UUID/" -i "$GRUB_CFG"
	add_grub_cfg
}

set_grub_password () {
	require_GRUB_CFG
	echo '> Set a GRUB password'
	until grep -qF 'grub.pbkdf2' grub_pw 2>/dev/null; do
		grub-mkpasswd-pbkdf2 | tee grub_pw
	done
	GRUB_PW=$(sed -nEe 's/.*(grub\.pbkdf2.*)/\1/p' grub_pw)
	sed -E -e 's/# (set superusers=)/\1/' -e 's/# (password_pbkdf2)/\1/' -e "s/grub\.pbkdf2.*/$GRUB_PW/" -i "$GRUB_CFG"
	rm grub_pw
	add_grub_cfg
}

secure_boot () {
	require_FIRMWARE_ROM
	require_GRUB_CFG

	# Enable secure boot.
	sed -E -e 's/# (trust)/\1/' -e 's/# (set secure_boot)/\1/' -e 's/# (export secure_boot)/\1/' -i "$GRUB_CFG"
	add_grub_cfg

	# Create signature.
	if ! gpg --list-keys 'Local boot <root@localhost>' 1>/dev/null 2>/dev/null; then
		gpg --batch --passphrase '' --quick-generate-key 'Local boot <root@localhost>' dsa3072 sign 0
	fi
	gpg --export 'Local boot <root@localhost>' > boot.key

	# Add signature public key to firmware.
	cbfstool "$FIRMWARE_ROM" remove -n boot.key 2>/dev/null
	cbfstool "$FIRMWARE_ROM" add -f boot.key -n boot.key -t raw
	rm boot.key

	# Sign GRUB files.
	cat <<- EOF > /usr/local/bin/sign-boot
	#!/bin/sh
	[ $(id -u) -eq 0 ] || exit
	echo 'Signing files in /boot'
	for pattern in '*.cfg' 'vmlinuz*' 'initr*' 'grubenv' '*.pf2' '*.png' '*.jpg'; do
	    find /boot \( -name "\$pattern" -a ! -name '*.sig' \) -exec gpg --local-user 'Local boot <root@localhost>' --detach-sign --yes {} \;
	done
	EOF
	chmod +x /usr/local/bin/sign-boot
	/usr/local/bin/sign-boot

	SECURE_BOOT=y # so we can display help when we finish
}

flash_firmware_rom () {
	require_FIRMWARE_ROM
	echo '> Flashing firmware'
	flashrom -p internal --ifd -i bios -w "$FIRMWARE_ROM" 1>/dev/null
	rm "$FIRMWARE_ROM"
}

flash_delayed () {
	flash_firmware_rom
	rm grubtest.cfg
	cat <<- EOF
	We are almost done! Please reboot your system and:

	- Enter into the GRUB console by pressing the letter \`c\` within 2 seconds after the boot image fades.
	- Type \`configfile (cbfsdisk)/grubtest.cfg\` and hit <Enter>.
	EOF

	if [ -n "$GRUB_PW" ]; then
		cat <<- EOF
		- Enter into the GRUB console again by pressing \`c\`.
		- It should ask for a user, type \`root\` as user, and your GRUB password.
		- Exit the console by pressing <Esc>.
		EOF
	fi

	if [ -n "$SECURE_BOOT" -o -n "$UUID" ]; then
		echo '- Try to boot your system.'
	fi

	if [ -n "$GRUB_PW" ]; then
		printf 'If your password is not accepted or '
	else
		printf 'If '
	fi

	cat <<- EOF
	you can't boot means that something went wrong. Worry not, changes will not be permanent until you make them permanent. Reboot and try to fix them by running =hwbtool= once again (only applicable to \`--grub\` options), or seek help.

	To make changes permanent: boot into your system with the kernel param \`iomem=relaxed\`, login as \`root\` and run \`hwbtool --commit-grub\`. Warning: If you use other GRUB related options previous to \`--commit-grub\` current changes will be discarded.
	EOF

	if [ -n "$SECURE_BOOT" ]; then
		echo ""
		echo "A utility for signing boot files has been provided in \`/usr/local/bin/sign-boot\`. You will need to run it after each kernel upgrade. You can automate this by creating a symbolic link on specific locations: On Arch based distros, \`/etc/initcpio/post/\`, and on Debian based distros, \`/etc/initramfs/post-update.d/\`."
	fi
}

commit_grub () {
	require_FIRMWARE_ROM
	cbfstool "$FIRMWARE_ROM" extract -n grubtest.cfg -f grubtest.cfg 2>/dev/null && \
	cbfstool "$FIRMWARE_ROM" remove -n grubtest.cfg && \
	cbfstool "$FIRMWARE_ROM" remove -n etc/grub.cfg && \
	cbfstool "$FIRMWARE_ROM" add -n etc/grub.cfg -f grubtest.cfg -t raw && \
	rm grubtest.cfg && \
	flash_firmware_rom && \
	echo 'Your GRUB firmware configuration has been commited succesfully.'
}

until [ $# -eq 0 ]; do
	case "$1" in
		--user) set_hostname && set_root_password && create_user ;;
		--luks*)
			LUKS_PASSWORD="$2"
			case "$1" in
				--luks-boot) set_luks_boot_password ;;
				--luks-root) set_luks_root_password ;;
				*) set_luks_boot_password && set_luks_root_password ;;
			esac
			shift
			;;
		--grub-uuid) append_boot_uuid && DELAY_FLASH=y ;;
		--grub-password) set_grub_password && DELAY_FLASH=y ;;
		--grub-secure-boot) secure_boot && DELAY_FLASH=y ;;
		--grub) append_boot_uuid && set_grub_password && secure_boot && DELAY_FLASH=y ;;
		--commit-grub) commit_grub ;;
	esac
	shift
done
if [ -n "$DELAY_FLASH" ]; then flash_delayed; fi