#!/bin/sh
# SPDX-License-Identifier: GPL-3.0+
# Copyright 2022-2025 Johannes Schauer Marin Rodrigues <josch@mister-muffin.de>

set -eu

usage() {
  echo "Set up a disk 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 on the first" >&2
  echo "partition of the eMMC or SD-card to load the rootfs from the " >&2
  echo "encrypted disk." >&2
  echo >&2
  echo "Usage: $0 [--help] [DEVICE_PATH]" >&2
  echo >&2
  echo "Options:" >&2
  echo "  --help           Display this help and exit." >&2
  echo "  --force          Avoid user interaction: run reform-migrate and" >&2
  echo "                   put /boot on eMMC if allowed. Requires passphrase" >&2
  echo "                   on standard input." >&2
  echo >&2
  echo "Arguments:" >&2
  echo "  [DEVICE_PATH]    Disk, partition or regular file to set up with " >&2
  echo "                   LUKS and LVM. If this argument is missing, defaults" >&2
  echo "                   to the SSD." >&2
  echo "                   If the eMMC device is provided, it will be" >&2
  echo "                   partitioned with the first partition for /boot and" >&2
  echo "                   the second partition for LUKS&LVM." >&2
  echo "                   The shorthands ssd, emmc, sd and usb can be used" >&2
  echo "                   to install to the SSD, the eMMC, the SD card or" >&2
  echo "                   to a USB drive respectively." >&2
  echo "                   To copy to a file of the same name as one of the" >&2
  echo "                   shorthands, prefix the filename with a dot and" >&2
  echo "                   and a slash for a relative path." >&2
  echo >&2
  echo "Example 1:" >&2
  echo "When booted from SD-card set up the first partition on eMMC as /boot" >&2
  echo "and set up the SSD with LUKS. Set up LVM inside LUKS with" >&2
  echo "one volume for swap and the other for the root file system. Copy" >&2
  echo "the current system to /dev/reformvg/root:" >&2
  echo >&2
  echo "    $0" >&2
  echo >&2
  echo "Example 2:" >&2
  echo "When booted from SD-card, set up eMMC with two partitions, the first" >&2
  echo "one for /boot and the second for LUKS. Set up LVM inside LUKS with" >&2
  echo "one volume for swap and the other for the root file system. Copy" >&2
  echo "the current system to /dev/reformvg/root. Suppress any interactivity" >&2
  echo "by using --force and obtain the LUKS passphrase from standard input." >&2
  echo >&2
  echo "    echo hunter2 | $0 --force emmc" >&2
  echo >&2
}

FORCE=false
while getopts :h-: OPTCHAR; do
  case "$OPTCHAR" in
    h)
      usage
      exit 0
      ;;
    -)
      case "$OPTARG" in
        help)
          usage
          exit 0
          ;;
        force) FORCE=true ;;
        *)
          echo "E: unrecognized option: --$OPTARG" >&2
          exit 1
          ;;
      esac
      ;;
    :)
      echo "E: missing argument for -$OPTARG" >&2
      exit 1
      ;;
    '?')
      echo "E: unrecognized option -$OPTARG" >&2
      exit 1
      ;;
    *)
      echo "E: error parsing options" >&2
      exit 1
      ;;
  esac
done
shift "$((OPTIND - 1))"

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

DEVICE_PATH=

if [ "$#" -eq 1 ]; then
  DEVICE_PATH="$1"
elif [ "$#" -gt 1 ]; then
  echo "E: invalid number of arguments" >&2
  usage
  exit 1
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 [ -z "$DEVICE_PATH" ]; then
  DEVICE_PATH="/dev/$DEV_SSD"
fi

case $DEVICE_PATH in "/dev/$DEV_MMC"*)
  case $EMMC_USE in
    false)
      echo "E: writing to eMMC not supported on $(cat /proc/device-tree/model)" >&2
      exit 1
      ;;
    warn)
      echo "W: Using eMMC on $(cat /proc/device-tree/model) is not without risk." >&2
      echo "W: For example, flashing the wrong u-boot or if the flashing process goes wrong, it is" >&2
      echo "W: possible to soft-brick your board. Restoring it might need some extra hardware." >&2
      echo "W: Please only proceed if you are sure that the benefits outweigh the risks for you." >&2
      if [ "$FORCE" = true ]; then
        echo "Proceeding without user interaction because of --force"
        response="y"
      else
        printf "Are you sure you want to proceed? [y/N] "
        read -r response
      fi

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

case $DEVICE_PATH in
  sd) DEVICE_PATH="/dev/${DEV_SD}" ;;
  ssd) DEVICE_PATH="/dev/${DEV_SSD}" ;;
  emmc) DEVICE_PATH="/dev/${DEV_MMC}" ;;
  usb) DEVICE_PATH="/dev/${DEV_USB}" ;;
esac

if [ -b "$DEVICE_PATH" ]; then
  case $DEVICE_PATH in
    /dev/sda*) HUMAN="SATA SSD" ;;
    /dev/nvme0n1*) HUMAN="NVMe SSD" ;;
    "/dev/$DEV_MMC"*) HUMAN="eMMC" ;;
    "/dev/$DEV_SD"*) HUMAN="SD card" ;;
    *) HUMAN="SSD" ;;
  esac

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

  echo ""

  if [ "$FORCE" = true ]; then
    echo "Proceeding without user interaction because of --force"
    response="y"
  else
    printf "Are you sure you want to proceed? [y/N] "
    read -r response
  fi

  if [ "$response" != "y" ]; then
    echo "Exiting."
    exit
  fi
elif [ -f "$DEVICE_PATH" ]; then
  # plain file -- do nothing
  HUMAN="reguar file"
else
  echo "E: $DEVICE_PATH is neither block device nor regular file" >&2
  exit 1
fi

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

trap error EXIT INT TERM

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

# If the user chose eMMC or SD-card as the target device, then it has to be
# partitioned first so that the unencrypted /boot partition can be in front
# of the luks partition. Users who do not want this behaviour can manually
# partition their eMMC or SD-card.
case $DEVICE_PATH in "/dev/$DEV_MMC" | "/dev/$DEV_SD")
  # Space in MB which needs to be kept empty for u-boot before the first
  # partition starts. Biggest requirement comes from rk3588 which expects
  # u-boot at an offset of 8 MiB and a 4 MiB trust (ATF, OP-TEE)
  # partition after that. So the first partition can only start at byte
  # 16777216 or 16 MiB: https://opensource.rock-chips.com/wiki_Partitions
  MAX_UBOOT=16
  BOOTSIZE=488
  /sbin/parted --script --machine "$DEVICE_PATH" "mklabel msdos"
  /sbin/parted --script --machine "$DEVICE_PATH" "mkpart primary ext4 ${MAX_UBOOT}MiB $((BOOTSIZE + MAX_UBOOT))MiB"
  /sbin/parted --script --machine "$DEVICE_PATH" "mkpart primary $((BOOTSIZE + MAX_UBOOT))MiB 100%"
  mkfs.ext4 "${DEVICE_PATH}p1"
  case $DEVICE_PATH in
    "/dev/$DEV_MMC") DEVICE_PATH="/dev/${DEV_MMC}p2" ;;
    "/dev/$DEV_SD") DEVICE_PATH="/dev/${DEV_SD}p2" ;;
  esac
  ;;
esac

# /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

if [ -f "$DEVICE_PATH" ]; then
  disksize="$(stat -c %s "$DEVICE_PATH")"
else
  disksize=$(lsblk --nodeps --bytes --noheadings --output=SIZE "$DEVICE_PATH" | head -1)
  case "$disksize" in *[!0123456789]* | 0?* | "")
    echo "E: unable to acquire disk size of $DEVICE_PATH" >&2
    exit 1
    ;;
  esac
fi

# convert disk size to gigabytes, rounding down
disksize="$((disksize / 1024 / 1024 / 1024))"

case $DEVICE_PATH in
  "/dev/$DEV_MMC"* | "/dev/$DEV_SD"*)
    : # do not limit swap size for eMMC and SD-card
    ;;
  *)
    # maximum swap size is 5% of disk size
    if [ "$mem" -gt "$((disksize * 5 / 100))" ]; then
      mem="$((disksize * 5 / 100))"
    fi
    ;;
esac

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

if [ "$FORCE" = true ]; then
  read -r PASSPHRASE
  printf "%s" "$PASSPHRASE" | cryptsetup luksFormat "$DEVICE_PATH" "-"
else
  cryptsetup luksFormat "$DEVICE_PATH"
fi
trap "cleanupluks; error" EXIT INT TERM
if [ "$FORCE" = true ]; then
  printf "%s" "$PASSPHRASE" | cryptsetup luksOpen --key-file=- "$DEVICE_PATH" reform_crypt
else
  cryptsetup luksOpen "$DEVICE_PATH" reform_crypt
fi
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=$(lsblk --nodeps --noheadings --output=UUID /dev/reformvg/swap)
CRYPTUUID=$(cryptsetup luksUUID "$DEVICE_PATH")

echo ""
if [ "$FORCE" = true ]; then
  echo "Running reform-migrate because of --force"
  response="y"
else
  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
fi

if [ "$response" != "y" ]; then
  echo "If you want to migrate this system to $HUMAN 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 $DEVICE_PATH 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-disk.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" != false ]; then
    if [ "$FORCE" = true ]; then
      echo "Placing /boot on eMMC because of --force"
      emmc_flag="--emmc"
    else
      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
  fi
  ret=0
  if [ "$FORCE" = true ]; then
    reform-boot-config --force $emmc_flag /dev/reformvg/root || ret=$?
  else
    reform-boot-config $emmc_flag /dev/reformvg/root || ret=$?
  fi
  if [ "$ret" -ne 0 ]; then
    echo "reform-boot-config failed. To re-run it manually, perform the following steps:" >&2
    echo "    $ cryptsetup luksOpen \"$DEVICE_PATH\" 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."
