#!/bin/bash
# Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

MY_VERSION="0.1.0"
MY_NAME="enclavectl"
MY_DESC="AWS Nitro Enclaves with K8s deployment tool"

source "$(dirname $(realpath $0))/scripts/common.sh"

# Configuration items
CONFIG_NAMES=(region instance_type eks_cluster_name eks_worker_node_name eks_worker_node_capacity k8s_version node_enclave_cpu_limit node_enclave_memory_limit_mib)

# Utility functions and definitions
source "$SCRIPTS_DIR/utils.sh"

# Scripts
readonly CREATE_LAUNCH_TEMPLATE="00_create_launch_template.sh"
readonly CREATE_EKS_CLUSTER="01_create_eks_cluster.sh"
readonly ENABLE_DEVICE_PLUGIN="02_enable_device_plugin.sh"
readonly BUILD_ENCLAVE_APPS="03_build_enclave_apps.sh"
readonly BUILD_IMAGE="04_build_image.sh"
readonly PUSH_IMAGE="05_push_image.sh"
readonly RUN_APP="06_run_app.sh"
readonly STOP_APP="07_stop_app.sh"
readonly CLEANUP_RESOURCES="99_cleanup_resources.sh"

USAGE="\
$MY_NAME v$MY_VERSION - $MY_DESC
Usage: $(basename "$0") <command> [arguments]

Commands:
    configure           Prepare the setup configuration
        --file              The file containing the settings for configuration (i.e. settings.json)

    setup               Setup a Nitro Enable enabled EKS cluster based on input configuration
                            - Generates a basic EC2 Launch Template for Nitro Enclaves and UserData
                            - Creates an EKS cluster with a managed node-group of configured capacity
                            - Deploys the Nitro Enclaves K8s Device plugin

    build               Build a Nitro Enclave based application for deployment
        --image             The application image name.

    push                Push the Nitro Enclaves application container to a remote auto-generated
                        private ECR repository.
        --image             The application image name

    run                 Generate the deployment specification for the Nitro Enclaves application
                        and deploy it
        --image             The application image name
	[--prepare-only]    Only generate the application deployment specification file without
	                    deploying it

    stop                Terminate the Nitro Enclaves with K8s application deployed via the 'run' command
        --image             The application image name

    cleanup             Clean up all the resources previously created via the 'setup' command
        [--force]           Ignores errors and force cleans all resources and configuration
"

# Validate number of arguments given to a function.
#
validate_arg_count() {
  local arg_count=$1; shift
  for arg in "$@"
  do
    [[ "$arg_count" == "$arg" ]] && { return; }
  done

  die "Invalid arguments. Please use \`$MY_NAME help\` for help."
}

# Print usage
#
cmd_help() {
  say "$USAGE"
}

# Ensure basic dependencies of the project are installed.
ensure_basic_deps() {
  which docker > /dev/null 2>&1
  ok_or_die "docker not found. Aborting." \
      "Please make sure you have docker installed. For more information, see" \
      "https://docs.docker.com/desktop/install/linux-install"

  which jq > /dev/null 2>&1
  ok_or_die "jq not found. Aborting." \
  "Please make sure you have jq package installed."
}

# Ensure eksctl and kubectl are available on this deployment machine
ensure_eks_deps() {
  which eksctl > /dev/null 2>&1
  ok_or_die "eksctl not found. Aborting." \
      "Please make sure you have eksctl installed. For more information, see" \
      "https://docs.aws.amazon.com/eks/latest/userguide/getting-started.html."

  which kubectl > /dev/null 2>&1
  ok_or_die "kubectl not found. Aborting." \
      "Please make sure you have kubectl installed. For more information, see" \
      "https://docs.aws.amazon.com/eks/latest/userguide/getting-started.html."
}

exec_subscript() {
  source "$1"
  shift
  main "$@"
}

apply_configuration() {
  local settings=$1
  local json_name=$2
  local ret

  truncate -s 0 "$WORKING_DIR/$FILE_CONFIGURATION"; ret=$?

  [[ $ret -eq 0 ]] && {
    for item in "${CONFIG_NAMES[@]}"
    do
      local value
      value=$(echo "$settings" | jq -r ".$item"); ret=$?
      [[ "$value" = "null" ]] && {
        say_err "$item value is not set in the $json_name file!";
        ret=$FAILURE;
        rm -f "$WORKING_DIR/$FILE_CONFIGURATION"
        break
      }
      echo "CONFIG_${item^^}=\"$value\"" >> "$WORKING_DIR/$FILE_CONFIGURATION"
    done
  }

  return $ret
}

try_create_setup_uuid() {
  local uuid_pattern='^\{?[A-Z0-9a-z]{8}-[A-Z0-9a-z]{4}-[A-Z0-9a-z]{4}-[A-Z0-9a-z]{4}-[A-Z0-9a-z]{12}\}?$'
  CONFIG_SETUP_UUID=""

  # Try loading the UUID from file.
  [[ -f $WORKING_DIR/$FILE_SETUP_ID ]] && { CONFIG_SETUP_UUID=$(<"$WORKING_DIR/$FILE_SETUP_ID"); }

  [[ "${CONFIG_SETUP_UUID}" != "" ]] || {
    say "Setup UUID doesn't exist. Creating one..."
    CONFIG_SETUP_UUID=$(<"/proc/sys/kernel/random/uuid")
  }

  # Check if the UUID is valid.
  [[ $CONFIG_SETUP_UUID =~ $uuid_pattern ]] || {
    die "Your existing configuration seems corrupted!" \
      "Run './$MY_NAME cleanup' to clean invalid setup configuration" \
      "and try restarting demo setup. If you already created some resources, you" \
      "need remove them manually."
  }

  echo "$CONFIG_SETUP_UUID" > "$WORKING_DIR/$FILE_SETUP_ID"
  ok_or_die "Cannot create session UUID file!" \
      "Please ensure that you have write access to the project folder."
  say "Using setup UUID: $CONFIG_SETUP_UUID"
}

try_load_configuration() {
  [[ -f "$WORKING_DIR/$FILE_CONFIGURATION" ]] && {
    source "$WORKING_DIR/$FILE_CONFIGURATION"

    for item in "${CONFIG_NAMES[@]}"
    do
      local var_name="CONFIG_${item^^}"
      local value=${!var_name}
      [[ $value = "" || $value = "null" ]] && {
        say_warn "The configuration seems corrupted! Ignoring existing configuration..."
        rm -f "$WORKING_DIR/$FILE_CONFIGURATION"
        return
      }
    done
  }

  [[ -f "$WORKING_DIR/$FILE_CONFIGURATION" ]] && \
      CONFIG_SETUP_UUID=$(<"$WORKING_DIR/$FILE_SETUP_ID");
}

cmd_configure() {
  validate_arg_count $# 2

  case $1 in
    --file)
      [ -f "$WORKING_DIR/$FILE_CONFIGURATION" ] && {
        say_warn "Project settings have already been configured." \
          "To apply new settings, please clean up the resources first" \
          "and try again."
        exit 0
      }
      settings_file="$2";
      ;;
    *)
      die "Invalid argument: $1. Please use \`$0 help\` for help.";;
  esac

  # Create a setup uuid. Load if it already exists.
  try_create_setup_uuid

  local settings
  settings=$(cat "$settings_file" 2> /dev/null)
  ok_or_die "Cannot open the settings file: $settings_file"

  echo "$settings" | jq '.' 2>&1 > /dev/null
  ok_or_die "Cannot parse the settings file."

  apply_configuration "$settings" "$settings_file"
  ok_or_die "Cannot create configuration from $settings_file!"

  say "Using configuration"
  echo "$settings" | jq '.'
  say "Configuration finished successfully."
}

cmd_setup() {
  validate_arg_count $# 0

  say "Running setup..."

  # Nitro Enclave Launch Template
  exec_subscript "$SCRIPTS_DIR/$CREATE_LAUNCH_TEMPLATE"
  ok_or_die "Cannot create EC2 Launch Template."

  # EKS Cluster
  exec_subscript "$SCRIPTS_DIR/$CREATE_EKS_CLUSTER"
  ok_or_die "Cannot create EKS Cluster."

  # Enable Device Plugin
  exec_subscript "$SCRIPTS_DIR/$ENABLE_DEVICE_PLUGIN"
  ok_or_die "Error while enabling the device plugin."

  say "Done."
}

cmd_build() {
  validate_arg_count $# 2

  case $1 in
    --image)
      exec_subscript "$SCRIPTS_DIR/$BUILD_ENCLAVE_APPS" "$2"
      ok_or_die "Cannot build enclave applications for $2!"

      exec_subscript "$SCRIPTS_DIR/$BUILD_IMAGE" "$2"
      ok_or_die "Cannot build docker image for $2!"
      ;;
    *)
      die "Invalid arguments. Please use \`$0 help\` for help."
  esac
}

cmd_push() {
  validate_arg_count $# 2

  case $1 in
    --image)
      exec_subscript "$SCRIPTS_DIR/$PUSH_IMAGE" "$2"
      ok_or_die "Cannot push docker image for $2!"
      ;;
    *)
      die "Invalid arguments. Please use \`$0 help\` for help."
  esac
}

cmd_run() {
  validate_arg_count $# 2 3

  local image=""
  local prepare_only=false

  while [[ $# -ge 1 ]]
  do
    case $1 in
      "--image")
        image=$2;
        shift;
        ;;
      "--prepare-only")
        prepare_only=true
        ;;
      *)
        die "Invalid arguments. Please use \`$0 help\` for help."
    esac
    shift;
  done

  exec_subscript "$SCRIPTS_DIR/$RUN_APP" "$image" "$prepare_only"
  ok_or_die "Error while running application!"
}

cmd_stop() {
  validate_arg_count $# 2

  case $1 in
    --image)
      exec_subscript "$SCRIPTS_DIR/$STOP_APP" "$2"
      ok_or_die "Error while stopping the application!"
      ;;
    *)
      die "Invalid arguments. Please use \`$0 help\` for help."
  esac
}

cmd_cleanup() {
  validate_arg_count $# 0 1
  local ignore_errors=false

  case $1 in
    --force)
      ignore_errors=true
      ;;
    "")
      ;;
    *)
      die "Invalid arguments. Please use \`$0 help\` for help."
  esac

  exec_subscript "$SCRIPTS_DIR/$CLEANUP_RESOURCES" "$ignore_errors"
  ok_or_die "Cannot clean resources due to previous errors."
}

main() {
  if [ "$#" -eq 0 ]; then
    cmd_help
    exit 1
  fi

  # Ensure basic dependencies
  ensure_basic_deps
  # Try loading applied settings.
  try_load_configuration

  local cmd="$1"
  case "$1" in
    -h|help)
     cmd_help
     exit 1
     ;;
   -c|configure)
     shift
     cmd_configure "$@"
     ;;
   *)
     declare -f "cmd_$cmd" > /dev/null
     ok_or_die "Unknown command: $1. Please use \`$MY_NAME help\` for help."

     case "$1" in
       setup|run|stop)
        ensure_eks_deps
        ;;
     esac

     [[ ! -f $WORKING_DIR/$FILE_CONFIGURATION ]] && \
       die "The demo hasn't been configured yet. Please use \`$MY_NAME help\` to know how to configure."

     cmd_"$@"
     ;;
  esac
}

main "${@}"