#!/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 "Unless you pass the --root and --boot options, this tool will try to" >&2
  echo "figure out your configuration automatically." >&2
  echo "Guessing the configuration only works when running a system from SD-card" >&2
  echo >&2
  echo "Usage: $0 [--help] [CMD...]" >&2
  echo >&2
  echo "Options:" >&2
  echo "  -h, --help       Display this help and exit." >&2
  echo " --boot PATH       Device/partition with /boot on it" >&2
  echo " --root PATH       Device/partition with / on it. If you pass a luks" >&2
  echo "                   device, it will be assumed to have been set up with" >&2
  echo "                   an LVM2 volume group reformvg inside of it."
  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
}

nth_arg() {
  shift "$1"
  printf "%s" "$1"
}

ROOT=
BOOT=
while getopts :h-: OPTCHAR; do
  case "$OPTCHAR" in
    h)
      usage
      exit 0
      ;;
    -)
      case "$OPTARG" in
        help)
          usage
          exit 0
          ;;
        boot)
          if [ "$OPTIND" -gt "$#" ]; then
            echo "E: missing argument for --boot" >&2
            exit 1
          fi
          BOOT="$(nth_arg "$OPTIND" "$@")"
          OPTIND=$((OPTIND + 1))
          ;;
        root)
          if [ "$OPTIND" -gt "$#" ]; then
            echo "E: missing argument for --root" >&2
            exit 1
          fi
          ROOT="$(nth_arg "$OPTIND" "$@")"
          OPTIND=$((OPTIND + 1))
          ;;
        *)
          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 [ -n "$BOOT" ] && [ -z "$ROOT" ]; then
  echo "E: setting --boot also requires --root to be set" >&2
  exit 1
fi

if [ -n "$ROOT" ] && [ -z "$BOOT" ]; then
  echo "E: setting --root also requires --boot to be set" >&2
  exit 1
fi

# 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

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

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"
}

# if --root and --boot were not specified, try to guess the setup
if [ -z "$ROOT" ] && [ -z "$BOOT" ]; then
  # 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: When not running from SD-card, the --root and --boot options are required." >&2
      exit 1
      ;;
  esac

  mmc_part_types="$(parted --script --json "/dev/${DEV_MMC}" print 2>/dev/null | jq --compact-output '[ .disk.partitions.[]? | .filesystem ]')"

  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

  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
      BOOT="/dev/${DEV_MMC}p1"
      ROOT="/dev/${DEV_SSD}"
      ;;
    '["ext4","ext4"];[]')
      echo "I: Assuming both /boot and root filesystem on eMMC." >&2
      BOOT="/dev/${DEV_MMC}p1"
      ROOT="/dev/${DEV_MMC}p2"
      ;;
    '["ext4"'*'];["ext4","swap"]')
      echo "I: Assuming /boot on eMMC and root filesystem on first partition of SSD" >&2
      BOOT="/dev/${DEV_MMC}p1"
      ROOT="/dev/${DEV_SSD_BASE}1"
      ;;
    '["ext4"'*'];["swap","ext4"]')
      echo "I: Assuming /boot on eMMC and root filesystem on second partition of SSD" >&2
      BOOT="/dev/${DEV_MMC}p1"
      ROOT="/dev/${DEV_SSD_BASE}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
fi

for device in "$ROOT" "$BOOT"; do
  if [ ! -e "$device" ]; then
    echo "E: $device does not exist" >&2
    exit 1
  fi

  if [ ! -b "$device" ]; then
    echo "E: $device is not a block device" >&2
    exit 1
  fi

  if [ -n "$(lsblk --noheadings --output=MOUNTPOINT "$device")" ]; then
    echo "E: $device has the following mounted volumes, unmount them before running this tool" >&2
    lsblk --noheadings --output=MOUNTPOINT "$device" | xargs --no-run-if-empty -I '{}' echo "E:   {}" >&2
    exit 1
  fi

  if ! get_exclusive_write_lock "$device"; then
    echo "E: unable to get exclusive write lock for $device" >&2
    exit 1
  fi
done

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

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"
}

cleanup_luks() {
  cleanup
  if [ -e /dev/reformvg ]; then
    vgchange -an reformvg >&2
  fi
  if [ -e /dev/mapper/reform_crypt ]; then
    cryptsetup luksClose reform_crypt
  fi
}

if cryptsetup isLuks "$ROOT"; then
  if [ -e /dev/mapper/reform_crypt ]; then
    echo "E: /dev/mapper/reform_crypt already exists" >&2
    exit 1
  fi

  if [ -e /dev/reformvg/root ]; then
    echo "E: /dev/reformvg/root already exists" >&2
    exit 1
  fi

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

  cryptsetup luksOpen "$ROOT" reform_crypt
  vgchange -ay reformvg >&2

  main /dev/reformvg/root "$BOOT" "$@"

  vgchange -an reformvg >&2
  cryptsetup luksClose reform_crypt
else
  trap cleanup EXIT INT TERM
  MOUNTROOT="$(mktemp --tmpdir --directory reform-emmc-bootstrap.XXXXXXXXXX)"

  main "$ROOT" "$BOOT" "$@"
fi

rmdir "$MOUNTROOT"

trap - EXIT INT TERM
