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

set -eu

usage() {
  echo "Run this command from a rescue system on SD-card." >&2
  echo "With no positional arguments, execute a bash shell inside the system" >&2
  echo "on eMMC/SSD." >&2
  echo "With positional arguments, instead of an interactive shell, run the" >&2
  echo "command and its arguments passed as positional arguments to this tool" >&2
  echo "inside the system on eMMC/SSD" >&2
  echo "This tool takes care of unlocking full disk encryption set up by" >&2
  echo "reform-setup-encrypted-disk and mounts /boot, /sys, /proc and /dev" >&2
  echo "It also copies /etc/resolv.conf from the current system into the" >&2
  echo "chroot, so if you have set up network connection on the outside," >&2
  echo "then the system on eMMC/SSD will also have network access" >&2
  echo >&2
  echo "Usage: $0 [--help] [CMD...]" >&2
  echo >&2
  echo "Options:" >&2
  echo "  -h, --help       Display this help and exit." >&2
  echo >&2
  echo "Examples:" >&2
  echo >&2
  echo "Get a shell inside the system you have on SSD or eMMC:" >&2
  echo >&2
  echo "  $0" >&2
  echo >&2
  echo "Run reform-check inside the system on SSD or eMMC" >&2
  echo >&2
  echo "  $0 reform-check" >&2
  echo >&2
}

while getopts :h-: OPTCHAR; do
  case "$OPTCHAR" in
    h)
      usage
      exit 0
      ;;
    -)
      case "$OPTARG" in
        help)
          usage
          exit 0
          ;;
        *)
          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))"

# by default, if no positional arguments were passed, run bash inside the chroot
if [ "$#" -eq 0 ]; then
  set -- /bin/bash
fi

if [ "$(id -u)" -ne 0 ]; then
  echo "reform-setup-encrypted-disk 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 "vgchange" >/dev/null 2>&1 || {
  echo >&2 'Please install "lvm2" using: apt install lvm2'
  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

# We need to wrap findmnt output in realpath because if the SD-card was mounted
# via its label, then we need to resolve /dev/disk/by-label to the device name
case "$(realpath "$(findmnt --noheadings --evaluate --mountpoint / --output SOURCE)")" in
  "/dev/${DEV_SD}p"*) : ;;
  *)
    echo "E: This script is meant to be run from a rescue system from SD-card" >&2
    exit 1
    ;;
esac

# eMMC device is being used (case 1): mount points are known, show them. Includes swap.
if [ -n "$(lsblk --noheadings --output=MOUNTPOINT "/dev/${DEV_MMC}")" ]; then
  echo "E: eMMC has the following mounted volumes, unmount them before running this tool" >&2
  lsblk --noheadings --output=MOUNTPOINT "/dev/${DEV_MMC}" | xargs --no-run-if-empty -I '{}' echo "E:   {}" >&2
  exit 1
fi

# eMMC device is being used (case 2): there are not file systems directly mounted on the block device
# but it is opened by consumers like device-mapper, raid or luks, to name some examples. In this situation
# it is not trivial to locate the consumer.
# reform-boot-config and reform-emmc-bootstrap do the same thing (could share code?)
get_exclusive_write_lock() {
  ret=0
  python3 - "$1" <<EOF || ret=$?
import errno, os, sys

try:
    os.open(sys.argv[1], os.O_WRONLY | os.O_EXCL)
except OSError as e:
    if e.errno == errno.EBUSY:
        sys.exit(1)
    raise
EOF
  return $ret
}

if ! get_exclusive_write_lock "/dev/${DEV_MMC}"; then
  echo "E: device /dev/${DEV_MMC} (eMMC) is still in use" >&2
  exit 1
fi

mmc_part_types="$(parted --script --json /dev/${DEV_MMC} print 2>/dev/null | jq --compact-output '[ .disk.partitions.[]? | .filesystem ]')"
if [ "$mmc_part_types" = "[]" ] && cryptsetup isLuks "/dev/${DEV_MMC}"; then
  mmc_part_types=luks
fi

ssd_part_types="[]"
if [ -b "/dev/${DEV_SSD}" ]; then
  ssd_part_types="$(parted --script --json /dev/${DEV_SSD} print 2>/dev/null | jq --compact-output '[ .disk.partitions.[]? | .filesystem ]')"
  if [ "$ssd_part_types" = "[]" ] && cryptsetup isLuks "/dev/${DEV_SSD}"; then
    ssd_part_types=luks
  fi
fi

DEV_SSD_BASE="$DEV_SSD"
if [ "$DEV_SSD" != "sda" ]; then
  DEV_SSD_BASE="${DEV_SSD}p"
fi

trap cleanup EXIT INT TERM
MOUNTROOT="$(mktemp --tmpdir --directory reform-emmc-bootstrap.XXXXXXXXXX)"

cleanup() {
  if [ -e "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak" ] || [ -L "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak" ]; then
    rm "$MOUNTROOT/etc/resolv.conf" || echo "W: removal of /etc/resolv.conf failed" >&2
    mv "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak" "$MOUNTROOT/etc/resolv.conf" || echo "W: restoring /etc/resolv.conf failed" >&2
  fi
  if mountpoint --quiet "$MOUNTROOT"; then
    umount --recursive "$MOUNTROOT"
  fi
  rmdir "$MOUNTROOT"
  if [ -e /dev/reformvg ]; then
    vgchange -an reformvg >&2
  fi
  if [ -e /dev/mapper/reform_crypt ]; then
    cryptsetup luksClose reform_crypt
  fi
}

main() {
  rootdev="$1"
  bootdev="$2"
  shift 2
  mount "$rootdev" "$MOUNTROOT"
  for dir in etc boot dev sys proc; do
    if [ ! -d "$MOUNTROOT/$dir" ]; then
      echo "E: The directory '$dir' does not exist in the filesystem on $rootdev" >&2
      exit 1
    fi
  done

  mount -o bind /dev "$MOUNTROOT/dev/"
  mount -t sysfs sys "$MOUNTROOT/sys/"
  mount -t proc proc "$MOUNTROOT/proc/"
  mount "$bootdev" "$MOUNTROOT/boot/"

  if [ ! -d "$MOUNTROOT/boot/extlinux" ] && [ ! -e "$MOUNTROOT/boot/boot.scr" ]; then
    echo "E: Neither extlinux directory nor boot.scr exist in filesystem on $bootdev" >&2
    printf "Do you still want to continue with %s mounted as /boot? [y/N] " "$bootdev" >&2
    read -r response
    if [ "$response" != "y" ]; then
      echo "Exiting."
      exit 1
    fi
  fi

  if [ -e /etc/resolv.conf ]; then
    if [ -e "$MOUNTROOT/etc/resolv.conf" ] || [ -L "$MOUNTROOT/etc/resolv.conf" ]; then
      mv "$MOUNTROOT/etc/resolv.conf" "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak"
    fi
    cp --dereference /etc/resolv.conf "$MOUNTROOT/etc/resolv.conf"
  fi

  chroot "$MOUNTROOT" "$@"

  if [ -e "$MOUNTROOT/etc/resolv.conf" ] && [ -e "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak" ] || [ -L "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak" ]; then
    rm "$MOUNTROOT/etc/resolv.conf"
    mv "$MOUNTROOT/etc/resolv.conf.reform-rescue-shell.bak" "$MOUNTROOT/etc/resolv.conf"
  fi

  umount --recursive "$MOUNTROOT"
}

case "$mmc_part_types;$ssd_part_types" in
  '["ext4"'*'];luks') echo "I: Assuming /boot on first partition of eMMC and root filesystem on encrypted SSD." >&2 ;;
  '["ext4","ext4"];[]') echo "I: Assuming both /boot and root filesystem on eMMC." >&2 ;;
  '["ext4"'*'];["ext4","swap"]') echo "I: Assuming /boot on eMMC and root filesystem on first partition of SSD" >&2 ;;
  '["ext4"'*'];["swap","ext4"]') echo "I: Assuming /boot on eMMC and root filesystem on second partition of SSD" >&2 ;;
  *)
    echo "E: your configuration $mmc_part_types;$ssd_part_types is currently unsupported, please file a ticket with details of your setup to get it supported" >&2
    exit 1
    ;;
esac

printf "Does that sound right? [y/N] " >&2
read -r response
if [ "$response" != "y" ]; then
  echo "Exiting."
  exit 1
fi

# for all cases where SSD is about to be used, check for mounted volumes and
# attempt to get an exclusive write lock
case "$mmc_part_types;$ssd_part_types" in '["ext4"'*'];luks' | '["ext4"'*'];["ext4","swap"]' | '["ext4"'*'];["swap","ext4"]')
  if [ -n "$(lsblk --noheadings --output=MOUNTPOINT "/dev/${DEV_SSD}")" ]; then
    echo "E: SSD has the following mounted volumes, unmount them before running this tool" >&2
    lsblk --noheadings --output=MOUNTPOINT "/dev/${DEV_SSD}" | xargs --no-run-if-empty -I '{}' echo "E:   {}" >&2
    exit 1
  fi
  if ! get_exclusive_write_lock "/dev/${DEV_SSD}"; then
    echo "E: device /dev/${DEV_SSD} (SSD) is still in use" >&2
    exit 1
  fi
  ;;
esac

case "$mmc_part_types;$ssd_part_types" in
  '["ext4"'*'];luks')
    cryptsetup luksOpen "/dev/${DEV_SSD}" reform_crypt
    vgchange -ay reformvg >&2

    main /dev/reformvg/root "/dev/${DEV_MMC}p1" "$@"

    vgchange -an reformvg >&2
    cryptsetup luksClose reform_crypt
    ;;
  '["ext4","ext4"];[]') main "/dev/${DEV_MMC}p2" "/dev/${DEV_MMC}p1" ;;
  '["ext4"'*'];["ext4","swap"]') main "/dev/${DEV_SSD_BASE}1" "/dev/${DEV_MMC}p1" ;;
  '["ext4"'*'];["swap","ext4"]') main "/dev/${DEV_SSD_BASE}2" "/dev/${DEV_MMC}p1" ;;
esac

rmdir "$MOUNTROOT"

trap - EXIT INT TERM
