#!/usr/bin/env atf-sh
# vim:set filetype=sh:
# shellcheck shell=sh

. $(atf_get_srcdir)/test_env.sh

export PREFIX="$srcdir"
PROVIDERS="aws azure digitalocean gcp hetzner incus oci nocloud scaleway"

init_tests \
	imds_help \
	imds_space \
	imds_endpoint_fallback \
	imds_endpoint_cache \
	imds_endpoint_ipv6 \
	imds_endpoint_route_skip \
	imds_endpoint_route_wait \
	\
	imds_hostname_aws \
	imds_hostname_azure \
	imds_hostname_digitalocean \
	imds_hostname_gcp \
	imds_hostname_hetzner \
	imds_hostname_incus \
	imds_hostname_nocloud \
	imds_hostname_oci \
	imds_hostname_scaleway \
	\
	imds_local_hostname_aws \
	imds_local_hostname_azure \
	imds_local_hostname_digitalocean \
	imds_local_hostname_gcp \
	imds_local_hostname_hetzner \
	imds_local_hostname_incus \
	imds_local_hostname_nocloud \
	imds_local_hostname_oci \
	imds_local_hostname_scaleway \
	\
	imds_ssh_keys_aws \
	imds_ssh_keys_azure \
	imds_ssh_keys_digitalocean \
	imds_ssh_keys_gcp \
	imds_ssh_keys_hetzner \
	imds_ssh_keys_incus \
	imds_ssh_keys_nocloud \
	imds_ssh_keys_oci \
	imds_ssh_keys_scaleway \
	\
	imds_aws_api_version_imdsv1 \
	imds_aws_api_version_imdsv2_explicit \
	imds_aws_api_version_imdsv2_latest \
	imds_aws_token_endpoint_port \
	\
	imds_nocloud_cmdline_local_hostname \
	imds_nocloud_smbios_local_hostname \
	\
	imds_userdata_incus \
	imds_userdata_scaleway

imds_help_body() {
	atf_check -o match:"Usage: imds" imds -h
}

imds_space_body() {
	for provider in $PROVIDERS; do
		CLOUD="$provider" atf_check -o match:'^ $' imds +s
		CLOUD="$provider" atf_check -o match:'^\t$' imds +t
		CLOUD="$provider" atf_check -o match:'^$' imds +n
	done
}

imds_endpoint_fallback_body() {
	IMDS_API_VERSION=2009-04-04 CLOUD=aws fake_metadata aws <<-EOF
		hostname: myhostname
	EOF
	IMDS_API_VERSION=2009-04-04 IMDS_ENDPOINT_WAIT_ATTEMPTS=0 \
		IMDS_ENDPOINTS="fail 169.254.169.254" CLOUD=aws atf_check \
		-o match:"myhostname" \
		imds @hostname
	atf_check -o match:"^169.254.169.254$" cat var/lib/cloud/.imds-endpoint
}

imds_endpoint_cache_body() {
	mkdir -p var/lib/cloud
	echo "cached.example" > var/lib/cloud/.imds-endpoint
	cat > cached.example.yaml <<-EOF
		hostname: cached-hostname
	EOF
	cat > first.example.yaml <<-EOF
		hostname: first-hostname
	EOF
	IMDS_API_VERSION=2009-04-04 IMDS_ENDPOINT_WAIT_ATTEMPTS=0 \
		WGET_STRIP_PREFIX="/2009-04-04/meta-data" \
		WGET_HOST_LOG="$PWD/hosts.log" \
		IMDS_ENDPOINTS="first.example cached.example" CLOUD=aws atf_check \
		-o match:"cached-hostname" \
		imds @hostname
	atf_check -o match:"^cached.example$" head -n 1 hosts.log
}

imds_endpoint_ipv6_body() {
	cat > "fd00:ec2::254.yaml" <<-EOF
		hostname: ipv6-hostname
	EOF
	IMDS_API_VERSION=2009-04-04 IMDS_ENDPOINT_WAIT_ATTEMPTS=0 \
		WGET_STRIP_PREFIX="/2009-04-04/meta-data" \
		WGET_HOST_LOG="$PWD/hosts.log" \
		IMDS_ENDPOINTS="[fd00:ec2::254]" CLOUD=aws atf_check \
		-o match:"ipv6-hostname" \
		imds @hostname
	atf_check -o match:"^\\[fd00:ec2::254\\]$" cat hosts.log
}

imds_endpoint_route_skip_body() {
	IMDS_API_VERSION=2009-04-04 CLOUD=aws fake_metadata aws <<-EOF
		hostname: myhostname
	EOF
	fake_bin ip <<-'EOF'
		#!/bin/sh
		[ "$3" = 169.254.169.254 ]
	EOF
	IMDS_API_VERSION=2009-04-04 \
		IMDS_ENDPOINTS="192.0.2.1 169.254.169.254" CLOUD=aws atf_check \
		-o match:"myhostname" \
		imds @hostname
	atf_check -o match:"^169.254.169.254$" cat var/lib/cloud/.imds-endpoint
}

imds_endpoint_route_wait_body() {
	IMDS_API_VERSION=2009-04-04 CLOUD=aws fake_metadata aws <<-EOF
		hostname: myhostname
	EOF
	fake_bin ip <<-'EOF'
		#!/bin/sh
		mkdir -p tmp
		count=$(cat tmp/route-count 2>/dev/null || echo 0)
		count=$((count + 1))
		echo "$count" > tmp/route-count
		[ "$count" -gt 1 ]
	EOF
	IMDS_API_VERSION=2009-04-04 IMDS_ENDPOINT_WAIT_ATTEMPTS=5 CLOUD=aws atf_check \
		-o match:"myhostname" \
		imds @hostname
	atf_check -o match:"^2$" cat tmp/route-count
}

check_hostname() {
	fake_metadata "$1" <<-EOF
		# aws, digitalocean, hetzner, nocloud
		hostname: myhostname
		# azure
		compute:
		  name: myhostname
		# gcp, oci
		instance:
		  hostname: myhostname
	EOF
	CLOUD="$1" atf_check -o match:"myhostname" imds @hostname
}
imds_hostname_aws_body()	{ check_hostname aws; }
imds_hostname_azure_body()	{ check_hostname azure; }
imds_hostname_digitalocean_body()	{ check_hostname digitalocean; }
imds_hostname_gcp_body()	{ check_hostname gcp; }
imds_hostname_hetzner_body()	{ check_hostname hetzner; }
imds_hostname_incus_body()	{
	fake_metadata incus <<-EOF
		#cloud-config
		local-hostname: myhostname
	EOF
	CLOUD="incus" atf_check -o match:"myhostname" imds @hostname
}
imds_hostname_nocloud_body()	{ check_hostname nocloud; }
imds_hostname_oci_body()	{ check_hostname oci; }
imds_hostname_scaleway_body()	{
	fake_metadata scaleway <<-EOF
HOSTNAME=myhostname
EOF
	CLOUD="scaleway" atf_check -o match:"myhostname" imds @hostname
}

check_local_hostname() {
	fake_metadata "$1" <<-EOF
		# nocloud, aws, incus
		local-hostname: myhostname
		# digitalocean, hetzner
		hostname: myhostname
		# azure
		compute:
		  name: myhostname
		# gcp, oci
		instance:
		  hostname: myhostname
	EOF
	CLOUD="$1" atf_check -o match:"myhostname" imds @local-hostname
}
imds_local_hostname_aws_body()		{ check_local_hostname aws; }
imds_local_hostname_azure_body()	{ check_local_hostname azure; }
imds_local_hostname_digitalocean_body()	{ check_local_hostname digitalocean; }
imds_local_hostname_gcp_body()		{ check_local_hostname gcp; }
imds_local_hostname_hetzner_body()	{ check_local_hostname hetzner; }
imds_local_hostname_incus_body()	{ check_local_hostname incus; }
imds_local_hostname_nocloud_body()	{ check_local_hostname nocloud; }
imds_local_hostname_oci_body()	{ check_local_hostname oci; }
imds_local_hostname_scaleway_body()	{
	fake_metadata scaleway <<-EOF
HOSTNAME=myhostname
FOO=bar
EOF
	CLOUD="scaleway" atf_check -o match:"^myhostname$" imds @local-hostname
}

check_ssh_keys() {
	local key="ssh-ed25519 keydata"
	fake_metadata "$1" <<-EOF
		# aws, alpine, nocloud
		public-keys:
		  0=testuser:
		  0:
		    openssh-key: $key
		# azure
		compute:
		  publicKeys:
		    - keyData: $key
		# gcp
		instance:
		  attributes:
		    ssh-keys: user1:$key
		  # oci
		  metadata:
		    ssh_authorized_keys: $key

	EOF
	CLOUD="$1" atf_check -o match:"$key" imds @ssh-keys
}
imds_ssh_keys_aws_body()	{ check_ssh_keys aws; }
imds_ssh_keys_azure_body()	{ check_ssh_keys azure; }
imds_ssh_keys_digitalocean_body() {
	local key="ssh-ed25519 keydata"
	fake_metadata "digitalocean" <<-EOF
		# digitalocean
		public-keys: $key
	EOF
	CLOUD="digitalocean" atf_check -o match:"$key" imds @ssh-keys
}
imds_ssh_keys_gcp_body()	{ check_ssh_keys gcp; }
imds_ssh_keys_hetzner_body()	{
	local key="ssh-ed25519 keydata"
	fake_metadata "hetzner" <<-EOF
		# hetzner
		public-keys: '["$key"]'
	EOF
	CLOUD="hetzner" atf_check -o match:"$key" imds @ssh-keys
}
imds_ssh_keys_incus_body()	{
	local key="ssh-ed25519 keydata"
	fake_userdata_incus <<-EOF
		ssh_authorized_keys:
		  - asdasdf
		  # other key
		  - $key
		foo: bar
	EOF
	CLOUD="incus" atf_check \
		-o match:"$key" \
		-o not-match:"other key" \
		imds @ssh-keys
}
imds_ssh_keys_nocloud_body()	{ check_ssh_keys nocloud; }
imds_ssh_keys_oci_body()	{ check_ssh_keys oci; }
imds_ssh_keys_scaleway_body()	{
	local key="ssh-ed25519 keydata"
	fake_metadata scaleway <<-EOF
SSH_PUBLIC_KEYS=1
SSH_PUBLIC_KEYS_0='ID KEY FINGERPRINT CREATION_DATE MODIFICATION_DATE DESCRIPTION IP'
SSH_PUBLIC_KEYS_0_ID=00000000-0000-0000-0000-000000000000
SSH_PUBLIC_KEYS_0_KEY='$key'
SSH_PUBLIC_KEYS_0_FINGERPRINT='256 MD5:00:00:00:00:00:00:00:00:00:00:00:00:00:00:00:00 ssh_test_key (ssh-ed25519)'
SSH_PUBLIC_KEYS_0_CREATION_DATE=2024-01-25T19:51:28.263483+00:00
SSH_PUBLIC_KEYS_0_MODIFICATION_DATE=2024-01-25T19:51:28.263483+00:00
SSH_PUBLIC_KEYS_0_DESCRIPTION=ssh_test_key
SSH_PUBLIC_KEYS_0_IP=
EOF
	CLOUD="scaleway" atf_check -o match:"$key" imds @ssh-keys
}

imds_aws_api_version_imdsv1_body() {
	# IMDSv1 (2009-04-04) should not attempt to get a token
	# Set IMDS_API_VERSION before fake_metadata so WGET_STRIP_PREFIX is correct
	IMDS_API_VERSION=2009-04-04 CLOUD=aws fake_metadata aws <<-EOF
		hostname: test-imdsv1
	EOF
	# For IMDSv1, nc should not be called at all
	# If it is called, the test will fail because we're not mocking it
	IMDS_API_VERSION=2009-04-04 CLOUD=aws atf_check \
		-o match:"test-imdsv1" \
		imds @hostname
}

imds_aws_api_version_imdsv2_explicit_body() {
	# IMDSv2 with explicit version (2009-04-05 or later)
	# Verify that metadata can be retrieved with explicit API version
	# Set IMDS_API_VERSION before fake_metadata so WGET_STRIP_PREFIX is correct
	IMDS_API_VERSION=2009-04-05 CLOUD=aws fake_metadata aws <<-EOF
		hostname: test-imdsv2-explicit
	EOF
	# Mock nc to provide a token (for IMDSv2 token request)
	fake_bin nc <<-'NCEOF'
		#!/bin/sh
		cat > /dev/null
		printf "HTTP/1.0 200 OK\r\n\r\nmock-token"
	NCEOF
	IMDS_API_VERSION=2009-04-05 CLOUD=aws atf_check \
		-o match:"test-imdsv2-explicit" \
		imds @hostname
}

imds_aws_api_version_imdsv2_latest_body() {
	# IMDSv2 with 'latest' (default behavior)
	# Verify that metadata can be retrieved with latest API version
	# Set IMDS_API_VERSION before fake_metadata so WGET_STRIP_PREFIX is correct
	IMDS_API_VERSION=latest CLOUD=aws fake_metadata aws <<-EOF
		hostname: test-imdsv2-latest
	EOF
	# Mock nc to provide a token (for IMDSv2 token request)
	fake_bin nc <<-'NCEOF'
		#!/bin/sh
		cat > /dev/null
		printf "HTTP/1.0 200 OK\r\n\r\nmock-token"
	NCEOF
	IMDS_API_VERSION=latest CLOUD=aws atf_check \
		-o match:"test-imdsv2-latest" \
		imds @hostname
}

imds_aws_token_endpoint_port_body() {
	cat > "fd00:ec2::254.yaml" <<-EOF
		hostname: test-imdsv2-port
	EOF
	fake_bin nc <<-'NCEOF'
		#!/bin/sh
		while [ -n "$1" ]; do
			case "$1" in
				-w) shift 2;;
				*) echo "$1" >> nc.args; shift;;
			esac
		done
		cat > /dev/null
		printf "HTTP/1.0 200 OK\r\n\r\nmock-token"
	NCEOF
	IMDS_API_VERSION=latest WGET_STRIP_PREFIX="/latest/meta-data" \
		IMDS_ENDPOINTS="[fd00:ec2::254]:8080" CLOUD=aws atf_check \
		-o match:"test-imdsv2-port" \
		imds @hostname
	atf_check -o match:"^fd00:ec2::254 8080 $" sh -c "tr '\n' ' ' < nc.args"
}

imds_nocloud_cmdline_local_hostname_body() {
	atf_require_prog yx
	mkdir proc
	for key in h local-hostname; do
		echo "BOOT_IMAGE=/boot/vmlinuz-lts ro ds=nocloud;$key=myhostname" > proc/cmdline
		CLOUD=nocloud atf_check \
			-o match:'^myhostname$' \
			imds @local-hostname
	done
}

imds_nocloud_smbios_local_hostname_body() {
	atf_require_prog yx
	mkdir -p sys/class/dmi/id
	for key in h local-hostname; do
		echo "ds=nocloud;$key=myhostname" > sys/class/dmi/id/product_serial
		CLOUD=nocloud atf_check \
			-o match:'^myhostname$' \
			imds @local-hostname
	done
}

imds_userdata_incus_body() {
	fake_userdata_incus <<-EOF
		#alpine-config
	EOF
	CLOUD="incus" atf_check -o match:"#alpine-config" imds @userdata
}

imds_userdata_scaleway_body() {
	local cmd="sh ./cmd"
	fake_userdata_scaleway <<-EOF
$cmd
EOF
	CLOUD="scaleway" atf_check -o match:"$cmd" imds @userdata
}