#!/usr/bin/env bash
# Copyright (c) 2023 Microsoft Corporation
#
# SPDX-License-Identifier: Apache-2.0

set -o errexit
set -o nounset
set -o pipefail

if [ -n "${K8S_TEST_POLICY_FILES:-}" ]; then
	K8S_TEST_POLICY_FILES=($K8S_TEST_POLICY_FILES)
else
	K8S_TEST_POLICY_FILES=( \
		"allow-all.rego" \
		"allow-all-except-exec-process.rego" \
    )
fi

kubernetes_dir=$(dirname "$(readlink -f "$0")")
source "${kubernetes_dir}/../../common.bash"

reset_workloads_work_dir() {
	rm -rf ${kubernetes_dir}/runtimeclass_workloads_work
	cp -R ${kubernetes_dir}/runtimeclass_workloads ${kubernetes_dir}/runtimeclass_workloads_work
	copy_test_policy_files
}

copy_test_policy_files() {
	local kata_opa_dir="${kubernetes_dir}/../../../src/kata-opa"

	for policy_file in ${K8S_TEST_POLICY_FILES[@]}
	do
		cp "${kata_opa_dir}/${policy_file}" ${kubernetes_dir}/runtimeclass_workloads_work/
	done
}

add_kernel_initrd_annotations_to_yaml() {
	local yaml_file="$1"
	local mariner_kernel_path="/usr/share/cloud-hypervisor/vmlinux.bin"
	local mariner_initrd_path="/opt/kata/share/kata-containers/kata-containers-initrd-mariner.img"
	local resource_kind="$(yq read ${yaml_file} kind)"

	case "${resource_kind}" in

	Pod)
		echo "Adding kernel and initrd annotations to ${resource_kind} from ${yaml_file}"
		yq write -i "${K8S_TEST_YAML}" 'metadata.annotations[io.katacontainers.config.hypervisor.kernel]' "${mariner_kernel_path}"
		yq write -i "${K8S_TEST_YAML}" 'metadata.annotations[io.katacontainers.config.hypervisor.initrd]' "${mariner_initrd_path}"
		;;

	Deployment|Job|ReplicationController)
		echo "Adding kernel and initrd annotations to ${resource_kind} from ${yaml_file}"
		yq write -i "${K8S_TEST_YAML}" 'spec.template.metadata.annotations[io.katacontainers.config.hypervisor.kernel]' "${mariner_kernel_path}"
		yq write -i "${K8S_TEST_YAML}" 'spec.template.metadata.annotations[io.katacontainers.config.hypervisor.initrd]' "${mariner_initrd_path}"
		;;

	List)
		echo "Issue #7765: adding kernel and initrd annotations to ${resource_kind} from ${yaml_file} is not implemented yet"
		;;

	ConfigMap|LimitRange|Namespace|PersistentVolume|PersistentVolumeClaim|RuntimeClass|Secret|Service)
		echo "Kernel and initrd annotations are not required for ${resource_kind} from ${yaml_file}"
		;;

	*)
		echo "k8s resource type ${resource_kind} from ${yaml_file} is not yet supported for kernel and initrd annotations testing"
		return 1
		;;
	esac
}

add_kernel_initrd_annotations() {
	if [[ "${KATA_HOST_OS}" = "cbl-mariner" ]]; then
		info "Add kernel and initrd annotations"
		for K8S_TEST_YAML in runtimeclass_workloads_work/*.yaml
		do
			add_kernel_initrd_annotations_to_yaml "${K8S_TEST_YAML}"
		done
	fi
}

main() {
	ensure_yq
	reset_workloads_work_dir
	add_kernel_initrd_annotations
}

main "$@"