#!/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
	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=""

for arg in $@; do
	case $arg in
	--layer=*)
		LAYERS="$LAYERS `echo $arg | sed -e 's/^--layer=//'`"
		;;
	--*)
		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}.vmdk"

echo 'creating a sparse copy of the image'
cp -a --sparse=always $SRC $SRC_FILE

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

echo 'creating vmdk file'
qemu-img convert -O vmdk $SRC_FILE $VMDK_FILE

# 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, monitor with: sudo kill -s USR1 `pidof dd`'
dd if=$SRC_FILE bs=64M | 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

# 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
	sync
	umount ${TMP_ROOT}${dest}${EXTRAPT}
done

# cleanup temp directory
rm -rf $TMP_ROOT $VMDK_FILE $SRC_FILE

echo 'done, no errors'