#!/bin/bash

PARTITION=$1
COMMANDS=$2
PASSPHRASE=$3

TARGET='/target'

if [ -z ${PARTITION} ] || [ ! -e ${PARTITION} ]; then
  echo "Chroot a partition or system directory"
  echo
  echo "Examples:"
  echo "$0 /dev/sda1"
  echo "$0 /path/to/systemdir"
  echo
  echo "Run a command:"
  echo "$0 /dev/sda1 \"command 1;command 2\""
  echo
  echo "Encrypted partition:"
  echo "$0 /dev/sda1 \"\" my_passphrase"
  echo "$0 /dev/sda1 \"command 1;command 2\" my_passphrase"
  exit 0
fi

# Be root to run this script
if [ $UID -ne 0 ]; then
  sudo $0 "$@"
  exit 0
fi

# Check if partition is already mounted and use that as the target
if [ "${PARTITION:0:4}" == '/dev' ]; then
    # Map luks encrypted partition
    FS=$(blkid -s TYPE -o value ${PARTITION})
    if [ "$FS" == 'crypto_LUKS' ]; then
        echo "Map encrypted partition ${PARTITION}"
        PNAME=$(basename ${PARTITION})
        if [ ! -z ${PASSPHRASE} ]; then
            printf "${PASSPHRASE}" | cryptsetup open --type luks ${PARTITION} ${PNAME}
        else
            cryptsetup open --type luks ${PARTITION} ${PNAME}
        fi
        if [ ! -e "/dev/mapper/${PNAME}" ]; then
            echo "Unable to map encrypted partition ${PARTITION} - exiting"
            exit
        fi
        PARTITION="/dev/mapper/${PNAME}"
    fi

    MOUNTCHK=$(grep "${PARTITION} " /proc/mounts | awk '{print $2}')
    if [ "${MOUNTCHK}" != "" ]; then
        TARGET=${MOUNTCHK}
    fi
    
    # Check if target is used as mount point by anything other than partition
    MOUNTCHK=$(grep " ${TARGET} " /proc/mounts | grep -v "${PARTITION}")
    if [ ! -z ${MOUNTCHK} ]; then
        echo "${TARGET} in use - exiting"
        exit 1
    fi
    
    # Create temporary dir
    if [ ! -d ${TARGET} ]; then
        mkdir -v ${TARGET}
    fi
    
    # Mounting
    mount -v ${PARTITION} ${TARGET}
else
    # No partition but directory is provided
    TARGET=${PARTITION}
fi

if ! $(ls ${TARGET}{/run,/sys,/proc,/dev} >/dev/null 2>&1); then
    echo "Missing ${TARGET}/{dev,proc,sys,run} - exiting"
    if [ ${TARGET} == "/target" ]; then
        umount -v ${TARGET}
        rmdir -v ${TARGET}
    fi
    exit 2
fi

TMP=$(mktemp)
trap 'rm -f ${TMP}' EXIT
chmod u+x "${TMP}"

set -e

if [ -z "${COMMANDS}" ]; then
    X='-x'
fi

# Create bash to mount temporary API filesystems
cat > "${TMP}"  <<END
#!/bin/bash
set -e ${X}
: Entered private mount namespace
if [ -h ${TARGET}/dev/shm ]; then mkdir -p ${TARGET}$(readlink ${TARGET}/dev/shm); fi
if [ -h ${TARGET}/var/lock ]; then mkdir -p ${TARGET}$(readlink ${TARGET}/var/lock); fi
mount -t devtmpfs devtmpfs ${TARGET}/dev
mount -t devpts devpts ${TARGET}/dev/pts
mount -t tmpfs tmpfs ${TARGET}/dev/shm
mount -t proc proc ${TARGET}/proc
mount -t sysfs sysfs ${TARGET}/sys
if [ -d /sys/firmware/efi/efivars ] && 
   [ -d ${TARGET}/sys/firmware/efi/efivars ]; then
    mount -t efivarfs efivarfs ${TARGET}/sys/firmware/efi/efivars
fi
END

# Enable networking in chroot environment
if [ -f "/etc/resolv.conf" ] && [ ! -L "${TARGET}/etc/resolv.conf" ]; then
    if [ -f "${TARGET}/etc/resolv.conf" ]; then
        mv -f "${TARGET}/etc/resolv.conf" "${TARGET}/etc/resolv.conf.bak"
    fi
    cp -f "/etc/resolv.conf" "${TARGET}/etc/resolv.conf"
fi

if [ -z "${COMMANDS}" ]; then
    echo chroot "${TARGET}" >> "$TMP"
else
    echo chroot "${TARGET}" "${COMMANDS}" >> "$TMP"
fi

# Run program in new namespaces
unshare -m -- "${TMP}"

if [ -f "${TARGET}/etc/resolv.conf.bak" ]; then
    mv -f "${TARGET}/etc/resolv.conf.bak" "${TARGET}/etc/resolv.conf"
fi

umount -q ${TARGET}
if [ -z "$(mount | grep ${TARGET})" ]; then
   rm -r ${TARGET}
fi
