#!/bin/bash

DOWNLOADED=0
declare -a exit_commands

get_field () {
	fieldname=$1
	shift
	VAL=$($@|grep $fieldname|sed s/$fieldname//g|sed s/[^a-zA-Z0-9\-]//g)
	echo $VAL
}

wait_field () {
	expected=$1
	shift
	value=$(get_field $@)
	while [ "$value" != "$expected" ]; do
		sleep 5
		value=$(get_field $@)
	done
}

function image_size()
{
	if is_tarball $1; then
		tar -vtf $1 | grep lh_hdd | awk '{print $3}'
	elif is_gzip $1; then
		tar -zvtf $1 | grep lh_hdd | awk '{print $3}'
	elif is_image $1; then
		ls -l $1 | cut -f5 -d' '
	else
		echo 0
	fi
}

function bytes_to_gib()
{
	echo $1 | awk '{print int(1+$1/1073741824)}'
}

function is_tarball()
{
	file $1 | grep -q "tar archive"
}

function is_gzip()
{
	file $1 | grep -q gzip
}

function is_image()
{
	file $1 | grep -q partition
}

function exit_trap()
{
	echo "Cleaning up..."
	# Run the exit_commands in reverse order
	for (( i=${#exit_commands[@]}-1 ; i>=0 ; i-- )) ; do
		eval "${exit_commands[i]}" >/dev/null
	done
}

function on_exit()
{
	local n=${#exit_commands[*]}
	exit_commands[$n]="$*"
	if [[ $n -eq 0 ]]; then
		trap exit_trap EXIT
	fi
}

usage () {
	echo "Usage: $(basename $0) [OPTION]..."
	echo "Creates an AMI from a Lighthouse image file."
	echo
	echo "  -f FILENAME      Use the specified local file to create the image"
	echo "  -r URI           Download the image file from the specified URI"
	echo "  -d DEVICE        Attach temporary disk images to the specified device (eg, xvde)"
	echo "  -n NAME          The name to use for generated images (default: Lighthouse)"
	echo "  -h               Display this message"
	echo
}

while getopts "r:f:d:n:" o; do
	case "${o}" in
		f)
			FILENAME=$OPTARG
			;;
		r)
			# If both -r and -f are given, -r will take priority
			DOWNLOAD_LINK=$OPTARG
			;;
		d)
			DEV=$OPTARG
			;;
		n)
			NAME=$OPTARG
			;;
		h)
			usage
			exit 0
			;;
		*)
			usage
			exit 1
			;;
	esac
done

if [ -z "$NAME" ]; then
	NAME=Lighthouse
fi

if [ -z "$DEV" ]; then
	DEV=xvdg
fi

# Check inside AWS EC2 instance before downloading image
# Get data about where we are
ID=$(ec2-metadata -i | awk '{print $2}')
AZ=$(ec2-metadata -z | awk '{print $2}')
REGION=$(aws configure get region)

if [ -z "$REGION" -o -z "$AZ" -o -z "$ID" ]; then
	echo "This script must be run from inside an AWS EC2 instance."
	exit 1
fi

if [ -n "$DOWNLOAD_LINK" ]; then
	echo "Downloading image..."
	on_exit rm -f /tmp/lh.img
	wget -q -O /tmp/lh.img $DOWNLOAD_LINK
	if [ $? -ne 0 ]; then
		echo "Download failed"
		exit 1
	fi
	FILENAME=/tmp/lh.img
fi

if [ -z "$FILENAME" ]; then
	echo "No image path was specified. Either -f or -r options must be provided."
	echo
	usage
	exit 1
fi

if [ ! -f "$FILENAME" ]; then
	echo "Could not locate an image file"
	echo
	usage
	exit 1
fi

SAFENAME=$(echo $NAME|sed s/\\s/_/g)

# Check it doesn't already exist
IID=$(get_field ImageId aws ec2 describe-images --filter Name=name,Values=${SAFENAME} --region ${REGION})
if [ ! -z "$IID" ]; then
	echo "An image with that name already exists"
	exit 1
fi

########################
# Create root LH image #
########################

# Calculate image size, so a large enough volume can be created
IMAGE_SIZE=$(image_size $FILENAME)
if [[ "$IMAGE_SIZE" -eq 0 ]]; then
	echo Could not determine image size
	exit 1
fi
IMAGE_SIZE_GIB=$(bytes_to_gib $IMAGE_SIZE)
echo "Image size is ${IMAGE_SIZE} bytes (${IMAGE_SIZE_GIB} GiB)"

# Create and attach an EBS volume
echo "Creating volume..."
VID=$(get_field VolumeId aws ec2 create-volume --volume-type gp3 --availability-zone $AZ --size ${IMAGE_SIZE_GIB} --region $REGION)
on_exit aws ec2 delete-volume --volume-id $VID --region $REGION
wait_field available State aws ec2 describe-volumes --filter Name=volume-id,Values=$VID --region $REGION

echo "Attaching volume ${VID} to EC2 instance..."
STATE=$(get_field State aws ec2 attach-volume --device $DEV --instance-id $ID --volume-id $VID --region $REGION)
on_exit wait_field available State aws ec2 describe-volumes --filter Name=volume-id,Values=$VID --region $REGION
on_exit aws ec2 detach-volume --force --volume-id $VID --region $REGION
wait_field "attached in-use" State aws ec2 describe-volumes --filter Name=volume-id,Values=$VID --region $REGION
sleep 5

# Copy the image onto the volume and snapshot it
echo "Cloning image onto volume..."

if is_tarball $FILENAME; then
	tar xf $FILENAME --to-stdout | sudo dd of=/dev/$DEV bs=1MB
elif is_gzip $FILENAME; then
	tar zxf $FILENAME --to-stdout | sudo dd of=/dev/$DEV bs=1MB
elif is_image $FILENAME; then
	sudo dd if=$FILENAME of=/dev/$DEV bs=1MB
else
	exit
fi

# Wait for the snapshot to complete
echo "Creating snapshot of volume..."
SNID=$(get_field SnapshotId aws ec2 create-snapshot --description ${SAFENAME} --volume-id $VID --region $REGION)
on_exit aws ec2 delete-snapshot --snapshot-id $SNID
echo "Waiting for snapshot ${SNID} to complete..."
wait_field completed State aws ec2 describe-snapshots --filter Name=snapshot-id,Values=$SNID --region $REGION

##########################
# Create Lighthouse AMI #
##########################

echo "Creating AMI from snapshot ${SNID}..."

aws ec2 \
register-image \
--name "$SAFENAME" \
--description "$NAME" \
--architecture x86_64 \
--root-device-name "/dev/sda1" \
--virtualization-type hvm \
--region $REGION \
--ena-support \
--block-device-mappings "[
    {
        \"DeviceName\": \"/dev/sda1\",
        \"Ebs\": {
            \"SnapshotId\": \"$SNID\",
            \"VolumeType\": \"gp3\"
        }
    }
]" > /dev/null

# Don't try to cleanup the snapshot now, as it's used by the AMI
unset 'exit_commands[${#exit_commands[@]}-1]'

echo "Done!"