#!/bin/sh

yes_no() {
    while true; do
        printf "$1 (y/n): "
        read -r yn
        case $yn in
            [Yy]* ) return 0;;
            [Nn]* ) return 1;;
            * ) echo "Please answer yes or no";;
        esac
    done
}

# Ensure we run as root
if [ "$EUID" -ne 0 ]; then
    echo "Must be ran as root"
    exit 1
fi

# Take NEWROOT as 1st argument
if [ $# -ge 1 ]; then
    NEWROOT="$1"
else
    echo "Provide newroot directory"
    exit 1
fi

# Take chroot user as 2nd argument, default to root
if [ $# -ge 2 ]; then
    USERNAME="$2"
else
    USERNAME="root"
fi

# Check if given NEWROOT is already mounted, if it is
# set REMOUNT to the mount source, so that we can remount
# it once we're done.
df_out=$(df --output=source,target | grep -w "$NEWROOT")
if [ -n "$df_out" ]; then
    REMOUNT="$(echo $df_out | awk '{print $1}')"
else
    # If the target isn't mounted already, check
    # if user gave $3 (mount location)
    if [ $# -ge 3 ]; then
        mount "$3" "$NEWROOT"
    else
        # If user didn't give mount location, try to
        # mount according to fstab
        if [ -n "$(grep -w "$NEWROOT" /etc/fstab)" ]; then
            mount "$NEWROOT"
        else
            # Ask for user confirmation to ensure that filesystem
            # is ready for chroot in given NEWROOT, exit if not
            yes_no "$NEWROOT wasn't mounted, is your filesystem in place?" || exit 1
        fi
    fi
fi

# Mount necessary directories for chroot to be possible
mount --types proc /proc "$NEWROOT/proc"
mount --rbind /sys "$NEWROOT/sys"
mount --make-rslave "$NEWROOT/sys"
mount --rbind /dev "$NEWROOT/dev"
mount --make-rslave "$NEWROOT/dev"

# Use /bin/su for chrooting with --login to also run
# /etc/profile and ~/.profile or ~/.zprofile
chroot "$NEWROOT" "/bin/su" "$USERNAME" --login

# Unmount recursively mounted directories
umount -l "$NEWROOT/dev"
umount -l "$NEWROOT/sys"
umount -l "$NEWROOT/proc"
umount -R "$NEWROOT"

# Remount partition according to fstab if REMOUT is set
# in order to leave the filesystem in the state it was
if [ -n "$REMOUNT" ]; then
    mount "$REMOUNT" "$NEWROOT"
fi