#!/sbin/openrc-run
# vim:set ft=sh noet ts=4:

description="Provides EC2 cloud bootstrap"

# override in /etc/conf.d/tiny-ec2-bootstrap
EC2_USER=${EC2_USER:-alpine}
IMDS2_TOKEN_TTL=${IMDS2_TOKEN_TTL:-5}

depend() {
	need net
	provide cloud-final
}

_get_metadata_token() {
	echo -ne "PUT /latest/api/token HTTP/1.0\r\nX-aws-ec2-metadata-token-ttl-seconds: $IMDS2_TOKEN_TTL\r\n\r\n" |
	nc 169.254.169.254 80 | tail -n 1
}

_get_metadata() {
	local uri="$1"
	wget -qO - --header "X-aws-ec2-metadata-token: $(_get_metadata_token)" \
		"http://169.254.169.254/latest/$uri" 2>/dev/null
}

_update_hostname() {
	local ec2_fqdn="$(_get_metadata meta-data/hostname)"
	local short_hostname="${ec2_fqdn%%\.*}"
	echo "$short_hostname" > /etc/hostname
	hostname -F /etc/hostname
	echo -e "127.0.1.1\t$ec2_fqdn $short_hostname" >> /etc/hosts
}

_set_ssh_keys() {
	local user="$1"
	local group="$(getent passwd "$user" | cut -d: -f4)"
	local ssh_dir="$(getent passwd "$user" | cut -d: -f6)/.ssh"
	local keys_file="$ssh_dir/authorized_keys"

	if [ ! -d "$ssh_dir" ]; then
		mkdir -p "$ssh_dir"
		chmod 755 "$ssh_dir"
	fi

	[ -f "$keys_file" ] && rm "$keys_file"

	touch "$keys_file"
	chmod 600 "$keys_file"
	chown -R "$user:$group" "$ssh_dir"

	for key in $(_get_metadata meta-data/public-keys/); do
		_get_metadata "meta-data/public-keys/${key%=*}/openssh-key/" >> "$keys_file"
	done
}

_run_userdata() {
	local user_data="$(_get_metadata user-data)"
	if printf '%s' "$user_data" | head -n1 | grep -q '^#!/'; then
		printf '%s' "$user_data" >/var/lib/cloud/user-data.sh
		chmod +x /var/lib/cloud/user-data.sh

		local log_file=/var/log/cloud-bootstrap.log
		local ec_file=/var/log/cloud-bootstrap.exit

		{ /var/lib/cloud/user-data.sh 2>&1 ; echo $? >"$ec_file"; } | tee "$log_file"
		ec=$(cat "$ec_file")

		echo "User Data Script Exit Status: $ec"
		return "$ec"
	fi
}

_resize_root_partition() {
	local mountpoint="$(busybox mountpoint -n / | cut -d' ' -f1)"

	# mountpoint is the second partition...
	if echo "$mountpoint" | cut -d' ' -f1 | grep -qE '/(nvme\d+n\d+p|xvd[a-z]+)2$'; then
		local volume="$(echo "$mountpoint" | sed -Ee 's/(nvme\d+n\d+|xvd[a-z]+)p?2/\1/')"
		einfo "Expanding root partition to volume size..."
		echo ", +" | sfdisk -q --no-reread -N 2 "$volume"
		einfo "Updating kernel with new partition table..."
		partx -u "$volume"
	fi
	einfo "Resizing..."
	resize2fs "$mountpoint"
}

_lock_root_user() {
	passwd -l root
}

_disable_password() {
	echo "$1:*" | chpasswd -e
}

start() {
	# Don't bootstrap if the host has already been bootstrapped
	[ -f "/var/lib/cloud/.bootstrap-complete" ] && return 0

	[ -d "/var/lib/cloud" ] || mkdir -p /var/lib/cloud

	ebegin "Locking root account"; _lock_root_account; eend $?
	ebegin "Disabling $EC2_USER password"; _disable_password "$EC2_USER"; eend $?
	ebegin "Expanding root partition"; _resize_root_partition; eend $?
	ebegin "Setting ec2 hostname"; _update_hostname; eend $?
	ebegin "Setting ec2 user ssh keys"; _set_ssh_keys "$EC2_USER"; eend $?
	ebegin "Running ec2 user data script"; _run_userdata; eend $?

	touch "/var/lib/cloud/.bootstrap-complete"
}
