schroot

#!/bin/bash
# vi:sw=2:ts=2:noexpandtab:nowrap

chroot=
chrootsu=1
chrootuser= # autodetect
chrootuser_grant_sudo=1
chrootshell="$SHELL"
chroot_pre_cmd=

# In bash style, we do some simple transformations for other shells
chrootPS1='\[\][chroot]\[\]\w \[\]\$\[\] '
chrootPS1root='\[\][chroot]\[\]\w \[\]#\[\] '

bindmounts="/dev /sys /proc /home"
tmpmounts="/tmp"
bindroot="/orig_root"
copyfiles="/etc/resolv.conf /etc/hostname"

# Ensure the correct ~/.schrootrc is expanded after sudo $0:
[ -z "$SCHROOTRC" ] && export SCHROOTRC=~/.schrootrc
if [ -f "$SCHROOTRC" ]; then
	. "$SCHROOTRC"
fi

warn()
{
	echo "$@"
}

die()
{
	echo "$@"
	exit 1
}

die_if_root()
{
	if [ `id -u` == 0 ]; then
		die "$@"
	else
		warn "$@"
		echo "Continuing anyway in case of a false-negative caused by insufficient permissions since we aren't root yet..."
	fi
}

usage()
{
	echo "Usage: $0 [-r] [-s shell] [chroot]"
	exit 1
}

parse_args()
{
	while [ $# -ge 1 ]; do
		case "$1" in
			'-r') chrootsu=0 ;;
			'-s')	chrootshell="$2";shift;;
			*) break
		esac
		shift
	done

	if [ -n "$1" ]; then
		chroot="$1"; shift
	fi

	[ $# -ge 1 ] && usage
}

check_params()
{
	[ -z "$chroot" ] && die No chroot specified
	[ "`readlink -f "$chroot"`" = "/" ] && die Attempted to chroot to /
	[ -d "$chroot" ] || die_if_root chroot target is not a directory
	[ -x "$chroot$chrootshell" ] || die_if_root "$chrootshell not found or not executable inside chroot"
}

become_root()
{
	if [ `id -u` != 0 ]; then
		echo Becomming root...
		if ( sudo -E 2>&1 | grep 'illegal option' > /dev/null ); then
			sudo `readlink -f "$0"` "$@"
		else
			sudo -E `readlink -f "$0"` "$@"
		fi
		exit
	fi
}

setup_user()
{
	if [ "$chrootsu" != "1" ]; then
		chrootuser=
		return
	fi
	if [ -z "$chrootuser" ]; then
		if [ -n "$SUDO_USER" ]; then
			chrootuser="$SUDO_USER"
		else
			echo SUDO_USER not set - remaining as root in chroot
		fi
	fi
	echo Setting up "$chrootuser" in chroot...

	tmp=`mktemp /tmp/tmp.XXXXXXXXXX`
	for f in /etc/passwd /etc/shadow; do
		if [ ! -f "$f" -o ! -f "$chroot$f" ]; then
			echo "Skipping update to $chroot$f - file does not exist"
			continue
		fi
		echo "Updating $chrootuser in $chroot$f"
		grep -v "^$chrootuser:" "$chroot$f" > "$tmp"
		grep "^$chrootuser:" "$f" | cat "$tmp" - > "$chroot$f"
	done
	if [ $chrootuser_grant_sudo -eq 1 ]; then
		if [ -f "/etc/sudoers" -a -f "$chroot/etc/sudoers" ]; then
			echo "Granting $chrootuser sudo access in chroot"
			grep -v "^$chrootuser ALL=(ALL) ALL$" "$chroot/etc/sudoers" > "$tmp"
			echo "$chrootuser ALL=(ALL) ALL" | cat "$tmp" - > "$chroot/etc/sudoers"
		else
			echo "Skipping update to /etc/sudoers - file does not exist"
		fi
	fi
	rm "$tmp"
}

chrootinuse()
{
	ls -l /proc/*/root|grep "\s$chroot$" >/dev/null
}

setup()
{
	chrootinuse && return
	if mount | grep "$chroot$bindroot"; then
		echo Stale chroot environment detected, cleaning up...
		cleanup
	fi
	echo Setting up chroot environment...
	if [ -n "$bindroot" ]; then
		echo "Binding / to $chroot$bindroot..."
		mkdir -p "$chroot$bindroot" 2>/dev/null
		mount --rbind / "$chroot$bindroot"
	fi
	for mountpoint in $bindmounts; do
		echo "Binding $mountpoint to $chroot$mountpoint..."
		mkdir -p "$chroot$mountpoint" 2>/dev/null
		mount --rbind "$mountpoint" "$chroot$mountpoint" || die "Failed to bind mount $mountpoint"
	done
	for mountpoint in $tmpmounts; do
		echo "Mounting tmpfs on $chroot$mountpoint..."
		mkdir -p "$chroot$mountpoint" 2>/dev/null
		mount -t tmpfs nodev "$chroot$mountpoint" || die "Failed to mount $mountpoint"
	done
	for file in $copyfiles; do
		echo "Copying $file into $chroot$file..."
		cp -f "$file" "$chroot$file" || die "Failed to copy $file"
	done

	echo "Updating $chroot/etc/mtab..."
	rm $chroot/etc/mtab
	sub_bindroot=$(echo $bindroot | sed 's/\//\\\//g')
	sub_chroot=$(echo $chroot | sed 's/\//\\\//g')
	sed "s/^\(\S*\) \(\S*\) /\1 $sub_bindroot\2 /g; s/^\(\S*\) $sub_bindroot$sub_chroot/\1 /" </proc/mounts > "$chroot/etc/mtab" || die "Failed to update $chroot/etc/mtab"

	echo "Updating $chroot/etc/hosts entry for $(hostname)..."
	tmp=`mktemp /tmp/tmp.XXXXXXXXXX`
	grep -v "^127.0.0.1 $(hostname)$" "$chroot/etc/hosts" > "$tmp"
	echo "127.0.0.1 $(hostname)" | cat "$tmp" - > "$chroot/etc/hosts"
	rm $tmp
}

cleanup()
{
	if chrootinuse; then
		echo chroot still in use, not cleaning up.
		return
	fi
	echo "Cleaning up chroot environment..."
	for mountpoint in $(echo $bindmounts $tmpmounts $bindroot | tr ' ' '\n' | tac); do
		echo "Unmounting $chroot$mountpoint..."
		umount -l "$chroot$mountpoint" || warn "Failed to unmount $mountpoint"
		# In case it has been mounted multiple times:
		while umount -l "$chroot$mountpoint" 2>/dev/null;do true;done
	done
}

_sanitise_PS1_other()
{
	# - Remove the '\[' and '\]' bashisms
	# - Replace '\w' with '$PWD'
	chrootPS1=$(echo "$chrootPS1" | sed 's/\\\[//g; s/\\\]//g; s/\\w/$PWD/g')
	chrootPS1root=$(echo "$chrootPS1root" | sed 's/\\\[//g; s/\\\]//g; s/\\w/$PWD/g')
}
_sanitise_PS1_zsh()
{
	# NOTE: Largely untested
	# - Replace '\[' and '\]' with zsh equivalents '%{' and '%}'
	# - Replace '\w' with '%~'
	chrootPS1=$(echo "$chrootPS1" | sed 's/\\\[/%{/g; s/\\\]/%}/g; s/\\w/%~/g')
	chrootPS1root=$(echo "$chrootPS1root" | sed 's/\\\[/%{/g; s/\\\]/%}/g; s/\\w/%~/g')
}
sanitise_PS1()
{
	s=$(basename $(readlink -f "$chroot$chrootshell"))
	case "$s" in
		'zsh')  _sanitise_PS1_zsh ;;
		'bash') ;; # Do nothing - strings are already in bash format
		*)      _sanitise_PS1_other;;
	esac
}

enter_chroot()
{
	if [ -n "$chrootuser" ]; then
		echo "Entering $chroot as $chrootuser..."
		# FIXME: This no longer gives a TTY (damn moving parts):
		# chroot "$chroot" su $chrootuser -l -s "$chrootshell" -c "env PS1='$chrootPS1' SUDO_PS1='$chrootPS1root' $chrootshell"
		# Workaround:
		$chroot_pre_cmd chroot "$chroot" su $chrootuser -l -s "$chrootshell"
	else
		echo "Entering $chroot..."
		# chroot "$chroot" su -l -c "env PS1='$chrootPS1root' $chrootshell"
		$chroot_pre_cmd chroot "$chroot" su -l
	fi
}

parse_args "$@"
check_params
become_root "$@"
sanitise_PS1
trap cleanup EXIT
setup
setup_user
enter_chroot