#!/bin/sh

# Ensure SSH public key and authorized_keys are set up for ultriproc

SSH_DIR="/root/.ssh"
ULTRIPROC_PUB_KEY_FILE="${SSH_DIR}/id_ultriproc_ed25519.pub"
AUTH_KEYS_FILE="${SSH_DIR}/authorized_keys"
HYPERGATOR_PUB_KEY="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAICBGvu5NDS1WBhCBKYLIDUoVsuXYL7ErkMvnpwv9agZL root@petalinux-sb-2021-sdi"
ULTRIPROC_RESTRICTED_KEY="command=\"/ross/bin/allowed_ultriproc_commands.sh\",no-pty,no-agent-forwarding,no-port-forwarding,no-X11-forwarding ${HYPERGATOR_PUB_KEY}"

# start up tcpsvd to listen on port 2222
TCPSVD_NAME="tcpsvd-ultriproc"
TCPSVD_COMMAND="/usr/bin/tcpsvd -E 192.168.12.1 2222 dropbear -R -F -i -p 192.168.12.1:2222"

start_tcpsvd_ultriproc() {
	echo "Starting $TCPSVD_NAME"
	if pgrep -f "$TCPSVD_COMMAND" > /dev/null 2>&1; then
		echo "$TCPSVD_NAME is already running"
		return 0
	fi

	$TCPSVD_COMMAND &
	echo "$TCPSVD_NAME started"
}

stop_tcpsvd_ultriproc() {
	echo "Stopping $TCPSVD_NAME"
    if pgrep -f "$TCPSVD_COMMAND" > /dev/null 2>&1; then
        pkill -f "$TCPSVD_COMMAND"
		echo "$TCPSVD_NAME stopped"
	else
		echo "$TCPSVD_NAME is not running"
	fi
}

restart_tcpsvd_ultriproc() {
	stop_tcpsvd_ultriproc
	sleep 1
	start_tcpsvd_ultriproc
}

check_and_create_pub_key() {
	if [ ! -d "$SSH_DIR" ]; then
		echo "Creating $SSH_DIR directory."
		mkdir -p "$SSH_DIR"
		chmod 700 "$SSH_DIR"
	fi

	if [ ! -f "$ULTRIPROC_PUB_KEY_FILE" ] || ! grep -Fxq "$HYPERGATOR_PUB_KEY" "$ULTRIPROC_PUB_KEY_FILE"; then
		echo "Creating or updating public key file: $ULTRIPROC_PUB_KEY_FILE"
		echo "$HYPERGATOR_PUB_KEY" > "$ULTRIPROC_PUB_KEY_FILE"
		chmod 600 "$ULTRIPROC_PUB_KEY_FILE"
	else
		echo "$ULTRIPROC_PUB_KEY_FILE is already up to date."
	fi
}

check_and_update_authorized_keys() {
	if [ ! -f "$AUTH_KEYS_FILE" ]; then
		echo "Creating authorized_keys file: $AUTH_KEYS_FILE"
		echo "$ULTRIPROC_RESTRICTED_KEY" > "$AUTH_KEYS_FILE"
		chmod 600 "$AUTH_KEYS_FILE"
	elif ! grep -Fxq "$ULTRIPROC_RESTRICTED_KEY" "$AUTH_KEYS_FILE"; then
		echo "Adding ultriproc key to authorized_keys file."
		echo "$ULTRIPROC_RESTRICTED_KEY" >> "$AUTH_KEYS_FILE"
		chmod 600 "$AUTH_KEYS_FILE"
	else
		echo "$AUTH_KEYS_FILE is already up to date."
	fi
}

case "$1" in
	start)
		echo "Checking and setting up SSH public key and authorized_keys..."
		check_and_create_pub_key
		check_and_update_authorized_keys
		start_tcpsvd_ultriproc
		;;
	stop)
		stop_tcpsvd_ultriproc
		;;
	restart|reload)
		echo "Restarting setup of SSH public key and authorized_keys..."
		check_and_create_pub_key
		check_and_update_authorized_keys
		restart_tcpsvd_ultriproc
		;;
	status)
		if [ -f "$ULTRIPROC_PUB_KEY_FILE" ]; then
			echo "ultriproc public key file exists with the following content:"
			cat "$ULTRIPROC_PUB_KEY_FILE"
		else
			echo "ultriproc public key file does not exist."
		fi
		if [ -f "$AUTH_KEYS_FILE" ]; then
			echo "authorized_keys file exists with the following content:"
			cat "$AUTH_KEYS_FILE"
		else
			echo "authorized_keys file does not exist."
		fi
		if pgrep -f "$TCPSVD_COMMAND" > /dev/null 2>&1; then
			echo "$TCPSVD_NAME is running"
		else
			echo "$TCPSVD_NAME is not running"
		fi
		;;
	*)
		echo "Usage: $0 {start|stop|restart|reload|status}"
		exit 1
		;;
esac

exit 0