#!/bin/bash

# Copyright (c) 2024 Microsoft Corporation
#
# SPDX-License-Identifier: Apache-2.0

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

# 1. Setup: start kata-agent as an individual process listening on the default port
# 2. Use a bats script to run all tests
# 3. Collect logs if possible (look into this to validate success/failures)
# 4. handle errors/exit in cleanup to remove all mounts, etc.
# The agent process runs locally as a standalone process for now.

agent_binary="/usr/bin/kata-agent"
agent_ctl_path="/opt/kata/bin/kata-agent-ctl"

# Name of socket file used by a local agent.
agent_socket_file="/tmp/kata-agent.socket"

# Kata Agent socket URI.
local_agent_server_addr="unix://@${agent_socket_file}"

# Log file that contains agent ctl output
ctl_log_file="${PWD}/agent-ctl.log"
# Log file that contains agent output.
agent_log_file="${PWD}/kata-agent.log"

agent_log_level="debug"
keep_logs=false

local_policy_file="/opt/kata/test.rego"

cleanup()
{
	info "cleaning resources..."

	local failure_ret="$?"

	stop_agent

	sudo rm -f "${local_policy_file}"

	local sandbox_dir="/run/sandbox-ns/"
	sudo umount -f "${sandbox_dir}/uts" "${sandbox_dir}/ipc" &>/dev/null || true
	sudo rm -rf "${sandbox_dir}" &>/dev/null || true

	if [[ "${failure_ret}" -eq 0 ]] && [[ "${keep_logs}" = 'true' ]]
	then
		info "SUCCESS: Test passed, but leaving logs:"
		info ""
		info "agent log file       : ${agent_log_file}"
		info "agent-ctl log file   : ${ctl_log_file}"
		return 0
	fi

	if [[ "${failure_ret}" -ne 0 ]]; then
		warn "ERROR: Test failed"
		warn ""
		warn "Not cleaning up to help debug failure:"
		warn ""

		info "agent-ctl log file   : ${ctl_log_file}"
		info "agent log file       : ${agent_log_file}"
		return 0
	fi

	sudo rm -f \
		"${agent_log_file}" \
		"${ctl_log_file}"
}

run_agent_ctl()
{
	local cmds="${1:-}"

	[[ -n "${cmds}" ]] || die "need commands for agent control tool"

	local redirect=">> ${ctl_log_file} 2>&1"

	local server_address="--server-address ${local_agent_server_addr}"

	info "Running agent-ctl with commands: ${cmds}"

	# shellcheck disable=SC2086
	if ! eval \
		sudo \
		RUST_BACKTRACE=full \
		"${agent_ctl_path}" \
		-l debug \
		connect \
		${server_address} \
		${cmds} \
		${redirect}
	then
		warn "agent-ctl command failed. Last 50 lines of agent-ctl log:"
		tail -n 50 "${ctl_log_file}" >&2 || true
		warn "Last 50 lines of agent log:"
		tail -n 50 "${agent_log_file}" >&2 || true
		return 1
	fi
}

get_agent_pid()
{
	local pids

	local name
	name=$(basename "${agent_binary}")

	pids=$(pgrep "${name}" || true)
	[[ -z "${pids}" ]] && return 0

	local count
	count=$(echo "${pids}"|wc -l)

	[[ "${count}" -gt 1 ]] && \
		die "too many agent processes running (${count}, '${pids}')"

	echo "${pids}"
}

check_agent_alive()
{
	local cmds=()

	cmds+=("-c Check")

	run_agent_ctl \
		"${cmds[@]}"
}

wait_for_agent_to_start()
{
	local cmd="check_agent_alive"

	local wait_time_secs=20
	local sleep_time_secs=1

	info "Waiting for agent process to start.."

	waitForProcess \
		"${wait_time_secs}" \
		"${sleep_time_secs}" \
		"${cmd}"

	info "Kata agent process running."
}

stop_agent() {
	info "Stopping agent"
	local cmds=()
	cmds+=("-c DestroySandbox")
	run_agent_ctl \
		"${cmds[@]}"
}

start_agent()
{
	local log_file="${1:-}"
	[[ -z "${log_file}" ]] && die "need agent log file"

	local running
	running=$(get_agent_pid || true)

	[[ -n "${running}" ]] && die "agent already running: '${running}'"

	eval \
		sudo \
			RUST_BACKTRACE=full \
			KATA_AGENT_LOG_LEVEL="${agent_log_level}" \
			KATA_AGENT_SERVER_ADDR="${local_agent_server_addr}" \
			KATA_AGENT_POLICY_FILE="${local_policy_file}" \
			"${agent_binary}" \
			&> "${log_file}" \
			&

    wait_for_agent_to_start
}

setup_agent() {
	info "Starting a single kata agent process."

	start_agent "${agent_log_file}"

	info "Setup done."
}

# The setup attempts to start kata-agent as standalone process instead of inside a UVM.
# Hence, the contents of the UVM rootfs is not available here.
# Initing Agent policy with a relaxed `allow-all.rego`
install_policy_doc()
{
	info "Installing local policy document"

	# shellcheck disable=SC2154
	allow_all_rego_file="${repo_root_dir}/src/kata-opa/allow-all.rego"
	[[ ! -f "${allow_all_rego_file}" ]] && die "Failed to locate allow-all.rego file"

	sudo cp "${allow_all_rego_file}" "${local_policy_file}"
}

# Same reason as above, we do not have the necessary components to start the coco processes
# in this setup. So removing them before starting kata agent process
try_and_remove_coco_attestation_procs()
{
	info "Removing coco attestation process for now"
	coco_procs=( "attestation-agent" "confidential-data-hub" "api-server-rest")
	local procs_path="/usr/local/bin/"

	for i in "${coco_procs[@]}"; do
		info "Moving ${i} to /tmp"
		if [[ -f "${procs_path}${i}" ]]; then
			sudo mv "${procs_path}${i}" /tmp || true
		fi
	done
}

deny_single_api_in_policy()
{
	info "Setting default deny for single API in policy"
	local pol_file=$1
	local deny_req=$2

	[[ ! -f "${local_policy_file}" ]] && install_policy_doc

	info "Not allowing ${deny_req}"
	sudo cp "${local_policy_file}" "${pol_file}"
	sed -i "s/\(.*${deny_req}.*:= \)\(.*\)/\1false/" "${pol_file}"
}
