#!/usr/bin/env sh
set -o errexit -o noclobber -o nounset

## GCE Custom Image
#
# This script intends to be "executable documention" for the process of setting
# up a Google Compute Engine (GCE) virtual machine from a custom image.


usage() {
	cat <<-USAGE
	Usage: $(basename "${0}") [options] <command> [<command> [..]]

	This script aids in setting up Google Compute Engine with an instance
	running a custom image.

	Execute commands in order provided at the command line. The procedure
	for a fresh setup follows the order listed below in Status Commands.

	Setup Commands:
	  mkconf          Configure files on raw image partition
	  mkraw           Convert image to raw format
	  mksize          Grow image to integer multiple of GBs
	  mktar           Archive raw image
	  mkgsb           Create GS bucket

	  togsa           Upload archive to GS bucket
	  togci           Create GCE image from GS archive
	  togvm           Create GCE instance from GCE image

	Status Commands:
	  ckconf          Check contents of qcow2 image
	  ckraw           Check integrity of raw image
	  cksize          Check size of raw image
	  cktar           Check contents of tar archive
	  ckgsb           Check existence of GS bucket
	  ckgsa           Check existence of tar archive in GS bucket
	  ckgci           Check existence of GCE image
	  ckgvm           Check status of GCE instance

	  ttysro          View GCE instance's serial output
	  ttysrw          Connect to instance serial port

	Removal Commands:
	  rmgvm           Remove GCE instance
	  rmgci           Remove GCE image
	  rmgsa           Remove archive from GS bucket
	  rmgsb           Remove GS bucket

	Metainfo Commands:
	  help            Show this help message
	  vars            Print value of configuration variables

	Options:
	  -f <family>     Family name for GCE image
	  -v <version>    Version string of GCE image
	  -n <name>       Name of GCE image

	  -N <name>       Name of GCE instance
	  -s <subnet>     Subnet name of GCE instance
	  -m <type>       Machine type of GCE instance

	  -B <uri>        URI of GS bucket
	  -A <uri>        URI of GS tar file

	  -p <number>     Partition of image on which / resides
	  -S <number>     Serial port of GCE instance

	  -b <base>       Common path base for files and names
	  -g <path>       Path on image of grub.cfg
	  -i <path>       Local path of source image file
	  -r <path>       Local path of raw image file
	  -a <path>       Local path of tar file

	  -h              Show this help message
	USAGE
}

fail() {
	msg=${1}; errno=${2-1}

	>&2 printf 'Error: %s\n' "${msg}"
	>&2 usage
	exit "${errno}"
}

show_vars() {
	printf 'gce_image_family=%s\n' "${gce_image_family}"
	printf 'gce_image_version=%s\n' "${gce_image_version}"
	printf 'gce_image_name=%s\n' "${gce_image_name}"
	printf '\n'

	printf 'gce_vm_name=%s\n' "${gce_vm_name}"
	printf 'gce_vm_subnet=%s\n' "${gce_vm_subnet}"
	printf 'gce_vm_type=%s\n' "${gce_vm_type}"
	printf '\n'

	printf 'ord_partition=%s\n' "${ord_partition}"
	printf 'ord_serial_port=%s\n' "${ord_serial_port}"
	printf '\n'

	printf 'path_base=%s\n' "${path_base}"
	printf 'path_grub_cfg=%s\n' "${path_grub_cfg}"
	printf 'path_qcow2=%s\n' "${path_qcow2}"
	printf 'path_raw=%s\n' "${path_raw}"
	printf 'path_tar_gz=%s\n' "${path_tar_gz}"
	printf '\n'

	printf 'gs_uri_bucket=%s\n' "${gs_uri_bucket}"
	printf 'gs_uri_tar=%s\n' "${gs_uri_tar}"
}

size_mod_gb() {
	img=${1}
	size=$(stat -c '%s' "${img}")
	mod_gb=$((size / (1024*1024*1024)))
	rem_gb=$((size % (1024*1024*1024)))

	printf '%s %s\n' "${mod_gb}" "${rem_gb}"
}

## GCE requires image sizes to be an integer number of GBs
grow_to_ciel_gbs() {
	img=${1}

	size_mod_gb "${img}" \
	| if read -r sz_mod_gb sz_rem_gb _; then
		[ "${sz_rem_gb}" -eq 0 ] && return
		dd if=/dev/zero of="${img}" \
		   bs=1G count=0 seek=$((sz_mod_gb + 1))
	fi
}

## GCE puts a few contraints on kernel cmdline arguments
#
# Remove unsupported:    `splashimage`, `rhgb`, and `quiet`
# Enable serial console: `console=ttyS0,38400n8d`
prepare_grub_cfg() {
	grub_cfg=${1}

	sudo sed -i -f - "${grub_cfg}" <<-'SED'
		/^\s\+linux/ { s/\s\+\(rhgb\|quiet\|splashimage\)//g }
		/^\s\+linux/ {
			s/\s\+console=ttyS0,38400n8d//g
			s/$/ console=ttyS0,38400n8d/
		}
	SED
}

check_grub_cfg() {
	grub_cfg=${1}

	sudo sed -n -f - "${grub_cfg}" <<-'SED'
		/^\s\+linux.*\(rhgb\|quiet\|splashimage\)/ { p }
		/^\s\+linux/ {
			s/\s\+console=ttyS0,38400n8d/&/g
			t; F; p
		}
	SED
}

free_nbd_dev() {
	sudo modprobe -v nbd
	for nbd in /sys/class/block/nbd*; do
		[ "$(cat "${nbd}"/size)" -ne 0 ] && continue
		dev_block=$(cat "${nbd}/dev")
		readlink -vf "/dev/block/${dev_block}"
		break
	done
}

mount_image_part() {
	img=${1}; part=${2}; mnt=${3}
	dev=$(free_nbd_dev)

	sudo qemu-nbd --format=qcow2 --connect="${dev}" "${img}"
	mkdir -vp "${mnt}"
	sudo mount -v "${dev}p${part}" "${mnt}"
}

umount_image() {
	mnt=${1}
	dev=$(mount | awk "\$3 ~ \"${mnt}\" {print \$1}")

	sudo umount -v "${mnt}"
	rmdir -v "${mnt}"

	sudo qemu-nbd --disconnect "${dev}"
}

configure_files() {
	img=${1}; part=${2}; grub_cfg=${3}
	mnt=$(mktemp -d)

	mount_image_part "${img}" "${part}" "${mnt}"
	prepare_grub_cfg "${mnt}/${grub_cfg}"
	umount_image "${mnt}"
}

check_files() {
	img=${1}; part=${2}; grub_cfg=${3}
	mnt=$(mktemp -d)

	mount_image_part "${img}" "${part}" "${mnt}"
	check_grub_cfg "${mnt}/${grub_cfg}"
	umount_image "${mnt}"
}

## GCE image source expects oldgnu gzip tar with image named `disk.raw`
archive_to_tar_gz() {
	raw=${1}; tar_gz=${2}
	dir=$(mktemp -d)

	ln -vs "${raw}" "${dir}/disk.raw"
	tar --verbose --sparse --dereference --format=oldgnu --create --gzip \
		--directory "${dir}" --file "${tar_gz}" disk.raw

	rm -v "${dir}/disk.raw"
	rmdir -v "${dir}"
}

## To debug boot, we must explicitly enable serial I/O
gce_create_vm() {
	name=${1}; image=${2}; subnet=${3}; type=${4}; port=${5}

	gcloud compute instances create "${name}" \
		--image "${image}" \
		--subnet "${subnet}" \
		--machine-type "${type}"
	gcloud compute instances add-metadata "${name}" \
		--metadata "serial-port-enable=${port}"
}

while getopts ':a:A:b:B:f:g:i:m:n:N:p:r:s:S:v:h' opt "${@}"; do
	case "${opt}" in
		f) opt_gce_image_family=${OPTARG};;
		v) opt_gce_image_version=${OPTARG};;
		n) opt_gce_image_name=${OPTARG};;

		N) opt_gce_vm_name=${OPTARG};;
		s) opt_gce_vm_subnet=${OPTARG};;
		m) opt_gce_vm_type=${OPTARG};;

		B) opt_gs_uri_bucket=${OPTARG};;
		A) opt_gs_uri_tar=${OPTARG};;

		p) opt_ord_partition=${OPTARG};;
		S) opt_ord_serial_port=${OPTARG};;

		b) opt_path_base=${OPTARG};;
		g) opt_path_grub_cfg=${OPTARG};;
		i) opt_path_qcow2=${OPTARG};;
		r) opt_path_raw=${OPTARG};;
		a) opt_path_tar_gz=${OPTARG};;

		h) usage; exit;;
		:) fail "Expected argument to option -${OPTARG}";;
		*) fail "Unrecognized option: -${OPTARG}";;
	esac
done
shift $((OPTIND - 1))


gce_image_family=${opt_gce_image_family:-gce-vm}
gce_image_version=${opt_gce_image_version:-}
gce_image_name=${opt_gce_image_name:-${gce_image_family}${gce_image_version:+-}${gce_image_version}}

gce_vm_name=${opt_gce_vm_name:-${gce_image_name}}
gce_vm_subnet=${opt_gce_vm_subnet:-default}
gce_vm_type=${opt_gce_vm_type:-f1-micro}

ord_partition=${opt_ord_partition:-1}
ord_serial_port=${opt_ord_serial_port:-1}

path_grub_cfg=${opt_path_grub_cfg:-/boot/grub/grub.cfg}
path_base=${opt_path_base:-${PWD}/imgs/${gce_image_name}}
path_qcow2=${opt_path_qcow2:-${path_base}.qcow2}
path_raw=${opt_path_raw:-${path_base}.raw}
path_tar_gz=${opt_path_tar_gz:-${path_base}.tar.gz}

gs_uri_bucket=${opt_gs_uri_bucket:-gs://$(date +%s.%N)}
gs_uri_tar=${opt_gs_uri_tar:-${gs_uri_bucket}/$(basename "${path_tar_gz}")}


[ $# -gt 0 ] || fail "Expected command"

for command in "${@}"; do
	case "${command}" in
		# Setup
		mkconf) configure_files "${path_qcow2}" \
		                        "${ord_partition}" \
		                        "${path_grub_cfg}";;
		mkraw)  qemu-img convert -O raw "${path_qcow2}" \
		                                "${path_raw}";;
		mksize) grow_to_ciel_gbs "${path_raw}";;
		mktar)  archive_to_tar_gz "${path_raw}" \
		                          "${path_tar_gz}";;
		mkgsb)  gsutil mb "${gs_uri_bucket}";;
		togsa)  gsutil cp "${path_tar_gz}" "${gs_uri_bucket}";;
		togci)  gcloud compute images create "${gce_image_name}" \
		               --source-uri "${gs_uri_tar}" \
		               --family "${gce_image_family}";;
		togvm)  gce_create_vm "${gce_vm_name}" \
		                      "${gce_image_name}" \
		                      "${gce_vm_subnet}" \
		                      "${gce_vm_type}" \
		                      "${ord_serial_port}";;
		ttysro) gcloud compute instances get-serial-port-output \
		               "${gce_vm_name}" \
		               --port "${ord_serial_port}";;
		ttysrw) gcloud compute connect-to-serial-port \
		               "${gce_vm_name}" \
		               --port "${ord_serial_port}";;

		# Status
		ckconf) check_files "${path_qcow2}" \
		                    "${ord_partition}" \
		                    "${path_grub_cfg}";;
		ckraw)  qemu-img compare "${path_qcow2}" "${path_raw}";;
		cksize) size_mod_gb "${path_raw}" | awk '$2 != 0 {exit 1}';;
		cktar)  tar -tf "${path_tar_gz}" | sed '/^disk.raw$/!q1';;
		ckgsb)  gsutil du "${gs_uri_bucket}";;
		ckgsa)  gsutil stat "${gs_uri_tar}";;
		ckgci)  gcloud compute images describe "${gce_image_name}";;
		ckgvm)  gcloud compute instances describe "${gce_vm_name}";;

		# Removal
		rmgvm)  gcloud compute instances delete "${gce_vm_name}";;
		rmgci)  gcloud compute images delete "${gce_image_name}";;
		rmgsa)  gsutil rm "${gs_uri_tar}";;
		rmgsb)  gsutil rb "${gs_uri_tar}";;

		# Metainfo
		help)   usage;;
		vars)   show_vars;;

		*) fail "Unrecognized command: ${command}";;
	esac
done