#!/bin/bash

VERBOSITY=0
TEMP_D=""
UMOUNTS=( )
QEMU_DISCONNECT=""

error() { echo "$@" 1>&2; }

Usage() {
	cat <<EOF
Usage: ${0##*/} [ options ] file cmd [ args ]

   mount a file to a temporary mount point and then
   invoke the provided cmd with args

   the temporary mountpoint will be put in an a environment variable
   named MOUNTPOINT.

   if any of the arguments are the literal string '_MOUNTPOINT_', then
   they will be replaced with the mount point. Example:
      ${0##*/} my.img chroot _MOUNTPOINT_ /bin/sh

   options:
    -v | --verbose             increase verbosity
         --read-only           use read-only mount.
    -m | --mountpoint MP       mount to directory MP rather than a temp dir
         --overlay             mount via overlayfs
    -P | --partition PARTNUM   mount partition PARTNUM (default 'auto')
                               if 'auto', then mount part 1 if image is 
                               partitioned otherwise mount image
    -p | --proc                bind mount /proc
    -s | --sys                 bind mount /sys
    -d | --dev                 bind mount /dev
         --system-mounts       bind mount /sys, /proc, /dev
         --system-resolvconf   copy host's resolvconf into /etc/resolvconf
         --format FMT          specify the format of the image.
                               default is to automatically determine
EOF
}

# umount_r(mp) : unmount any filesystems under r
#  this is useful to unmount a chroot that had sys, proc ... mounted
umount_r() {
	local p
	for p in "$@"; do
		[ -n "$p" ] || continue
		tac /proc/mounts | sh -c '
			p=$1
			didumount=0
			while read s mp t opt a b ; do
				[ "${mp}" = "${p}" -o "${mp#${p}/}" != "${mp}" ] ||
					continue
				umount "$mp" || exit 1
				didumount=1
			done
			[ $didumount -eq 1 ] || exit 1
			exit 0' umount_r "${p%/}"
		[ $? -eq 0 ] || return
	done
}

bad_Usage() { Usage 1>&2; [ $# -eq 0 ] || error "$@"; exit 1; }

disconnect_qemu() {
	[ -n "$QEMU_DISCONNECT" ] || return 0
	local out="" nbd="$QEMU_DISCONNECT"
	local pid="" pfile="/sys/block/${nbd#/dev/}/pid"
	{ read pid < "$pfile" ; } >/dev/null 2>&1
	[ -n "$pid" -a ! -d "/proc/$pid" ] && 
		error "qemu-nbd process seems to have died. was '$pid'"
	out=$(qemu-nbd --disconnect "$nbd" 2>&1) &&
		QEMU_DISCONNECT="" || {
			error "failed to disconnect $nbd";
			error "$out"
			return 1;
	}
}

do_umounts() {
	local um="" fails=0 mydir="$PWD/"
	for um in "$@"; do
		um=$(readlink -f "$um") || {
			error "WARNING: failed to get full path to '$um'";
			fails=$(($fails+1))
			continue;
		}
		[ "${mydir#${um}/}" != "${mydir}" ] && {
			error "WARNING: leaving '$mydir' to unmount $um";
			cd /
		}
		umount_r "$um" || {
			error "WARNING: unmounting filesystem at $um failed!"
			fails=$(($fails+1))
		}
	done
	return $fails
}

cleanup() {
	if [ "${#UMOUNTS[@]}" -ne 0 ]; then
		debug 2 "umounts: ${UMOUNTS[*]}"
		do_umounts "${UMOUNTS[@]}"
	fi
	disconnect_qemu
	[ -z "${TEMP_D}" -o ! -d "${TEMP_D}" ] ||
		rm --one-file-system -Rf "${TEMP_D}" ||
		error "removal of temp dir failed!"
}

debug() {
	local level="$1"; shift;
	[ "${level}" -gt "${VERBOSITY}" ] && return
	error "${@}"
}

get_image_format() {
	local img="$1" out=""
	out=$(qemu-img info "$img") &&
		out=$(echo "$out" | awk '$0 ~ /^file format:/ { print $3 }') &&
		_RET="$out"
}

get_partition() {
	# return in _RET the 'auto' partition for a image.
	# return partition number for a partitioned image
	# return 0 for unpartitioned
	# return 0 if image is partitioned, 1 if not
	local img="$1"
	out=$(LANG=C sfdisk --list -uS "$img" 2>&1) || {
		error "failed determining if partitioned: $out";
		return 1;
	}
	if echo "$out" | grep -q 'Device.*Start.*End'; then
		_RET=1
	else
		_RET=0
	fi
}

mount_callback_umount() {
	local img_in="$1" dev="" out="" mp="" ret="" img="" readonly=""
	local opts="" bmounts="" system_resolvconf=false ptnum=auto
	local cd_mountpoint=false fmt="" mp_is_tmp=false overlay=false
	local img_mp="" workd=""

	short_opts="CdhmPpsv"
	long_opts="cd-mountpoint,dev,help,format:,mountpoint:,overlay,partition:,proc,read-only,sys,system-mounts,system-resolvconf,verbose"
	getopt_out=$(getopt --name "${0##*/}" \
		--options "${short_opts}" --long "${long_opts}" -- "$@") &&
		eval set -- "${getopt_out}" ||
		{ bad_Usage; return 1; }

	while [ $# -ne 0 ]; do
		cur=${1}; next=${2};
		case "$cur" in
			-C|--cd-mountpoint) cd_mountpoint=true;;
			-d|--dev) bmounts="${bmounts:+${bmounts} }/dev";;
			   --format) fmt=$next;;
			-h|--help) Usage ; exit 0;;
			-m|--mountpoint) mp=$next;;
			-P|--partition) ptnum=$next;;
			-O|--overlay) overlay=true;;
			-p|--proc) bmounts="${bmounts:+${bmounts} }/proc";;
			-s|--sys) bmounts="${bmounts:+${bmounts} }/sys";;
			   --system-mounts) bmounts="/dev /proc /sys";;
			   --system-resolvconf) system_resolvconf=true;;
			-v|--verbose) VERBOSITY=$((${VERBOSITY}+1));;
			   --opts) opts="${opts} $next"; shift;;
			   --read-only) readonly="ro";;
			--) shift; break;;
		esac
		shift;
	done

	[ $# -ge 2 ] || { bad_Usage "must provide image and cmd"; return 1; }

	[ -n "$readonly" ] && { $system_resolvconf && ! $overlay; } && {
		error "--read-only is incompatible with system-resolvconf";
		error "maybe try with --overlay"
		return 1;
	}

	img_in="$1"
	shift 1

	img=$(readlink -f "$img_in") ||
		{ error "failed to get full path to $img_in"; return 1; }

	[ -f "$img" ] ||
		{ error "$img: not a file"; return 1; }

	[ "$(id -u)" = "0" ] ||
		{ error "sorry, must be root"; return 1; }

	trap cleanup EXIT
	TEMP_D=$(mktemp -d "${TMPDIR:-/tmp}/${0##*/}.XXXXXX") ||
		{ error "failed to make tempdir"; return 1; }
	if [ -z "$mp" ]; then
		mp="${TEMP_D}/mp"
		mkdir "$mp" || return
		mp_is_tmp=true
	else
		[ -d "$mp" ] ||
			{ error "mountpoint '$mp': not a directory"; return 1; }
		mp=$(readlink -f "$mp") || {
			error "failed to get full path to provided mountpoint";
			return 1;
		}
	fi
	if $overlay; then
		img_mp="${TEMP_D}/underlay"
		mkdir -p "$img_mp" || return
	else
		img_mp=$mp
	fi

	local cmd="" arg="" found=false
	cmd=( )
	for arg in "$@"; do
		if [ "${arg}" = "_MOUNTPOINT_" ]; then
			debug 1 "replaced string _MOUNTPOINT_ in arguments arg ${#cmd[@]}"
			arg=$mp
		fi
		cmd[${#cmd[@]}]="$arg"
	done

	if [ "${cmd[0]##*/}" = "bash" -o "${cmd[0]##*/}" = "sh" ] &&
	   [ ${#cmd[@]} -eq 0 ]; then
		debug 1 "invoking shell ${cmd[0]}"
		error "MOUNTPOINT=$mp"
	fi

	local hasqemu=false
	command -v "qemu-nbd" >/dev/null 2>&1 && hasqemu=true

	if out=$(set -f; mount -o loop${readonly:+,$readonly} $opts \
			 "$img" "$img_mp" 2>&1); then
		debug 1 "mounted simple filesystem image '$img_in'"
		UMOUNTS[${UMOUNT[@]}]="$img_mp"
	else
		if ! $hasqemu; then
			error "simple mount of '$img_in' failed."
			error "if this not a raw image, or it is partitioned"
			error "you must have qemu-nbd (apt-get install qemu-utils)"
			error "mount failed with: $out"
			return 1
		fi
	fi

	if [ "${#UMOUNTS[@]}" -eq 0 ]; then
		if [ ! -e /sys/block/nbd0 ] && ! grep -q nbd /proc/modules; then
			debug 1 "trying to load nbd module"
			modprobe nbd >/dev/null 2>&1
			udevadm settle >/dev/null 2>&1
		fi
		[ -e /sys/block/nbd0 ] || {
			error "no nbd kernel support, but simple mount failed"
			return 1;
		}

		if [ -z "$fmt" ]; then
			get_image_format "$img" && fmt="$_RET" || {
				error "failed to get image format for '$img' (try --format)"
				return 1
			}
		fi

		local f nbd="" pidfile="" pid="" roflag=""
		for f in /sys/block/nbd*; do
			[ -d "$f" -a ! -f "$f/pid" ] && nbd=${f##*/} && break
		done
		if [ -z "$nbd" ]; then
			error "failed to find an nbd device"
			return 1;
		fi
		nbd="/dev/$nbd"

		[ -n "$readonly" ] && roflag="--read-only"
		pidfile="/sys/block/${nbd##*/}/pid" pid=""
		if ! qemu-nbd $roflag "--format=$fmt" --connect "$nbd" "$img"; then
			{ read pid < "$pidfile"; } >/dev/null 2>&1 &&
				[ -d "/proc/$pid" ] &&
				QEMU_DISCONNECT="$nbd"
			error "failed to qemu-nbd connect $img to $nbd"
			return 1
		fi
		QEMU_DISCONNECT="$nbd"

		local pfile="/sys/block/${nbd#/dev/}/pid"
		if [ ! -f "$pfile" ]; then
			debug 1 "waiting on pidfile for $nbd in $pfile"
			local i=0
			while [ ! -f "$pfile" ] && i=$(($i+1)); do
				if [ $i -eq 200 ]; then
					error "giving up on pidfile $pfile for $nbd"
					return 1
				fi
				sleep .1
				debug 2 "."
			done
		fi

		local nptnum=""
		debug 1 "connected $img_in ($fmt) to $nbd. waiting for device."
		i=0
		while i=$(($i+1)):; do
			get_partition "$nbd" && nptnum="$_RET" && break
			[ $i -eq 40 ] && {
				error "gave up on $nbd"
				return 1
			}
			[ $(($i%10)) -eq 0 ] &&
				debug 1 "waiting for $nbd to be ready."
			sleep .1
		done

		udevadm settle >/dev/null 2>&1
		if [ "${ptnum}" = "auto" ]; then
			if [ "$nptnum" = "0" ]; then
				debug 1 "unpartitioned disk."
			else
				debug 1 "partitioned disk."
			fi
			ptnum=$nptnum
		fi
		if [ "$ptnum" -ne 0 ]; then
			mdev="${nbd}p${ptnum}"
		else
			mdev="${nbd}"
		fi
		i=0
		while :; do
			[ -b "$mdev" ] && break
			i=$(($i+1))
			[ $i -eq 100 ] && {
				error "gave up on waiting for $mdev"
				return 1
			}
			[ $(($i%10)) -eq 0 ] &&
				debug 1 "waiting for $mdev part=$ptnum to be ready."
			sleep .1
		done

		if ( set -f; mount ${ro:+-o ${ro}} $opts "$mdev" "$img_mp" ) &&
			UMOUNTS[${#UMOUNTS[@]}]="$img_mp"; then
			debug 1 "mounted $mdev via qemu-nbd $nbd"
		else
			local pid="" pfile="/sys/block/${nbd#/dev/}/pid"
			{ read pid < "$pfile" ; } >/dev/null 2>&1
			[ -n "$pid" -a ! -d "/proc/$pid" ] ||
				error "qemu-nbd process seems to have died. was '$pid'"

			error "failed to mount $mdev"
			return 1
		fi

	fi

	if $overlay; then
		local olayopts="lowerdir=$img_mp,upperdir=$mp"
		workdir="${TEMP_D}/workdir"
		mkdir "$workdir"
		# 3.18+ require 'workdir=' option.
		case "$(uname -r)" in
			2*|3.1[01234567]*|3.[0-9].*) :;;
			*) olayopts="${olayopts},workdir=$workdir";;
		esac
		mount -t overlayfs -o "$olayopts" "$img_mp" "$mp" || {
			error "failed mount -t overlayfs -o '$olayopts' '$img_mp' '$mp'"
			return 1;
		}
		UMOUNTS[${#UMOUNTS[@]}]="$mp"
	fi
	local bindmp=""
	for bindmp in $bmounts; do
		[ -d "$mp${bindmp}" ] || mkdir "$mp${bindmp}" ||
			{ error "failed mkdir $bindmp in mount"; return 1; }
		mount --bind "$bindmp" "$mp/${bindmp}" ||
			{ error "failed bind mount '$bindmp'"; return 1; }
	done

	if ${system_resolvconf}; then
		local rcf="$mp/etc/resolv.conf"
		debug 1 "replacing /etc/resolvconf"
		if [ -e "$rcf" -o -L "$rcf" ]; then
			local trcf="$rcf.${0##*/}.$$"
			rm -f "$trcf" &&
				mv "$rcf" "$trcf" && ORIG_RESOLVCONF="$trcf" ||
				{ error "failed mv $rcf"; return 1; }
		fi
		cp "/etc/resolv.conf" "$rcf" ||
			{ error "failed copy /etc/resolv.conf"; return 1; }
	fi

	local startwd="$PWD"
	debug 1 "invoking: MOUNTPOINT=$mp" "${cmd[@]}"

	${cd_mountpoint} && cd "$mp"
	MOUNTPOINT="$mp" "${cmd[@]}"
	ret=$?
	cd "$startwd"

	if ${system_resolvconf}; then
		local rcf="$mp/etc/resolv.conf"
		cmp --quiet "/etc/resolv.conf" "$rcf" >/dev/null ||
			error "WARN: /etc/resolv.conf changed in image!"
		rm "$rcf" &&
			{ [ -z "$ORIG_RESOLVCONF" ] || mv "$ORIG_RESOLVCONF" "$rcf"; } ||
			{ error "failed to restore /etc/resolv.conf"; return 1; }
	fi

	debug 1 "cmd returned $ret. unmounting $mp"
	do_umounts "${UMOUNTS[@]}" && UMOUNTS=( ) ||
		{ error "failed umount $img"; return 1; }

	if "$mp_is_tmp"; then
		rmdir "$img_mp" || error "WARN: failed 'rmdir $img_mp'"
	fi

	
	if [ -n "$QEMU_DISCONNECT" ]; then
		disconnect_qemu || return 1;
	fi
	return $ret
}

mount_callback_umount "$@"

# vi: ts=4 noexpandtab