#!/bin/sh
set -e

# This script only supports GPT partitions for the lx-trainer image.
ROOTPT=3
EXTRAPT=4

usage_exit()
{
	echo
	echo "usage: sudo $0 src [--layer=tarball]... dest [dest]..."
	echo
	echo "example: copying image to devices: sde, sdf, sdg"
	echo "$0 lx-trainer.img /dev/sde /dev/sdf /dev/sdg"
	echo
	echo "    use --keep-files to speed up multiple calls"
	echo "    (so that vmdk and img can be re-used later)"
	echo
	exit 1
}

# minimum 2 arguments
if [ $# -lt 2 ]; then
	echo "error: invalid arguments"
	usage_exit
fi

if [ `id -u` -ne 0 ]; then
	echo "error: sorry, this script must be run as root"
	exit 1
fi

if [ -z "`which mkexfatfs`" ]; then
	echo "error: mkexfatfs program not found"
	echo
	echo "try this:"
	echo "sudo apt install exfat-utils exfat-fuse"
	exit 1
fi

if [ -z "`which qemu-img`" ]; then
	echo "error: qemu-img program not found"
	echo
	echo "try this:"
	echo "sudo apt install qemu-utils"
	exit 1
fi

# gather arguments
SRC=""
LAYERS=""
DESTS=""

KEEP_FILES=0
for arg in $@; do
	case $arg in
	--layer=*)
		LAYERS="$LAYERS `echo $arg | sed -e 's/^--layer=//'`"
		;;
	--keep-files)
		KEEP_FILES=1
		;;
	--*)
		echo "error: unknown option: $arg"
		usage_exit
		;;
	*)
		if [ -z "$SRC" ]; then
			SRC="$arg"
		else
			DESTS="$DESTS $arg"
		fi
		;;
	esac
done

# source must be file
if [ ! -f "$SRC" ]; then
	echo "error: invalid src"
	usage_exit
fi

# destination must be block device
for dest in $DESTS; do
	if [ ! -b $dest ]; then
		echo "error: invalid dest: $dest"
		usage_exit
	fi
done

# all layers must be available
for tarball in $LAYERS; do
	if [ ! -f $tarball ]; then
		echo "error: missing layer file: $tarball"
		usage_exit
	fi
done

# detect partion type of source
if [ "`env - /sbin/parted --script $SRC print | \
	grep 'Partition Table' | awk '{print $3}'`" != "gpt" ]; then
	echo "error: only gpt partition tables supported"
	usage_exit
fi

SRC_FILE="${SRC}_copy.img"
VMDK_FILE="${SRC_FILE}.vmdk"

# setup temp directory for mountpoints
TMP_ROOT="/tmp/dd-multi-`date +%s`.$$"
rm -rf $TMP_ROOT
mkdir -p $TMP_ROOT
echo "using directory $TMP_ROOT for mount points"

if [ $KEEP_FILES -eq 1 -a -f $SRC_FILE ]; then
	echo 'using existing prepared copy of the image'
	if [ -n "$LAYERS" ]; then
		echo '(SPECIFIED LAYERS WILL NOT BE RE-UNPACKED)'
	fi
else
	echo 'creating a sparse copy of the image'
	cp -a --sparse=always $SRC $SRC_FILE

	if [ -n "$LAYERS" ]; then
		LOOPDEV=`losetup --show -P -f $SRC_FILE`
		LOOPPT="${LOOPDEV}p${ROOTPT}"

		mkdir -p ${TMP_ROOT}${LOOPPT}
		mount $LOOPPT ${TMP_ROOT}${LOOPPT}

		for tarball in $LAYERS; do
			echo "unpacking layer: $tarball"
			tar -x -f $tarball --numeric-owner \
				-C ${TMP_ROOT}${LOOPPT}
		done

		sync
		umount ${LOOPPT}

		losetup -d $LOOPDEV
	fi
fi

if [ $KEEP_FILES -eq 1 -a -f $VMDK_FILE ]; then
	echo 'using existing prepared vmdk file'
else
	echo 'creating vmdk file'
	qemu-img convert -O vmdk $SRC_FILE $VMDK_FILE
fi

# delete mbr's
dd if=/dev/zero bs=10M count=1 2> /dev/null | tee $DESTS > /dev/null
partprobe $DESTS

# efficiently copy source to all destinations and rescan devices
echo 'writing image to disks (64GiB)'
dd if=$SRC_FILE bs=128M status=progress | tee $DESTS > /dev/null

# fix partition table and add extra partition
for dest in $DESTS; do
	echo fix | sudo parted ---pretend-input-tty $dest print > /dev/null 2>&1
	parted --script --align optimal -- $dest mkpart /extra ntfs 128GB 256GB
done

# rescan devices
while ! partprobe $DESTS 2> /dev/null; do
	echo "waiting for partprobe..."
	sleep 1
done
sleep 1

# wait for partitions to appear
for dest in $DESTS; do
	while [ ! -b ${dest}${ROOTPT} ]; do
		echo "waiting for ${dest}${ROOTPT}..."
		sleep 1
	done
	while [ ! -b ${dest}${EXTRAPT} ]; do
		echo "waiting for ${dest}${EXTRAPT}..."
		sleep 1
	done
done
sleep 1

# create exfat on extra partition
for dest in $DESTS; do
	mkexfatfs ${dest}${EXTRAPT} > /dev/null 2>&1
done

for dest in $DESTS; do
	mkdir -p ${TMP_ROOT}${dest}${EXTRAPT}
	mount ${dest}${EXTRAPT} ${TMP_ROOT}${dest}${EXTRAPT}
	cp -v $VMDK_FILE ${TMP_ROOT}${dest}${EXTRAPT}/training-hd.vmdk &
done
wait
sync

for dest in $DESTS; do
	umount ${TMP_ROOT}${dest}${EXTRAPT}
done

# cleanup temp directory
rm -rf $TMP_ROOT
if [ $KEEP_FILES -ne 1 ]; then
	rm -f $VMDK_FILE $SRC_FILE
else
	echo "as requested, not removing:"
	echo "     $VMDK_FILE"
	echo "     $SRC_FILE"
fi

echo 'done, no errors'