#!/bin/bash
# Runs LOCALLY. One-time setup on a fresh machine:
#   1. Installs dropbear-initramfs
#   2. Copies SSH authorized key into dropbear's keystore
#   3. Configures initramfs networking (IP=dhcp, DEVICE=eth0)
#   4. Rebuilds the Ubuntu initramfs with dropbear included
#   5. Adds the 'dropbear-rescue' GRUB entry (idempotent)
#
# No reboot occurs. After this script, use 02_reboot_to_ramdisk.sh.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "$SCRIPT_DIR/.env"

: "${REMOTE:?Edit .env and set REMOTE}"

ENTRY_ID="dropbear-rescue"

echo "==> Installing dropbear-initramfs on $REMOTE"
ssh $SSH_OPTS "$REMOTE" "DEBIAN_FRONTEND=noninteractive apt-get install -y -q dropbear-initramfs"

echo ""
echo "==> Configuring dropbear"
ssh $SSH_OPTS "$REMOTE" "
  set -euo pipefail

  # Authorized key for SSH access in rescue mode
  mkdir -p /etc/dropbear/initramfs
  cp /root/.ssh/authorized_keys /etc/dropbear/initramfs/authorized_keys
  chmod 600 /etc/dropbear/initramfs/authorized_keys

  # IP=dhcp as a kernel param is what dropbear's init-premount hook reads.
  # The DEVICE= and IP= lines in initramfs.conf are for network-root boots
  # (NFS/iSCSI) and have no effect here — the kernel param is what matters.
  # We set them anyway as documentation.
  grep -q '^DEVICE=' /etc/initramfs-tools/initramfs.conf \
    || echo 'DEVICE=eth0' >> /etc/initramfs-tools/initramfs.conf
  grep -q '^IP=' /etc/initramfs-tools/initramfs.conf \
    || echo 'IP=dhcp' >> /etc/initramfs-tools/initramfs.conf

  echo '  dropbear configured'
"

echo ""
echo "==> Rebuilding Ubuntu initramfs with dropbear"
ssh $SSH_OPTS "$REMOTE" "update-initramfs -u -k \$(uname -r) 2>&1 | tail -2"

echo ""
echo "==> Adding GRUB entry '$ENTRY_ID'"
ROOT_DEV=$(ssh $SSH_OPTS "$REMOTE" "findmnt -n -o SOURCE /")

# If /boot is a separate partition, GRUB must search that partition's UUID
# and use paths without the /boot prefix (e.g. /vmlinuz not /boot/vmlinuz).
# If /boot is on the root filesystem, search root UUID and use /boot/vmlinuz.
BOOT_IS_SEPARATE=$(ssh $SSH_OPTS "$REMOTE" "findmnt -n /boot >/dev/null 2>&1 && echo yes || echo no")
if [ "$BOOT_IS_SEPARATE" = "yes" ]; then
  BOOT_UUID=$(ssh $SSH_OPTS "$REMOTE" "findmnt -n -o UUID /boot")
  VMLINUZ_PATH="/vmlinuz"
  INITRD_PATH="/initrd.img"
else
  BOOT_UUID=$(ssh $SSH_OPTS "$REMOTE" "findmnt -n -o UUID /")
  VMLINUZ_PATH="/boot/vmlinuz"
  INITRD_PATH="/boot/initrd.img"
fi
echo "  /boot separate: $BOOT_IS_SEPARATE — search UUID: $BOOT_UUID"

# Use vmlinuz/initrd.img symlinks (not versioned filenames) so the entry
# survives kernel upgrades without needing to be regenerated.
ssh $SSH_OPTS "$REMOTE" "
  set -euo pipefail

  # Remove any stale entry so it gets rewritten with current values.
  sed -i '/^menuentry.*$ENTRY_ID/,/^}/d' /etc/grub.d/40_custom 2>/dev/null || true

  cat >> /etc/grub.d/40_custom << 'GRUBEOF'

menuentry 'Dropbear Rescue (boot_custom_os)' --id '$ENTRY_ID' {
    set gfxpayload=text
    insmod gzio
    insmod part_gpt
    insmod ext2
    search --no-floppy --fs-uuid --set=root $BOOT_UUID
    echo 'Starting dropbear rescue (Ubuntu initramfs + SSH)...'
    linux   $VMLINUZ_PATH root=$ROOT_DEV ro net.ifnames=0 biosdevname=0 IP=dhcp nomodeset break=mount console=ttyS0,115200n8 console=tty1
    initrd  $INITRD_PATH
}
GRUBEOF
  echo '  Entry written to /etc/grub.d/40_custom.'
  update-grub 2>&1 | tail -2
"

echo ""
echo "Setup complete on $REMOTE."
echo "  dropbear-initramfs installed and keyed"
echo "  Ubuntu initramfs rebuilt with dropbear"
echo "  GRUB entry: '$ENTRY_ID'"
echo ""
echo "Next step:"
echo "  ./02_reboot_to_ramdisk.sh"