#!/bin/bash

# I simply dumped the contents of the Leap-15.4 vagrant image for this:
TREE=/tmp/leap-15.4-image.tar

# Both the installer and the eventual system image will be at 2G
IMAGE_SIZE=2048

# Just for prototyping. For the final implementation, we need the user to
# supply the passphrase, or we use a randomly generated one
# DEFAULT_PASSPHRASE="kingcrimson"

# List of PCRs to seal the LUKS key to
SEAL_PCR_LIST=0,1,2,3,4,5,6,7
#SEAL_PCR_LIST=0,4
SEAL_PCR_LIST=0,2
SEAL_PCR_BANK=sha256

# These need to match exactly what grub2 uses to create the SRK
TPM2_SRK_ATTRS="userwithauth|restricted|decrypt|fixedtpm|fixedparent|noda|sensitivedataorigin"

# Distro name used in the /boot/efi/EFI/$distro path
DISTRO=opensuse

BUILDDIR=/tmp/build
INSTALL_IMAGE=$BUILDDIR/install-image
SYSTEM_IMAGE=$BUILDDIR/system-image
MOUNTDIRS="/dev /sys /proc"

opt_quiet=false
opt_testing=true

function create_empty_image {

	image_path=$1

	echo "Creating empty image at $image_path; size=$IMAGE_SIZE MiB"
	dd if=/dev/zero of=$image_path bs=1MiB count=$IMAGE_SIZE status=none
}

function partition_disk {

	disk_dev=$1
	fstype=${2:-ext2}

	parted $disk_dev mklabel gpt
	parted $disk_dev mkpart fat32 1 100
	parted $disk_dev name 1 "p.UEFI"
	parted $disk_dev set 1 esp on
	parted $disk_dev mkpart $fstype 100 2048
	parted $disk_dev name 2 "p.lnxroot"

	parted $disk_dev print
}

function partition_image {

	image_path=$1
	fstype=${2:-ext2}

	echo "Creating partition table for $image_path"
	loop_dev=$(losetup --show -f $image_path)
	partition_disk $loop_dev $fstype
	losetup -d $loop_dev
}

function attach_image {

	declare -g LOOPDEV BOOTDEV ROOTDEV

	image_path=$1
	dev_path=$2

	LOOPDEV=$(losetup -P --show -f $image_path)
	fdisk -l $LOOPDEV

	echo "Attaching image $image_path"
	mkdir -p $dev_path
	ln -sf $LOOPDEV $dev_path/disk
	BOOTDEV=${LOOPDEV}p1
	ROOTDEV=${LOOPDEV}p2

	ln -sf $BOOTDEV $dev_path/boot
	ln -sf $ROOTDEV $dev_path/root
	find $dev_path -ls

	mkfs -t fat $BOOTDEV
	mkfs -t ext4 $ROOTDEV
}

function mount_device {

	device=$1
	mountpoint=$2

	declare -g MOUNTED_DIRS

	echo "mount $device $mountpoint"
	mount $device $mountpoint

	MOUNTED_DIRS="$mountpoint $MOUNTED_DIRS"
}

function bind_host_fs {

	host_dir=$1
	mountpoint=$2

	echo "bind mount $host_dir"
	mount -o bind $host_dir $mountpoint

	MOUNTED_DIRS="$mountpoint $MOUNTED_DIRS"
}

function mount_image {

	build_root=$1
	root_dev=$2
	boot_dev=$3

	mount_device $root_dev $build_root
	mkdir -p $build_root/boot/efi
	mount_device $boot_dev $build_root/boot/efi
}

function get_path_device {

	df --output=source $1 | tail -1
}

function get_partition_uuid {

	partx --show -o UUID -g $*
}

function get_device_uuid {

	shortname=$(basename $1)
	lsblk -nlo uuid,name "$1" | while read uuid name; do
		if [ $name = $shortname ]; then
			echo $uuid
			break
		fi
	done
}

function get_os_release {

	root_dev=$1

	. $root_dev/etc/os-release
	echo "$PRETTY_NAME"
}

function install_system_image {

	tarball=$1
	root_path=$2

	for name in boot dev sys proc run; do
		mkdir -p $root_path/$name
	done

	echo "Unpacking $tarball to $root_path"
	tar -C $root_path -xf $tarball
}

function install_system_repo {

	root_path=$1
	repo_name=$2
	repo_url=$3
	key_url=$4

	rpm --root $root_path --import $key_url

	zypper --root $root_path ar $repo_url $repo_name
	zypper --root $root_path refresh FDE

	# interactive_mode
}

function install_package {

	root_path=$1; shift

#	cache=$(mktemp -d /tmp/pkgcache.XXXXXX)
	zypper --root $root_path install -f -y "$@"
#	rm -rf $cache
}

function detect_kernel_version {

	boot_path=$1

	set -- $(readlink $boot_path/initrd)
	if [ $# -eq 1 ]; then
		set -- $(expr $1 : 'initrd-\(.*\)')
	fi

	if [ $# -ne 1 ]; then
		echo "Unable to detect default image" >&2
		return 1
	fi

	echo $1
}

function grub_mount_encrypted_commands {

	sealed_luks_key=$1

	# We use the UUID of the underlying partition that holds
	# the luks volume.
	uuid=$(get_device_uuid $CRYPTDEV)
	uuid=${uuid//-/}

	cat <<EOF
set uuid=$uuid

insmod tpm2
sleep 1
tpm2_key_protector_init -b $SEAL_PCR_BANK -p $SEAL_PCR_LIST -k \$prefix/$sealed_luks_key

insmod luks2
sleep 1
cryptomount -u $uuid -k tpm2
sleep 2

set root="cryptouuid/$uuid"
set prefix="(\$root)/boot/grub2"
sleep 2
configfile (\$root)/boot/grub2/grub.cfg
EOF
}

function grub_mount_unencrypted_commands {

	uuid=$1

	cat <<EOF
tpm_record_pcrs
set btrfs_relative_path="yes"
search --fs-uuid --set=root $uuid
set prefix=(\$root)/boot/grub2
configfile (\$root)/boot/grub2/grub.cfg
EOF
}

function install_grub {

	root_path=$1
	sealed_luks_key=$2

	boot_path="$root_path/boot"

	kernel_version=$(detect_kernel_version $boot_path)
	if [ -z "$kernel_version" ]; then
		return 1
	fi

	echo "Installing grub"
	grub2-install $LOOPDEV --target=x86_64-efi \
		--disable-shim-lock \
		--boot-directory $boot_path

#	. /etc/default/grub

	root_dev=$(get_path_device $root_path)
	uuid=$(get_device_uuid $root_dev)
	os_release=$(get_os_release $root_path)

	if [ -n "$sealed_luks_key" ]; then
		kernel_root="rd.luks.uuid=$uuid rd.luks.name=$uuid=root"
		kernel_root+=" rd.luks.key=efivar:LuksBootKey"
		kernel_root+=" rd.luks.options=timeout=10s,discard,tries=1"
	else
		kernel_root="root=UUID=$uuid"
	fi
	echo "Describing root fs as: $kernel_root"

	echo "Creating grub.cfg for kernel $kernel_version"
	sed -e "s|@KERNEL_ROOT_SPEC@|$kernel_root|" \
	    -e "s|@OS_RELEASE@|$os_release|" \
	    -e "s|@KERNEL_VERSION@|$kernel_version|" \
	    < grub.cfg.in > $boot_path/grub2/grub.cfg

	efi_distro_path=$boot_path/efi/EFI/$DISTRO
	efi_boot_path=$boot_path/efi/EFI/BOOT

	mkdir -p $efi_distro_path $efi_boot_path

	if [ -f "$efi_distro_path/grub.cfg" ]; then
		mv "$efi_distro_path/grub.cfg" "$efi_distro_path/grub.cfg.orig"
	fi

	if [ -n "$sealed_luks_key" ]; then
		cp -v $sealed_luks_key $efi_distro_path/sealed.key
		grub_mount_encrypted_commands sealed.key >$efi_distro_path/grub.cfg
	else
		grub_mount_unencrypted_commands $uuid >$efi_distro_path/grub.cfg
	fi

	echo "Contents of $efi_distro_path/grub.cfg"
	cat $efi_distro_path/grub.cfg
	echo

	mv $efi_distro_path/* $efi_boot_path

	if $opt_testing; then
		if [ -f bootx64.efi ]; then
			echo "Installing local bootx64.efi for testing"
			cp -v bootx64.efi $efi_boot_path/bootx64.efi
		fi
	fi

	find $boot_path/efi -ls

	echo "Installing openSUSE theme"
	mkdir -p $boot_path/grub2/themes/openSUSE
	cp -a /boot/grub2/themes/openSUSE/* $boot_path/grub2/themes/openSUSE
}

function install_fstab {

	root_path=$1
	root_dev=$2
	boot_dev=$3

	echo "Creating /etc/fstab"
	root_blk_uuid=$(get_device_uuid $root_dev)
	boot_blk_uuid=$(get_device_uuid $boot_dev)

	cat >$root_path/etc/fstab <<EOF
UUID=$root_blk_uuid / ext4 defaults 0 1
UUID=$boot_blk_uuid /boot/efi vfat defaults 0 0
EOF

	cat $root_path/etc/fstab
	echo
}

function rebuild_mkinitrd {

	root_path=$1

	echo "Re-creating initrd"

	# -B tells mkinitrd to not update the bootloader, lest we
	# end up overwriting our carefully crafted grub.cfg
	chroot $root_path mkinitrd -B -m "virtio_blk virtio_pci"
}

function forget_persistent_device_names {

	root_path=$1

	rm -f $root_path/etc/udev/rules.d/70-persistent-net.rules
}

function change_user_password {

	root_path=$1
	user_name=$2
	user_pass=$3

	usermod --root $root_path --password "$user_pass" $user_name
}

function install_packages {

	zypper in -y "$@"
}

function copy_installer {

	destdir=$1

	# requires libtss2-fapi libtss2-tcti-device0
	mkdir -p $destdir
	cp -v $0 $destdir
	cp install-luks grub.cfg.in $destdir
	cp -v $TREE $destdir/os-image.tar

	if $opt_testing; then
		if [ -f bootx64.efi ]; then
			cp -v bootx64.efi $destdir
		fi
		if [ -f systemd-cryptsetup-generator ]; then
			cp -v systemd-cryptsetup-generator $destdir
		fi
	fi
}

function unmount_image {

	if [ -n "$MOUNTED_DIRS" ]; then
		echo "Unmounting image ($MOUNTED_DIRS)"
		umount $MOUNTED_DIRS
		MOUNTED_DIRS=
	fi
}

function cleanup {

	unmount_image

	if [ -n "$LOOPDEV" ]; then
		losetup -d $LOOPDEV
		LOOPDEV=
	fi
	if [ -n "$LUKSNAME" ]; then
		echo "closing LUKS device $LUKSNAME"
		cryptsetup close $LUKSNAME
		LUKSNAME=
	fi
}

# Danger zone
function force_cleanup {

	losetup -a|sed 's,:.*,,'|while read dev; do
		losetup -d $dev
	done
}

function create_install_image {

	mkdir -p $BUILDDIR

	create_empty_image $SYSTEM_IMAGE

	create_empty_image $INSTALL_IMAGE
	partition_image $INSTALL_IMAGE
	attach_image $INSTALL_IMAGE $BUILDDIR/dev

	ROOT=$BUILDDIR/root
	mount_image $ROOT $ROOTDEV $BOOTDEV

	BOOTDIR=$ROOT/boot

	install_system_image $TREE $ROOT

	install_system_repo $root_path FDE \
		https://download.opensuse.org/repositories/home:/okir:/FDE/15.4 \
		https://download.opensuse.org/repositories/home:/okir:/FDE/15.4/repodata/repomd.xml.key

	# Install EFI boot loader stuff without triggering an update-bootloader run
	cp $root_path/etc/sysconfig/bootloader{,.save}
	echo > $root_path/etc/sysconfig/bootloader
	install_package $root_path \
		-r FDE \
		grub2-x86_64-efi shim
	mv $root_path/etc/sysconfig/bootloader{.save,}

	df $ROOT

	# shim-install only seems to work when running in the installled system,
	# as it uses efibootmgr which wants to scribble over the system's EFI
	# variables...
	if false; then
		shim-install --efi-directory $BOOTDIR/efi $BOOTDEV

		echo
		echo "After shim-install"
		find $BOOTDIR
		echo
	fi

	install_fstab $ROOT $ROOTDEV $BOOTDEV

	echo
	df $BOOTDEV $ROOTDEV
	echo

	for dir in $MOUNTDIRS; do
		bind_host_fs $dir $ROOT$dir
	done

	if ! install_grub $ROOT; then
		exit 1
	fi

	rebuild_mkinitrd $ROOT
	forget_persistent_device_names $ROOT

	# Change the password of root to "root"
	change_user_password $ROOT root '$1$iGdfrvuE$3eOcnGP0Z5blS1cHZ.Vk/1'

	copy_installer $ROOT/root

	if false; then
		echo "Entering image for manual tweaking"
		chroot $ROOT
	fi
}

function prepare_luks_partition {

	declare -g CRYPTDEV BOOTDEV ROOTDEV LUKSNAME

	device=$1
	luks_name=$2
	master_key=$3
	luks_key=$4

	echo "Generating random LUKS partition key"
	dd if=/dev/random bs=1 count=32 of=$luks_key

	# Alternatively, the installed could just use a user supplied
	# passphrase (and protect that with the TPM)
	# Beware, cryptsetup will read the whole file and consider an existing
	# newline part of the pass phrase. Which makes it kinda hard to enter
	# that passphrase interactively later.
	if [ -n "$DEFAULT_PASSPHRASE" ]; then
		echo -n $DEFAULT_PASSPHRASE > $luks_key
	fi

	echo "Setting up $device as an encrypted partition"
	cryptsetup luksFormat --type luks2 --pbkdf PBKDF2 -q --verbose $device $luks_key
	cryptsetup luksDump --dump-master-key --key-file $luks_key -q --master-key-file $master_key $device

	dmsetup table --showkeys

	cryptsetup open --key-file $luks_key $device $luks_name

	CRYPTDEV=$device
	LUKSNAME=$luks_name
	ROOTDEV=/dev/disk/by-id/dm-name-$luks_name

	mkfs -t ext4 $ROOTDEV || exit 1
}

function tpm2 {

	cmd="tpm2_$1"; shift

	echo "$cmd $@"
	$cmd "$@"
}

function seal_key {

	pcr_list=$1
	secret=$2
	pubkey=$3
	privkey=$4

	dir=$(mktemp -d /dev/shm/seal.XXXXXX)

	context=$dir/primary.ctx
	session=$dir/session.ctx
	policy=$dir/policy.tpm

	tpm2 createprimary --quiet -c $context -a "$TPM2_SRK_ATTRS"
	tpm2 startauthsession --session $session 
	tpm2 policypcr --session $session --pcr-list $pcr_list --policy $policy
	tpm2 flushcontext $session 
	tpm2 create --quiet -C $context -u $pubkey -r $privkey -L $policy -i $secret 
	cp $pubkey $privkey /tmp

	rm -rf $dir
}

function unseal_key {

	pcr_list=$1
	pubkey=$2
	privkey=$3
	unsealed_key_file=$4

	dir=$(mktemp -d /dev/shm/seal.XXXXXX)

	context=$dir/primary.ctx
	session=$dir/session.ctx
	policy=$dir/policy.tpm
	key_context=$dir/key.ctx

	tpm2 createprimary --quiet -c $context -a "$TPM2_SRK_ATTRS"
	tpm2 startauthsession --policy-session -S $session
	tpm2 policypcr --session $session --pcr-list $pcr_list --policy $policy
	tpm2 load -C $context -u $pubkey -r $privkey -c $key_context
	tpm2 unseal -c $context -p session:$session -c $key_context -o $unsealed_key_file
	tpm2 flushcontext $session 

	rm -rf $dir
}

function seal_key_verify {

	pcr_list=$1
	secret=$2
	sealed_key_file=$3

	tmpdir=$(mktemp -d /dev/shm/seal.XXXXXX)
	pubkey=$tmpdir/key.pub
	privkey=$tmpdir/key.priv

	seal_key $pcr_list $secret $pubkey $privkey

	# The sealed key that grub expects seems to be just
	# the concatenation of TPM2B_PUBLIC and a TPM2B containing
	# the private key portion
	cat $pubkey $privkey > $sealed_key_file

	unsealed_file=$tmpdir/unsealed.key
	unseal_key $pcr_list $pubkey $privkey $unsealed_file

	if ! cmp --quiet $secret $unsealed_file; then
		echo "Problem: TPM seal/unseal did not work!"
		retval=1
	else
		echo "Splendid, we were able to unseal the TPM protected key"
		retval=0
	fi

	rm -rf $tmpdir
	return $retval
}

function mounted_virtio_disks {

	grep '/dev/vd[a-z]' /proc/mounts |
		sed 's:\(/dev/vd.\).*$:\1:' | sort -u
}

function install_requisite_packages {

	# Copy all files required to prepare the system image
	# We do this _first_ because cryptsetup insists on rebuilding
	# initrd. Without the virtio modules.
	install_packages \
		e2fsprogs \
		btrfsprogs \
		dosfstools \
		system-user-tss \
		libtss2-fapi1 \
		libtss2-tcti-device0 \
		tpm2.0-tools \
		cryptsetup

	if [ -x /root/systemd-cryptsetup-generator ]; then
		echo "Overwriting systemd-cryptsetup-generator with test version"
		cp /root/systemd-cryptsetup-generator /usr/lib/systemd/system-generators/systemd-cryptsetup-generator
	fi
}

function install_system {

	set -- /dev/vd?

	mounted=$(mounted_virtio_disks|tr '\15' '|')
	echo $mounted
	for dev in /dev/vd?; do
		case $dev in
		$mounted)
			echo "$dev is in use";;
		*)	echo "Using $dev as installation target"
			install_dev=$dev
			break;;
		esac
	done

	# Wipe any existing disk label
	dd if=/dev/zero of=$install_dev bs=1MiB count=1 status=none
	partition_disk $install_dev

	boot_dev=${install_dev}1
	root_dev=${install_dev}2

	install_requisite_packages

	master_key=/dev/shm/luks.master
	luks_key=/dev/shm/luks.key
	sealed_luks_key=/root/luks.key.sealed

	prepare_luks_partition $root_dev "root" $master_key $luks_key

	# Use the dm-crypt target instead of the underlying raw device
	root_dev=$ROOTDEV

	pcr_list="$SEAL_PCR_BANK:$SEAL_PCR_LIST"
	if ! seal_key_verify $pcr_list $luks_key $sealed_luks_key; then
		echo "Refusing to create encrypted partition - unlocking would fail"
		rm -f $luks_key $master_key
		return 1
	fi

	# for debugging purposes. A real installer could offer the user
	# to back up the master key to a USB stick or some other device
	cp $master_key /root
	cp $luks_key /root

	rm -f $master_key $luks_key

	mkfs -t fat $boot_dev

	root_path=/mnt
	boot_path=$root_path/boot

	mount_image $root_path $root_dev $boot_dev

	install_system_image $TREE $root_path
	df $root_path

	if ! install_grub $root_path $sealed_luks_key; then
		exit 1
	fi

	install_fstab $root_path $root_dev $boot_dev

	echo
	df $BOOTDEV $ROOTDEV
	echo

	for dir in $MOUNTDIRS; do
		bind_host_fs $dir $root_path$dir
	done

	rebuild_mkinitrd $root_path
	forget_persistent_device_names $root_path

	# Change the password of root to "root"
	change_user_password $root_path root '$1$iGdfrvuE$3eOcnGP0Z5blS1cHZ.Vk/1'

}

function test_tpm {

	echo "super secret" > /tmp/secret

	pcr_list="$SEAL_PCR_BANK:$SEAL_PCR_LIST"
	if ! seal_key_verify $pcr_list /tmp/secret /tmp/sealed; then
		echo "Seal/unseal is not working"
		exit 1
	fi
}

# For testing purposes
function interactive_mode {

	echo
	echo "*** Entering interactive mode ***"
	PS1=">> " bash
	echo "*** Done with interactive mode ***"
}

MOUNTED_DIRS=""
trap "cleanup" 0 1 2 15

# set -x
# set -- $(getopt qT: "$@")
eval set -- $(getopt -l quiet,tree:,no-testing -o qT: -- "$@")

while [ $# -gt 0 ]; do
	opt=$1; shift
	case $opt in
	-q|--quiet)
		opt_quiet=true;;
	-T|--tree)
		TREE=$1; shift;;
	--no-testing)
		opt_testing=false;;
	--)	break;;
	*)	echo "Unexpected option $opt" >&2
		exit 1;;
	esac
done

case $# in
0)	action=create-installer;;
1)	action=$1;;
*)	echo "Too many arguments; expected one action name" >&2
	exit 1;;
esac

case $action in
create-installer)
	create_install_image;;
install-system)
	install_system;;
test-tpm)
	test_tpm;;
*)
	echo "Unsupported action $action" >&2
	exit 1;;
esac
	
cleanup

if [ -f /root/luks.master ]; then
	echo "***"
	echo " /root contains a copy of the master key. This is for debugging; make sure to disable this code later"
	echo "***"
	sleep 1
fi

echo Done.
