#!/bin/sh

set -eu

usage() {
  echo "Set up the SSD with LUKS encryption containing LVM2 volumes for" >&2
  echo "swap and the rootfs. Offers to run reform-migrate and " >&2
  echo "reform-boot-config to set up the initramfs in /boot to load the" >&2
  echo "rootfs from the encrypted SSD." >&2
  echo >&2
  echo "Usage: $0 [--help]" >&2
  echo >&2
  echo "Options:" >&2
  echo "  --help           Display this help and exit." >&2
}

if [ "$#" -gt 0 ]; then
  if [ "$1" != "--help" ]; then
    echo "E: too many arguments" >&2
    usage
    exit 1
  fi
  usage
  exit 0
fi

if [ "$(id -u)" -ne 0 ]; then
  echo "reform-setup-encrypted-nvme has to be run as root / using sudo."
  exit
fi

command -v "cryptsetup" >/dev/null 2>&1 || {
  echo >&2 'Please install "cryptsetup" using: apt install cryptsetup'
  exit 1
}
command -v "lvchange" >/dev/null 2>&1 || {
  echo >&2 'Please install "lvm2" using: apt install lvm2'
  exit 1
}
command -v "mkfs.ext4" >/dev/null 2>&1 || {
  echo >&2 'Please install "e2fsprogs" using: apt install e2fsprogs'
  exit 1
}

# shellcheck source=/dev/null
if [ -e "./machines/$(cat /proc/device-tree/model).conf" ]; then
  . "./machines/$(cat /proc/device-tree/model).conf"
elif [ -e "/usr/share/reform-tools/machines/$(cat /proc/device-tree/model).conf" ]; then
  . "/usr/share/reform-tools/machines/$(cat /proc/device-tree/model).conf"
else
  echo "E: unable to find config for $(cat /proc/device-tree/model)" >&2
  exit 1
fi

if [ "$DEV_SSD" = "sda" ]; then
  HUMAN="SATA"
else
  HUMAN="NVMe"
fi

echo "This will ERASE ALL DATA from your $HUMAN SSD."

echo ""

printf "Are you sure you want to proceed? [y/N] "
read -r response

if [ "$response" != "y" ]; then
  echo "Exiting."
  exit
fi

cleanupvg() { vgchange -an reformvg; }
cleanupluks() { cryptsetup luksClose reform_crypt; }
error() { echo "$0 FAILED to run" >&2; }

trap error EXIT INT TERM

if [ -n "$(lsblk --noheadings --output=MOUNTPOINT "/dev/$DEV_SSD")" ]; then
  echo "/dev/$DEV_SSD is still in use" >&2
  exit 1
fi

# /proc/meminfo contains the sizes in kibibytes
mem="$(awk '/^MemTotal:/ {print $2}' /proc/meminfo)"
case "$mem" in *[!0123456789]* | 0?* | "")
  echo "E: unable to acquire total memory from /proc/meminfo" >&2
  exit 1
  ;;
esac
# convert memory size to gigabytes, rounding up
mem="$(awk 'BEGIN {printf("%.f",'"$mem"'/1024/1024+0.5)}')"
# minimum swap size is 4G
if [ "$mem" -lt 4 ]; then
  mem=4
fi

disksize=$(lsblk -bno SIZE "/dev/$DEV_SSD" | head -1)
case "$disksize" in *[!0123456789]* | 0?* | "")
  echo "E: unable to acquire disk size of /dev/$DEV_SSD" >&2
  exit 1
  ;;
esac
# convert disk size to gigabytes, rounding down
disksize="$((disksize / 1024 / 1024 / 1024))"

# maximum swap size is 5% of disk size
if [ "$mem" -gt "$((disksize * 5 / 100))" ]; then
  mem="$((disksize * 5 / 100))"
fi

if [ "$mem" -le 1 ]; then
  echo "E: your disk is too small for swap" >&2
  exit 1
fi

cryptsetup luksFormat "/dev/$DEV_SSD"
trap "cleanupluks; error" EXIT INT TERM
cryptsetup luksOpen "/dev/$DEV_SSD" reform_crypt
pvcreate /dev/mapper/reform_crypt
vgcreate reformvg /dev/mapper/reform_crypt
trap "cleanupvg; cleanupluks; error" EXIT INT TERM
lvcreate --name swap --size "${mem}G" reformvg
mkswap /dev/reformvg/swap
lvcreate --name root --extents 100%FREE reformvg
mkfs.ext4 /dev/reformvg/root
SWAPUUID=$(blkid -s UUID -o value /dev/reformvg/swap)
CRYPTUUID=$(blkid -s UUID -o value "/dev/$DEV_SSD")

echo ""
printf "The encrypted %s is now set up. Do you want me to run reform-migrate now as well? [y/N] " "$HUMAN"
read -r response

if [ "$response" != "y" ]; then
  echo "If you want to migrate this system to $HUMAN SSD you can now run:"
  echo ""
  echo "echo RESUME=UUID=$SWAPUUID > /etc/initramfs-tools/conf.d/resume"
  echo "echo reform_crypt UUID=$CRYPTUUID none luks,discard,x-initrd.attach > /etc/crypttab"
  echo "echo UUID=$SWAPUUID none swap sw 0 0 >> /etc/fstab"
  echo "cryptsetup luksOpen /dev/$DEV_SSD reform_crypt"
  echo "vgchange -ay reformvg"
  echo "reform-migrate /dev/reformvg/root"
  echo "vgchange -an reformvg"
  echo "cryptsetup luksClose reform_crypt"
else
  # we are not really running reform-migrate but imitate what it does
  # instead because we want to write out some files after the rsync but
  # before running reform-boot-config
  ROOTMNT="$(mktemp --tmpdir --directory reform-setup-encrypted-nvme.XXXXXXXXXX)"
  trap 'umount $ROOTMNT; cleanupvg; cleanupluks; error' EXIT INT TERM
  mount /dev/reformvg/root "$ROOTMNT"
  rsync -axHAWXS --numeric-ids --info=progress2 / "$ROOTMNT"
  echo "RESUME=UUID=$SWAPUUID" >"$ROOTMNT/etc/initramfs-tools/conf.d/resume"
  echo "reform_crypt UUID=$CRYPTUUID none luks,discard,x-initrd.attach" >"$ROOTMNT/etc/crypttab"
  echo "UUID=$SWAPUUID none swap sw 0 0" >>"$ROOTMNT/etc/fstab"
  trap "cleanupvg; cleanupluks; error" EXIT INT TERM
  umount "$ROOTMNT"

  emmc_flag=
  if [ "$EMMC_USE" = true ]; then
    printf "Your /boot partition will be on eMMC by default. Do you want it on the SD-Card instead? [y/N] "
    read -r response
    if [ "$response" != "y" ]; then
      emmc_flag="--emmc"
    fi
  fi
  ret=0
  reform-boot-config $emmc_flag /dev/reformvg/root || ret=$?
  if [ "$ret" -ne 0 ]; then
    echo "reform-boot-config failed. To re-run it manually, perform the following steps:" >&2
    echo "    $ cryptsetup luksOpen \"/dev/$DEV_SSD\" reform_crypt" >&2
    echo "    $ vgchange -ay reformvg" >&2
    echo "    $ reform-boot-config $emmc_flag /dev/reformvg/root" >&2
    echo "    $ vgchange -an reformvg" >&2
    echo "    $ cryptsetup luksClose reform_crypt" >&2
    exit "$ret"
  fi
fi

trap "cleanupluks; error" EXIT INT TERM
cleanupvg

trap "error" EXIT INT TERM
cleanupluks

trap - EXIT INT TERM

echo "You can now reboot into your encrypted System."
