#!/bin/bash

VERBOSITY=0
TEMP_D=""
UMOUNTS=( )
QEMU_DISCONNECT=""
CR=$'\n'
TAB=$'\t'

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

   supported 'file' are:
      file      : any disk format supported by qemu-nbd

   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

   Additionally, the helper program 'mchroot' will be added to the path
   and can be used effectively as 'chroot _MOUNTPOINT_':
      ${0##*/} my.img mchroot

   options:
    -v | --verbose             increase verbosity
    -h | --help                print this message.

         --read-only           use read-only mount.
    -C | --cd-mountpoint       change dir to mountpoint before executing cmd.
    -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"
	debug 1 "disconnecting $nbd"
	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/" mounts="" i=0
	mounts=( "$@" )
	for((i=${#mounts[@]}-1;i>=0;i--)); do
		um=${mounts[$i]}
		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.
	# _RET=partition number for a partitioned image
	# _RET=0 for unpartitioned
	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
}

add_bin() {
	cat > "$1" || { error "failed to write to $1"; return 1; }
	chmod 755 "$1" || { error "failed to set perms on $1"; return 1; }
}

add_helpers() {
	local d="$1"
	local umap="$1" gmap="$2"
	[ -d "$1" ] || mkdir -p "$1"
	add_bin "$d/mchroot" <<"EOF" || return 1
#!/bin/sh
exec chroot "$MOUNTPOINT" "$@"
EOF
	return
}

mount_overlay() {
	local lower="$1" upper="$2" workdir="$3"
	local olayopts="lowerdir=$lower,upperdir=$upper"
	# 3.18+ require 'workdir=' option.
	case "$(uname -r)" in
		2*|3.1[01234567]*|3.[0-9].*) :;;
		*) olayopts="${olayopts},workdir=$workdir"
			mkdir -p "$workdir" ||
				{ _ERR="Failed to create workdir '$workdir'"; return 1; }
			;;
	esac

	local cmd="" fstype="" ret="" out="" fsfile="/proc/filesystems"
	_ERR=""
	for fstype in overlay overlayfs; do
		cmd=( mount -t "$fstype" -o "$olayopts" "$lower" "$upper" )
		debug 2 "attempting '$fstype' mount with: ${cmd[*]}"
		out=$("${cmd[@]}" 2>&1)
		ret=$?
		if [ $ret -eq 0 ]; then
			debug 1 "mounted '$fstype' via $fstype: ${cmd[*]}"
			return 0
		fi
		_ERR="${_ERR}Failed [$ret]: ${cmd[*]}:${CR}"
		_ERR="${_ERR}$out${CR}"
		if [ -r "$fsfile" ] && grep -q "${TAB}${fstype}$" "$fsfile"; then
			# this failed and we have support in kernel. do not try further.
			return $ret
		fi
	done
	return $ret
}

assert_nbd_support() {
	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 "Cannot use nbd: no nbd kernel support."
		return 1;
	}
}

find_unused_nbd() {
	# return a path to an unused nbd device (/dev/nbd?)
	local f roflag=""
	for f in /sys/block/nbd*; do
		[ -d "$f" -a ! -f "$f/pid" ] &&
			{ _RET="/dev/${f##*/}"; return 0; }
	done
	error "failed to find an nbd device"
	return 1;
}

connect_nbd() {
	local img="$1" fmt="$2" ptnum="${3:-auto}" rwmode="${3:-ro}"
	local nbd="" pidfile="" pid="" cmd="" ret="" roflag="" nptnum="" i=""
	if [ "$rwmode" = "ro" ]; then
		roflag="--read-only"
	fi
	# yes, there is a race condition here.
	find_unused_nbd || return
	nbd="$_RET"
	cmd=( qemu-nbd $roflag "--format=$fmt" --connect "$nbd" "$img" )
	"${cmd[@]}" && QEMU_DISCONNECT="$nbd"
	ret=$?
	if [ $ret -ne 0 ]; then
		error "Failed [$ret]: $*"
		return $ret
	fi
	pidfile="/sys/block/${nbd#/dev/}/pid"
	if [ ! -f "$pidfile" ]; then
		debug 1 "waiting on pidfile for $nbd in $pidfile"
		i=0
		while [ ! -f "$pidfile" ] && i=$(($i+1)); do
			if [ $i -eq 200 ]; then
				error "giving up on pidfile $pidfile for $nbd"
				disconnect_qemu
				return 1
			fi
			sleep .1
			debug 2 "."
		done
	fi
	read pid < "$pidfile" && debug 2 "pid for $nbd is $pid" || {
		error "reading pid from $pidfile for $nbd failed!";
		disconnect_qemu
		return 1
	}

	debug 1 "connected $img_in ($fmt) ${rwmode} to $nbd. waiting for device."

	local out=""
	# This can fail due to udev events, but we ignore that.  We need to ensure
	# it happens for where it doesnt happen automatically (LP: #1741300)
	out=$(blockdev --rereadpt "$nbd" 2>&1) ||
		debug 1 "blockdev rereadpt $nbd failed"
	udevadm settle

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

	if [ "${ptnum}" = "auto" ]; then
		if [ "$nptnum" = "0" ]; then
			debug 1 "unpartitioned disk."
		else
			debug 1 "partitioned disk."
		fi
		ptnum=$nptnum
	else
		if [ "$nptnum" = "0" -a "$ptnum" != "0" ]; then
			error "img $img does not appear partitioned but ptnum=$ptnum provided."
			return 1
		fi
	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"
			disconnect_qemu
			return 1
		}
		[ $(($i%10)) -eq 0 ] &&
			debug 1 "waiting for $mdev part=$ptnum to be ready."
		sleep .1
	done

	_RET_NBD="$nbd"
	_RET_PT="$ptnum"
	_RET_DEV="$mdev"
}

mount_nbd() {
	local img="$1" mp="$2" fmt="$3" ptnum="$4" rwmode="${5:-rw}" opts="$6"
	if [ -z "$fmt" ]; then
		get_image_format "$img" && fmt="$_RET" || {
			error "failed to get image format for '$img' (try --format)"
			return 1
		}
	fi
	assert_nbd_support || return
	connect_nbd "$img" "$fmt" "$ptnum" "$rwmode" || return
	local ptnum="$_RET_PT" mdev="$_RET_DEV" nbd="$_RET_NBD"
	if ( set -f; mount -o "$rwmode" $opts "$mdev" "$img_mp" ); then
		debug 1 "mounted $mdev via qemu-nbd $nbd at $img_mp"
	else
		error "failed to mount $mdev"
		return 1
	fi
}

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

	short_opts="Cdhm:P:psSv"
	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;;
			   --overlay) overlay=true;;
			-p|--proc) bmounts="${bmounts:+${bmounts} }/proc";;
			-s|--sys) bmounts="${bmounts:+${bmounts} }/sys";;
			-S|--system-mounts) bmounts="/dev /proc /sys";;
			   --system-resolvconf) system_resolvconf=true;;
			-v|--verbose) VERBOSITY=$((${VERBOSITY}+1));;
			   --opts) opts="${opts} $next"; shift;;
			   --read-only) readonly=true; rwmode="ro";;
			--) shift; break;;
		esac
		shift;
	done

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

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

	img_in="$1"
	shift 1

	if [ "${img_in#lxd:}" != "${img_in}" -a ! -f "${img_in}" ]; then
		error "${img_in}: lxd is no longer supported."
		return 1;
	fi

	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
	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

	out=""
	if [ "$ptnum" = "auto" -o "$ptnum" = "0" ] &&
		out=$(set -f; mount -o "loop,$rwmode" $opts "$img" "$img_mp" 2>&1); then
		debug 1 "mounted simple fs image $rwmode in '$img_in' at $img_mp"
		UMOUNTS[${#UMOUNTS[@]}]="$img_mp"
	else
		local hasqemu=false
		command -v "qemu-nbd" >/dev/null 2>&1 && hasqemu=true
		if ! $hasqemu; then
			error "simple mount of '$img_in' failed."
			error "if this is not a simple unpartitioned raw image, then"
			error "you must have qemu-nbd (apt-get install qemu-utils)"
			if [ -n "$out" ]; then
				error "mount failed with: $out"
			fi
			return 1
		fi
		mount_nbd "$img" "$img_mp" "$fmt" "$ptnum" "$rwmode" "$opts" || return
		UMOUNTS[${#UMOUNTS[@]}]="$img_mp"
	fi

	if $overlay; then
		mount_overlay "$img_mp" "$mp" "${TEMP_D}/workdir" || {
			[ -n "${_ERR}" ] && error "${_ERR}"
			error "Unable to mount overlay filesystem.  Maybe no kernel support?"
			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; }
		UMOUNTS[${#UMOUNTS[@]}]="$mp${bindmp}"
		debug 1 "mounted $bindmp to $mp${bindmp}"
	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 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

	add_helpers "$TEMP_D/bin" "$SUBUID" "$SUBGID" || {
		error "failed to add helpers to $TEMP_D";
		return 1;
	}
	PATH="$TEMP_D/bin:$PATH"

	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 [ -n "$QEMU_DISCONNECT" ]; then
		disconnect_qemu || return 1;
	fi
	return $ret
}

mount_callback_umount "$@"

# vi: ts=4 noexpandtab