#!/bin/sh
# vim:set filetype=sh:

# Tiny Cloud - Instance MetaData Service client

### configuration and connection setup

: "${PREFIX:=/usr}"
: "${LIBDIR:=$PREFIX/lib}"
. "$LIBDIR/tiny-cloud/common"

if [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
	cat <<-EOT
		Usage: imds [-h|--help] { -e | +e | -E | +E | -A | +A | +n | +s | +t | @<alias> | <imds-path> | -- <imds-path> ... } ...
		  -h           : help
		  -e / +e      : ignore / catch errors
		  -E / +E      : hide / show STDERR
		  -A / +A      : hide / show errors from every request attempt
		  +n / +s / +t : insert newline / space / tab
		  --           : treat remaining arguments as literal IMDS paths
		  <alias> :-
		    hostname       : instance hostname
		    local-hostname : instance local hostname
		    ssh-keys       : instance SSH keys
		    userdata       : instance user data
		    network-interfaces : interfaces(5) configuration
		    nameservers    : DNS nameservers
		    nics           : instance NICs
		    nic:<iface>[,<nic-key> ...] : specific NIC interface
		      <iface>   : network interface (i.e. eth1)
		      <nic-key> :- { -e | +e | +n | +s | +t | @<nic-alias> | <nic-path> }
		        <nic-alias> :-
		          mac         : mac address
		          ipv4        : ipv4 address(es)
		          ipv6        : ipv6 address(es)
		          ipv4-net    : subnet ipv4 network(s)
		          ipv6-net    : subnet ipv6 network(s)
		          ipv4-prefix : delegated ipv4 CIDR(s)
		          ipv6-prefix : delegated ipv6 CIDR(s)
	EOT
	exit 0
fi

# Provider mappings and transient state.
unset \
	IMDS_HEADER \
	IMDS_QUERY \
	IMDS_HOSTNAME \
	IMDS_LOCAL_HOSTNAME \
	IMDS_SSH_KEYS \
	IMDS_USERDATA \
	IMDS_NETWORK_INTERFACES \
	IMDS_NICS \
	IMDS_MAC \
	IMDS_IPV4 \
	IMDS_IPV6 \
	IMDS_IPV4_NET \
	IMDS_IPV6_NET \
	IMDS_IPV4_PREFIX \
	IMDS_IPV6_PREFIX \
	IMDS_SELECTED_SOURCE \
	IMDS_SELECTED_FORMAT \
	IMDS_DOCUMENT_DIR

if [ ! -d "$LIBDIR/tiny-cloud/cloud/$CLOUD" ]; then
	echo "ERROR: Unknown Cloud '$CLOUD'" >&2
	exit 1
fi

. "$LIBDIR/tiny-cloud/imds-connection"
load_imds_connection

# Defaults may be overridden in /etc/tiny-cloud.conf
: "${IMDS_ENDPOINT_CACHE:=$TINY_CLOUD_VAR/.imds-endpoint}"
: "${IMDS_ENDPOINT_WAIT_ATTEMPTS:=10}"

: "${IMDS_NAMESERVERS:=}"

### internal helpers

_imds_unsupported() {
	echo "imds: '$1' is not supported by '$CLOUD'" >&2
	return 1
}

_wget() {
	local error rv
	error=$(mktemp "${TMPDIR:-/tmp}/wget.XXXXXX") || return 1
	wget --quiet --server-response --timeout 1 --output-document - \
		"$@" 2>"$error"
	rv=$?
	WGET_ERROR="$(cat "$error")"
	WGET_STATUS="$(
		awk '$1 ~ /^HTTP\/[0-9.]+$/ { status=$2 }
			END { print status }' "$error"
	)"
	rm -f "$error"
	return "$rv"
}

_imds_endpoints() {
	local cached e
	cached="$(cat "$IMDS_ENDPOINT_CACHE" 2>/dev/null)" || :
	for e in $IMDS_ENDPOINTS; do
		[ "$e" = "$cached" ] && echo "$e"
	done
	for e in $IMDS_ENDPOINTS; do
		[ "$e" != "$cached" ] && echo "$e"
	done
}

_imds_cache_endpoint() {
	local cached
	cached="$(cat "$IMDS_ENDPOINT_CACHE" 2>/dev/null)" || :
	[ "$1" = "$cached" ] ||
		echo "$1" 2>/dev/null > "$IMDS_ENDPOINT_CACHE" || :
}

_imds_http() {
	local path="$1" endpoint response_error transport_error
	shift

	for endpoint; do
		IMDS_ENDPOINT="$endpoint"
		IMDS_CURRENT_ENDPOINT="$endpoint"
		WGET_ERROR=
		WGET_STATUS=
		if [ -n "$IMDS_OUTPUT" ]; then
			_wget --header "$(_imds_header)" "http://$endpoint/$path" \
				> "$IMDS_OUTPUT"
		else
			_wget --header "$(_imds_header)" "http://$endpoint/$path"
		fi && {
			_imds_cache_endpoint "$endpoint"
			return 0
		}
		case "$WGET_STATUS" in
			4??)
				printf '%s\n' "$WGET_ERROR" >&2
				return 1
				;;
			'') transport_error="$WGET_ERROR" ;;
			*) response_error="$WGET_ERROR" ;;
		esac
	done
	[ -z "${response_error:-$transport_error}" ] ||
		printf '%s\n' "${response_error:-$transport_error}" >&2
	return 1
}

_imds_cache_init() {
	[ -z "$IMDS_CACHE_DIR" ] || return 0
	IMDS_CACHE_DIR=$(mktemp -d "$TINY_CLOUD_RUN/imds.XXXXXX") || return
	trap 'rm -rf "$IMDS_CACHE_DIR"' 0
	trap 'exit 1' HUP INT TERM
}

# Fetch command receives the destination filename as its last argument.
# Callers validate the cache path before using this helper.
_imds_fetch_cached() {
	local file="$1" tmpfile
	shift

	[ ! -f "$file" ] || return 0
	mkdir -p "${file%/*}" || return
	tmpfile=$(mktemp "$file.XXXXXX") || return
	"$@" "$tmpfile" && mv "$tmpfile" "$file" || {
		rm -f "$tmpfile"
		return 1
	}
}

_imds_nameservers_default() {
	local found server

	[ -n "${DEFAULT_NAMESERVERS:-}" ] || return 1
	for server in $DEFAULT_NAMESERVERS; do
		case "$server" in
			*:*) ip -6 route get "$server" >/dev/null 2>&1 || continue;;
			*) ip route get "$server" >/dev/null 2>&1 || continue;;
		esac
		printf '%s\n' "$server"
		found=1
	done
	if [ -n "$found" ]; then
		return 0
	fi

	# A network that is temporarily down provides no useful family signal.
	# Keep every documented default for the permanent configuration to use.
	for server in $DEFAULT_NAMESERVERS; do
		printf '%s\n' "$server"
	done
}

### default provider hooks

# Prepare mappings before each query; repeat calls must be safe.
# No stdout; return nonzero on failure.
_imds_init() { :; }

# Print a request header, if required.
_imds_header() { :; }

# Query a provider-relative metadata path and print its data.
# The HTTP implementation also accepts an internal destination filename.
_imds() {
	local endpoint endpoints attempts=1 delay IMDS_OUTPUT="$2"

	# Native document selectors are resolved locally. A destination argument
	# fetches the document itself through the same route and retry handling.
	if [ -n "$IMDS_DOCUMENT_DIR" ] && [ "$#" -eq 1 ]; then
		_imds_document "$1"
		return
	fi

	while :; do
		endpoints=
		if [ "$IMDS_ENDPOINT_WAIT_ATTEMPTS" -gt 0 ]; then
			for endpoint in $(_imds_endpoints); do
				imds_target_has_route "$endpoint" &&
					endpoints="$endpoints $endpoint"
			done
		fi

		if [ "$IMDS_ENDPOINT_WAIT_ATTEMPTS" -le 0 ] ||
				[ "$attempts" -ge "$IMDS_ENDPOINT_WAIT_ATTEMPTS" ]; then
			[ -n "$endpoints" ] || endpoints="$(_imds_endpoints)"
			_imds_http "$IMDS_URI/$1$IMDS_QUERY" $endpoints
			return
		fi

		if [ -n "$endpoints" ]; then
			if [ "${attempt_stderr:-0}" -eq 1 ]; then
				_imds_http "$IMDS_URI/$1$IMDS_QUERY" \
					$endpoints && return
			else
				_imds_http "$IMDS_URI/$1$IMDS_QUERY" \
					$endpoints 2>/dev/null && return
			fi
			case "$WGET_STATUS" in
				4??)
					[ "${attempt_stderr:-0}" -eq 1 ] ||
						printf '%s\n' "$WGET_ERROR" >&2
					return 1
					;;
			esac
			delay=0.1
		else
			delay=1
		fi

		sleep "$delay"
		attempts=$((attempts + 1))
	done
}

_imds_ssh_keys() { _imds_unsupported @ssh-keys; }

# Print decoded user data.
_imds_userdata() {
	[ -n "$IMDS_USERDATA" ] || {
		_imds_unsupported @userdata
		return 1
	}
	_imds "$IMDS_USERDATA"
}

# Print interfaces(5) configuration; fail if unavailable.
_imds_network_interfaces() {
	[ -n "$IMDS_NETWORK_INTERFACES" ] || return 1
	_imds "$IMDS_NETWORK_INTERFACES"
}

# Print nameservers, one per line; empty/failure uses defaults.
_imds_nameservers_metadata() {
	[ -n "$IMDS_NAMESERVERS" ] || return 1
	_imds "$IMDS_NAMESERVERS"
}

_imds_nics() {
	[ -n "$IMDS_NICS" ] || {
		_imds_unsupported @nics
		return 1
	}
	_imds "$IMDS_NICS"
}

_imds_nic_index() { _imds_unsupported "@nic:$1"; }

# Read a NIC value using its metadata identifier, original key, and mapping.
_imds_nic_value() { _imds "$IMDS_NICS/$1/$3"; }

### load provider

. "$LIBDIR/tiny-cloud/cloud/$CLOUD/imds"

### common query handling

# Run a command under the dispatcher's err/stderr flags. Keep its actual status
# in the caller's local rv even when errors are ignored.
_imds_call() {
	if [ "$stderr" -eq 1 ]; then
		"$@"
	else
		"$@" 2>/dev/null
	fi
	rv=$?
	[ "$err" -eq 0 ] || return "$rv"
}

_imds_nameservers() {
	local nameservers

	nameservers="$(_imds_nameservers_metadata)" || nameservers=
	if [ -n "$nameservers" ]; then
		printf '%s\n' "$nameservers"
		return
	fi
	_imds_nameservers_default
}

_imds_nic() {
	local request="$1" keys key alias nic

	[ -n "$IMDS_NICS" ] || {
		_imds_unsupported "@nic:$request"
		return 1
	}
	case "$request" in
		*,*) keys="${request#*,}";;
		*) keys=;;
	esac

	set --
	while [ -n "$keys" ]; do
		key="${keys%%,*}"
		case "$keys" in
			*,*) keys="${keys#*,}";;
			*) keys=;;
		esac
		[ -n "$key" ] || continue

		alias="$key"
		case "$key" in
			@mac)          key="$IMDS_MAC";;
			@ipv4)         key="$IMDS_IPV4";;
			@ipv6)         key="$IMDS_IPV6";;
			@ipv4-net)     key="$IMDS_IPV4_NET";;
			@ipv6-net)     key="$IMDS_IPV6_NET";;
			@ipv4-prefix)  key="$IMDS_IPV4_PREFIX";;
			@ipv6-prefix)  key="$IMDS_IPV6_PREFIX";;
			# error/output control passthrough
			-e|+[enst])
				set -- "$@" "$key" ''
				continue
				;;
		esac
		[ -n "$key" ] || {
			_imds_unsupported "$alias"
			return 1
		}
		set -- "$@" "$alias" "$key"
	done

	# Validate mappings before a provider lookup can fetch metadata.
	nic=$(_imds_nic_index "${request%%,*}") || return 1
	[ -n "$nic" ] || return 1

	# Nested controls apply within this NIC query.
	local err=1 stderr=1

	[ "$#" -gt 0 ] || {
		_imds_nic_value "$nic" '' ''
		return $?
	}
	while [ "$#" -gt 0 ]; do
		alias="$1" key="$2"
		shift 2
		case "$alias" in
			-e) err=0;;
			+e) err=1;;
			+n) printf '\n';;
			+s) printf ' ';;
			+t) printf '\t';;
			*)
				_imds_call _imds_nic_value "$nic" "$alias" "$key" ||
					return $?
				;;
		esac
	done
}

imds() {
	local cmd args key rv err=1 stderr=1 attempt_stderr=0 options=1
	while [ -n "$1" ]; do
		cmd=_imds
		args=
		key="$1"; shift
		if [ "$options" -eq 1 ]; then
			case $key in
				--) options=0; continue ;;
				# error handling/output
				-e) err=0; continue ;;	# ignore
				+e) err=1; continue ;;	# return
				-E) stderr=0; continue ;; # hide
				+E) stderr=1; continue ;; # show
				-A) attempt_stderr=0; continue ;; # final attempt only
				+A) attempt_stderr=1; continue ;; # every attempt
				# output control
				+n) printf "\n"; continue ;;	# insert newline
				+s) printf " "; continue ;;		# insert space
				+t) printf "\t"; continue ;;	# insert tab
				-*)
					if [ "$stderr" -eq 1 ]; then
						echo "imds: unknown option '$key'" >&2
						echo "Try 'imds --help' for usage." >&2
					fi
					return 2
					;;
			esac
		fi
		_imds_call _imds_init || return $?
		[ "$rv" -eq 0 ] || continue
		if [ "$options" -eq 0 ]; then
			args="$key"
		else
			case $key in
				# key aliasing
				@hostname)			args="$IMDS_HOSTNAME" ;;
				@local-hostname)	args="$IMDS_LOCAL_HOSTNAME" ;;
				@ssh-keys)			cmd=_imds_ssh_keys ;;
				@userdata)			cmd=_imds_userdata ;;
				@network-interfaces)	cmd=_imds_network_interfaces ;;
				@nameservers)			cmd=_imds_nameservers ;;
				@nics)				cmd=_imds_nics ;;
				@nic:*)
					cmd=_imds_nic
					args="${key#@nic:}"
					;;
				# use key verbatim
				*) args="$key" ;;
			esac
		fi
		if [ "$options" -eq 1 ] && [ "$cmd" = _imds ]; then
			case "$key" in
				@hostname|@local-hostname)
					if [ -z "$args" ]; then
						_imds_call _imds_unsupported "$key" ||
							return $?
						continue
					fi
					;;
			esac
		fi
		_imds_call "$cmd" "$args" || return $?
	done
}

imds "$@"
