#!/usr/bin/env bash

DEVICE="${1}"
BACKUP_DIR="./backups"
DATE_FIX=$(date '+%Y%m%d-%H%M%S')
DD_OPTS="bs=512 iflag=fullblock conv=notrunc"
TARGET_BOOT="qubes_dom0-boot"

welcome() {
	echo "################################"
	echo "This script will encrypt an unencrypted /boot partition"
	echo "Confirmation will be asked before writing"
	echo "################################"
}

warning() {
	echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
	echo "The following steps may corrupt and lose your data, continue at your own risk"
	echo "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
	read -r
}

adios() {
	echo "[+] Procedure completed!"
}

panic() {
	echo "[*] Something went wrong in a write operation, system may be in a corrupted state. Attempting recovery"
	restore
	exit 1
}

restore() {
	echo "[*] Attempting to restore original partition scheme"
	dd if=${BACKUP_DIR}/mbr-${DATE_FIX}.img of=${DEVICE} bs=512 iflag=fullblock conv=notrunc status=progress
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Something went wrong restoring, hope you made a backup as advised ☠"
	fi
}

check_params() {
	if [[ "${1}" -ne 1 ]]; then
		echo "Usage: ./encryptboot.sh <device>"
		echo "Example: ./encryptboot.sh /dev/sda"
		exit
	fi
}

check_root() {
    if [[ "${EUID}" -ne 0 ]]; then
        echo "[-] This script must be run as root; re-run prefixed with sudo"
        exit 1
    fi
}

check_device() {
	if [[ ! -b "${DEVICE}" ]]; then
		echo "[-] Device ${DEVICE} does not exists"
		exit 1
	fi
}

backup_boot() {
	echo "[+] Backing up boot device"
	mkdir -p "${BACKUP_DIR}"
	dd if=${DEVICE}1 of=${BACKUP_DIR}/boot-${DATE_FIX}.img ${DD_OPTS} status=progress
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Something went wrong backing up boot partition, exiting"
		exit 1
	fi
	BOOT_HASH=$(sha256sum ${DEVICE}1 | cut -d ' ' -f 1)
	BOOT_BACKUP_HASH=$(sha256sum ${BACKUP_DIR}/boot-${DATE_FIX}.img | cut -d ' ' -f 1)
	if [[ ${BOOT_HASH} != ${BOOT_BACKUP_HASH} ]]; then
		echo "[-] Backup ${BACKUP_DIR}/boot-${DATE_FIX}.img hash is not equal to ${DEVICE}1 hash, exiting"
		exit 1
	fi
	echo "[+] Backup successful"
}

backup_partition_table() {
	echo "[+] Backing up partition table"
	mkdir -p "${BACKUP_DIR}"
	dd if=${DEVICE} of=${BACKUP_DIR}/mbr-${DATE_FIX}.img ${DD_OPTS} count=1
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Something went wrong backing up partition table, exiting"
		exit 1
	fi
}

check_headers() {
	BOOT_HEADER=$(dd if=${DEVICE}1 ${DD_OPTS} count=16 2>/dev/null | file -s -)
	LUKS_HEADER=$(dd if=${DEVICE}2 ${DD_OPTS} count=16 2>/dev/null | file -s -)
	if [[ "${BOOT_HEADER}" != *"ext4"* ]]; then
		echo "[-] ${DEVICE}1 is not an ext4 filesystem"
		exit 1
	fi
	if [[ "${LUKS_HEADER}" != *"LUKS"* ]]; then
		echo "[-] ${DEVICE}2 is not a LUKS container"
		exit
	fi
	echo "[+] Headers check completed"
}

get_offsets() {
	echo "[+] Getting boot partition offsets"
	START_OFFSET=$(parted -s ${DEVICE} unit s print | grep boot | tr -s ' ' | cut -d ' ' -f 3 | tr -d 's')
	END_OFFSET=$(parted -s ${DEVICE} unit s print | grep boot | tr -s ' ' | cut -d ' ' -f 4 | tr -d 's')
	if [[ "${START_OFFSET}" -le 0 ]] || [[ "${END_OFFSET}" -le 0 ]] || [[ "${END_OFFSET}" -le ${START_OFFSET} ]]; then
		echo "[-] Error parsing boot partition get_offsets"
		exit 1
	fi
	#OFFSET=$((${END_OFFSET}-${START_OFFSET}))
	OFFSET=$((${END_OFFSET}+1))
}

delete_partitions() {
	echo "[+] Deleting old partition scheme"
	parted "${DEVICE}" rm 1
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Something went wrong deleting boot partition"
		panic
	fi
	parted "${DEVICE}" rm 2
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Something went wrong deleting LUKS partition"
		panic
	fi
}

create_partition() {
	echo "[+] Creating new full disk partition"
	parted -s ${DEVICE} mkpart primary luks 0% 100%
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Something went wrong creatig the new partition"
		panic
	fi
}

check_offsets() {
	echo ${START_OFFSET}
	echo ${END_OFFSET}
	echo ${OFFSET}
	LUKS_HEADER=$(dd if=${DEVICE}1 ${DD_OPTS} skip=${OFFSET} seek=0 count=16 2>/dev/null | file -s -)
	if [[ "${LUKS_HEADER}" != *"LUKS"* ]]; then
		echo "[-] Luks header not found at given offset "
		exit
	fi

}

move_data() {
	dd if=${DEVICE}1 of=${DEVICE}1 ${DD_OPTS} skip=${OFFSET} seek=0 status=progress
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Failed moving data backwards, hope you had backups because this is most likely total corruption. MBR and boot.img backups are in ${BACKUP_DIR}"
		exit
	fi
}

config_luks_lvm() {
	echo "[+] Extending LVM pool"
	cryptsetup luksOpen ${DEVICE}1 qubespv
	pvresize qubespv
	echo "[+] Creating LVM boot partition"
	lvcreate -n boot -l100%FREE ${TARGET_BOOT}
}

restore_boot() {
	echo "[+] Copying old boot image in new encrypted LVM volume "
	dd if=${BACKUP_DIR}/boot-${DATE_FIX}.img of=/dev/mapper/${TARGET_BOOT} ${DD_OPTS} status=progress
	if [[ "${?}" -ne 0 ]]; then
		echo "[-] Failed to copy back boot.img to LVM, probably a recoverable state but needs manual intervention"
		exit
	fi
	LVM_BOOT_HASH=$(sha256sum ${TARGET_BOOT} | cut -d ' ' -f 1)
	if [[ ${BOOT_HASH} != ${LVM_BOOT_HASH} ]]; then
		echo "[-] "
		exit 1
	fi
	echo "[+] Boot partition written back successfully"
}

check_params "${#}"
welcome
check_root
check_device
backup_partition_table
backup_boot
check_headers
get_offsets
check_offsets
warning
#delete_partitions
#create_partition
#move_data
#config_luks_lvm
#restore_boot
adios