#!/bin/bash
set -Eeuo pipefail
# disable command path hashing as we are going to move fast and break things
set +h

WORKDIR="/tmp/menhera"
MIRROR="https://iso.meson.cc/ultralite/archlinux"
ROOTFS="${MIRROR}/iso/latest/arch/x86_64/airootfs.sfs"

# internal global variables
OLDROOT="/"
NEWROOT=""

[[ "$(uname -m)" == "x86_64" ]] || { echo "error: wrong arch $(uname -m) (not x86_64)"; exit 1; }

PUBKEY="./key.pub"
[ -f "$PUBKEY" ] || { echo "error: you should put your ssh public key at ${PUBKEY} in order to ssh with root account and this key later"; exit 1; }

if [[ $EUID -ne 0 ]]; then
    echo "This script must be run as root"
    exit 1
fi

systemctl status > /dev/null || { echo "error: no systemd"; exit 1; }

# helper functions
# https://stackoverflow.com/a/3232082/2646069
menhera::confirm() {
    # call with a prompt string or use a default
    read -r -p "${1:-Are you sure? [y/N]} " response
    case "$response" in
        [yY][eE][sS]|[yY])
            true
            ;;
        *)
            false
            ;;
    esac
}

echo -e "We will start a temporary RAM system as your recovery environment."
echo -e "Note that this script will kill programs and umount filesystems without prompting."
echo -e "Please confirm:"
echo -e "\tYou have closed all programs you can, and backed up all important data"
echo -e "\tYou can SSH into your system as root user"
menhera::confirm || exit -1

# fix possible PATH problems
export PATH="$PATH:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

echo "Loading kernel modules..."
modprobe overlay
modprobe squashfs
sysctl kernel.panic=10

echo "Creating workspace in '${WORKDIR}'..."
# workspace
mkdir -p "${WORKDIR}"
mount -t tmpfs -o size=100% tmpfs "${WORKDIR}"

mkdir -p "${WORKDIR}/newroot" # where iso is stored
mkdir -p "${WORKDIR}/newrootro" # lower
mkdir -p "${WORKDIR}/newrootrw" # upper
mkdir -p "${WORKDIR}/overlayfs_workdir" # work

echo "Downloading temporary rootfs..."
wget -O "${WORKDIR}/rootfs.squashfs" "${ROOTFS}"

echo "Mounting temporary rootfs..."
mount -t squashfs "${WORKDIR}/rootfs.squashfs" "${WORKDIR}/newrootro"
mount -t overlay overlay -o rw,lowerdir="${WORKDIR}/newrootro",upperdir="${WORKDIR}/newrootrw",workdir="${WORKDIR}/overlayfs_workdir" "${WORKDIR}/newroot"
NEWROOT="${WORKDIR}/newroot"

rm -f "${NEWROOT}/etc/resolv.conf"
cat "${OLDROOT}/etc/resolv.conf" > "${NEWROOT}/etc/resolv.conf"
install -DTm0600 <(cat "$PUBKEY") "${NEWROOT}/root/.ssh/authorized_keys"
chmod 0700 "${NEWROOT}/root/.ssh"
sync && sync # why sync twice? idk

echo "Swapping rootfs..."
# prepare future mount point for our old rootfs
mkdir -p "${WORKDIR}/newroot/mnt/oldroot"
mount --make-rprivate /

# swap root
pivot_root "${WORKDIR}/newroot" "${WORKDIR}/newroot/mnt/oldroot"

OLDROOT="/mnt/oldroot"
NEWROOT="/"

# move mounts
for i in dev proc sys run; do
    if [ -d "${OLDROOT}/$i" ]; then
        mount --move "${OLDROOT}/$i" "${NEWROOT}/$i"
    fi
done
mount -t tmpfs -o size=100% tmpfs "${NEWROOT}/tmp"

mkdir -p "${WORKDIR}"
mount --move "${OLDROOT}/${WORKDIR}" "${WORKDIR}"

echo "Restarting SSH daemon..."
systemctl daemon-reload
systemctl is-active ssh && systemctl stop ssh || true
systemctl restart sshd

echo "Disabling swap..."
swapoff -a

echo "Restarting init process..."
systemctl daemon-reexec

systemctl --quiet is-system-running || true

echo -e "If you are connecting from SSH, please create a second session to this host use root and"
echo -e "confirm you can get a shell."
echo -e "After your confirmation, we are going to kill the old SSH server."

if menhera::confirm; then
    echo "Killing all programs still using the old root..."
    fuser -kvm "${OLDROOT}" -15
    # in most cases the parent process of this script will be killed, so goodbye
else
    echo -e "Please manually issue a reboot to recover your old OS. If you believe there is a bug in menhera.sh, "
    echo -e "raise a ticket at https://github.com/Jamesits/menhera.sh/issues ."
fi