#!/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}
		       imds <argument> ...
		       imds [<argument> ...] -- <path> ...

		  --  Treat remaining arguments as literal metadata paths.

		  <argument> :- <path> | @<alias> | <control>

		  <path>       Provider-specific metadata key path.

		  <control> :-
		    -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

		  <alias> :-
		    hostname, local-hostname, ssh-keys, userdata,
		    network-interfaces, nameservers, nics,
		    nic:<iface>[,<nic-arg> ...]

		  <iface>     Local network interface, e.g. eth0.

		  <nic-arg> :- <nic-control> | @<nic-alias> | <nic-path>

		  <nic-control> :- -e | +e | +n | +s | +t

		  <nic-path>  Metadata key path relative to the selected NIC.

		  <nic-alias> :-
		    mac, ipv4, ipv6, ipv4-net, ipv6-net, ipv4-prefix, ipv6-prefix
	EOT
	exit 0
fi

prepare_dirs || exit 1
cloud_resolve

# 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

# Allocate in the main shell so command substitutions share one temporary tree.
# Never use an inherited path as a cleanup target.
IMDS_TMP_DIR=$(private_tmpdir imds) || exit 1
_imds_cleanup() {
	local directory
	# Volume loaders remove mountpoints only after a successful unmount.
	# Keep the tree if one remains; recursive removal could cross a live mount.
	for directory in "$IMDS_TMP_DIR"/metadata-volume.*; do
		[ ! -d "$directory" ] || return 1
	done
	rm -rf "$IMDS_TMP_DIR"
}
trap '_imds_cleanup' 0
trap 'exit 1' HUP INT TERM

. "$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
	_imds_transport_init || return
	error=$(mktemp "$IMDS_TMP_DIR/transport/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" ] ||
		(
			if [ "$IMDS_ENDPOINT_CACHE" = "$TINY_CLOUD_VAR/.imds-endpoint" ]; then
				# Ordinary users can reuse the root-owned hint without writing it.
				[ "$(id -u)" = 0 ] || exit 0
				make_file "$IMDS_ENDPOINT_CACHE" || exit
			fi
			# Explicit cache overrides remain caller-managed storage.
			umask 077
			echo "$1" > "$IMDS_ENDPOINT_CACHE"
		) 2>/dev/null || :
}

_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
	private_mkdir "$IMDS_TMP_DIR/cache" || return
	IMDS_CACHE_DIR="$IMDS_TMP_DIR/cache"
}

_imds_transport_init() {
	private_mkdir "$IMDS_TMP_DIR/transport"
}

# 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
	private_mkdir "${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"
}

# Fetch first so a decoder cannot hide a failed request.
_imds_base64_userdata() {
	local data
	data=$(_imds "$IMDS_USERDATA") || return
	printf '%s' "$data" | base64 -d
}

# 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 "$@"
