#!/usr/bin/env bash
#
# Copyright (c) 2024, Cloudera, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

VAULT_APP_DATA_DIR=${VAULT_APP_DATA_DIR:-${HOME}/.vault_data}

DEFAULT_VAULT_CLUSTER_PORT=8201
DEFAULT_VAULT_VER=1.14.8

PYTHON_EXECUTABLE=python3
VAULT_START_WAIT_SECONDS=5
RUN_IN_BACKGROUND=${RUN_IN_BACKGROUND:-true}

download_vault() {
  if [[ ! -f "${VAULT_DIR}/${VAULT_ZIP}" ]] ; then
      cd $VAULT_DIR
      # download Vault
      echo "WARNING... Downloading Vault WITHOUT certs verification" | tee -a "$LOG_FILE"
      wget "https://releases.hashicorp.com/vault/${VAULT_VER}/${VAULT_ZIP}" --no-check-certificate

      if [[ $? != 0 ]] ; then
        echo "Cannot download Vault" | tee -a "$LOG_FILE"
        exit 1
      fi
  fi
}

install_vault() {
  download_vault

  echo "Extracting the binaries.. " | tee -a "$LOG_FILE"
  if [[ -f ${VAULT_BIN} ]] ; then
    rm ${VAULT_BIN}
  fi

  unzip -o ${VAULT_DIR}/${VAULT_ZIP} -d ${VAULT_DIR}
  chmod a+x ${VAULT_BIN}

  # check
  ${VAULT_BIN} --version | tee -a "$LOG_FILE"
}

create_vault_config() {
  echo 'Create vault config...' | tee -a "$LOG_FILE"
  export VAULT_SERVER_TLS VAULT_SERVER_TLS_CERT_FILE VAULT_SERVER_TLS_KEY_FILE
  export VAULT_DATA_DIR VAULT_SERVER_HOST VAULT_SERVER_PORT VAULT_ADDR VAULT_SERVER_PID_FILE
  export VAULT_DISABLE_MLOCK=$(grep -q "Locked" /proc/meminfo && echo "false" || echo "true")
  create_config_file_from_j2_template "${VAULT_CONF_FILE}"
}

initialize_vault() {
  if [ -f "${VAULT_TOKEN_FILE}" -a -s "${VAULT_TOKEN_FILE}" ]; then
    echo "Vault server has already been initialized!" | tee -a "$LOG_FILE"
  else
    VAULT_STATUS=$(${VAULT_BIN} operator init --status >/dev/null 2>&1; echo $?)
    if [ "$VAULT_STATUS" -eq 2 ]; then
      echo "Initializing vault.." | tee -a "$LOG_FILE"
      ${VAULT_BIN} operator init -key-shares=1 -key-threshold=1 -format=json > ${VAULT_TOKEN_FILE}
      chmod 400 ${VAULT_TOKEN_FILE}

      ${VAULT_BIN} operator init --status >/dev/null 2>&1
      if [ $? -ne 0 ]; then
        echo "An error occurred while initializing the vault" | tee -a "$LOG_FILE"
        exit 1
      fi
    else
      echo "[ERROR]: Vault server has already been initialized but the token file is missing or empty! Please see the troubleshooting section of the documentation!" | tee -a "$LOG_FILE"
      exit 1
    fi
  fi
}

unseal_vault() {
  VAULT_STATUS=$(${VAULT_BIN} status >/dev/null 2>&1; echo $?)
  if [ "$VAULT_STATUS" -eq 2 ]; then

    echo "Unsealing vault.." | tee -a "$LOG_FILE"
    source "${VENV_DIR}"/bin/activate
    UNSEAL_KEY=$(cat ${VAULT_TOKEN_FILE} | ${VENV_DIR}/bin/python -c "import sys, json; print(json.load(sys.stdin)['unseal_keys_b64'][0])")
    deactivate
    ${VAULT_BIN} operator unseal ${UNSEAL_KEY}

    ${VAULT_BIN} status >/dev/null 2>&1
    if [ $? -ne 0 ]; then
      echo "An error occurred while unsealing the vault! " | tee -a "$LOG_FILE"
      exit 1
    fi
  fi
  if [ ! -f "${VAULT_TOKEN_FILE_BAK}" ]; then
    backup_vault_token_file
  fi
}

init_secret_engine() {
  ROOT_TOKEN=$(get_vault_token)
  export VAULT_TOKEN=$(${VAULT_BIN} login -token-only "${ROOT_TOKEN}")
  VAULT_STATUS=$(${VAULT_BIN} secrets list | grep secret/; echo $?)
  if [ "$VAULT_STATUS" != 0 ]; then
    echo "Initializing kv secret engine" | tee -a "$LOG_FILE"
    ${VAULT_BIN} secrets enable -path=secret -version=2 kv || true
  fi
}

stop_vault() {
  if is_vault_server_running; then
    echo "Stopping vault server..." | tee -a "$LOG_FILE"
    local pid=$(get_vault_server_pid)
    kill "${pid}"

    sleep 2
    if ! is_vault_server_running; then
      echo "Done!" | tee -a "$LOG_FILE"
    else
      echo "Failed!" | tee -a "$LOG_FILE"
      exit 1
    fi

  else
    echo "Vault server is not running!" | tee -a "$LOG_FILE"
  fi

}

start_vault_server() {
  if ! is_vault_server_running; then
    create_vault_config

    echo "Starting vault server.." | tee -a "$LOG_FILE"
    CMD="${VAULT_BIN} server -config=${VAULT_CONF_FILE}"
    if [ -n "${VAULT_SERVER_TLS_KEY_PASS}" ]; then
      echo ${VAULT_SERVER_TLS_KEY_PASS} | nohup $CMD >> "${LOG_FILE}" 2>&1 &
    else
      nohup $CMD >> "${LOG_FILE}" 2>&1 &
    fi
    sleep "${VAULT_START_WAIT_SECONDS}"
    tail -n30 "${LOG_FILE}"
  fi
}

start_vault() {
  VAULT_SERVER_TLS=false
  if [ $(to_lower "${VAULT_SERVER_SCHEME}" ) = "https" ]; then

    # check if server ssl config file exists
    if [ ! -f "${VAULT_SERVER_SSL_CONF_FILE}" ]; then
      echo "ERROR: ${VAULT_SERVER_SSL_CONF_FILE} does not exist!"
      exit 1
    fi

    # check if cert file is specified
    if [ -z "${VAULT_SERVER_TLS_CERT_FILE}" ]; then
      echo "ERROR: server.ssl.cert.pem is missing from ${VAULT_SERVER_SSL_CONF_FILE}!"
      exit 1
    fi

    # check if cert file exists
    if [ ! -f "${VAULT_SERVER_TLS_CERT_FILE}" ]; then
      echo "ERROR: ${VAULT_SERVER_TLS_CERT_FILE} does not exist!"
      exit 1
    fi

    # check if key file is specified
    if [ -z "${VAULT_SERVER_TLS_CERT_FILE}" ]; then
      echo "ERROR: server.ssl.key.pem is missing from ${VAULT_SERVER_SSL_CONF_FILE}!"
      exit 1
    fi

    # check if key file exists
    if [ ! -f "${VAULT_SERVER_TLS_CERT_FILE}" ]; then
      echo "ERROR: server.ssl.cert.pem does not exist!"
      exit 1
    fi

    VAULT_SERVER_TLS=yes
    export VAULT_SKIP_VERIFY=true
  fi
  echo "TLS configured: ${VAULT_SERVER_TLS}"

  # start server
  start_vault_server

  # init
  initialize_vault

  # unseal
  unseal_vault

  # init secret engine
  init_secret_engine
}

create_vault_dirs() {
  mkdir -p ${VAULT_DIR}
  mkdir -p ${VAULT_CONF_DIR}
  mkdir -p ${VAULT_DATA_DIR}
}

get_current_vault_version() {
  local version=$(${VAULT_BIN} --version | awk -F ' ' {'print $2'})
  echo "${version:1}"
}

# if called as source code then return
[[ $_ != $0 ]] && return 0>/dev/null

# TODO:
set -x
set -e

while [[ $# -gt 0 ]]; do
  case $1 in
  start)
    ACTION_START=true
    shift # past argument with no value
    ;;
  stop)
    ACTION_STOP=true
    shift # past argument with no value
    ;;
  restart)
    ACTION_STOP=true
    ACTION_START=true
    shift # past argument with no value
    ;;
  setup)
    ACTION_INSTALL=true
    shift # past argument with no value
    ;;
  -vsc=* | --vault-server-conf=*)
    VAULT_SERVER_CONF_FILE=${1#*=}
    shift # past argument=value
    ;;
  -vsc | --vault-server-conf)
    VAULT_SERVER_CONF_FILE=$2
    shift # past argument
    shift # past value
    ;;
  -ssc=* | --vault-ssl-server-conf=*)
    VAULT_SERVER_SSL_CONF_FILE=${1#*=}
    shift # past argument=value
    ;;
  -ssc | --vault-ssl-server-conf)
    VAULT_SERVER_SSL_CONF_FILE=$2
    shift # past argument
    shift # past value
    ;;
  -h | --help)
    HELP=true
    shift # past argument with no value
    ;;
  *)
    echo "Unknown option: $1"
    HELP=true
    shift # past argument with no value
    ;;
  esac
done

__usage="
 Usage: $(basename "$0") [options]
 Options:
    start: Start the Vault Server
    stop: Stop the Vault Server
    setup: Install the Vault Server
    -vsc, --vault-server-conf: The config file of the vault-server (Default: ${VAULT_APP_DATA_DIR}/vault-server.properties)
    -ssc, --vault-ssl-server-conf: The server ssl config file (Default: ${VAULT_APP_DATA_DIR}/ssl-server.properties)
    -h,   --help: Print this help
 "

if [ "${HELP}" = true ]; then
  echo "$__usage"
  exit 0
fi

SCRIPT_PATH="$(
  cd "$(dirname "$0")" >/dev/null 2>&1 || exit
  pwd -P
)"

# source utils
source "${SCRIPT_PATH}"/utils.sh

VAULT_CONF_DIR=${VAULT_CONF_DIR:-${VAULT_APP_DATA_DIR}/conf}

# reload existing vault-server-conf file if it exists and is not specified
if [ -z "${VAULT_SERVER_CONF_FILE}" ]; then
  DEFAULT_VAULT_SERVER_CONF_FILE="${VAULT_CONF_DIR}/vault-server.properties"
  if [ -f "${DEFAULT_VAULT_SERVER_CONF_FILE}" ]; then
    echo "Since the vault-server-conf file is not specified, it uses the one in the ${VAULT_CONF_DIR} folder."
    VAULT_SERVER_CONF_FILE="${DEFAULT_VAULT_SERVER_CONF_FILE}"
  else
    echo "No vault-server-conf file found in the ${VAULT_CONF_DIR} folder."
  fi
fi

if [ -n "${VAULT_SERVER_CONF_FILE}" ]; then

  if [ ! -f "${VAULT_SERVER_CONF_FILE}" ]; then
    echo "Could not open ${VAULT_SERVER_CONF_FILE}!"
    exit 1
  fi

  echo "Loading vault-server-conf file (${VAULT_SERVER_CONF_FILE})"

  if [ -z "${VAULT_SERVER_PID_FILE}" ]; then
    VAULT_SERVER_PID_FILE=$(get_properties_conf_value "vault.server.pid.file" ${VAULT_SERVER_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_LOG_DIR}" ]; then
    VAULT_SERVER_LOG_DIR=$(get_properties_conf_value "vault.server.log.dir" ${VAULT_SERVER_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_PORT}" ]; then
    VAULT_SERVER_PORT=$(get_properties_conf_value "vault.server.port" ${VAULT_SERVER_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_SCHEME}" ]; then
    VAULT_SERVER_SCHEME=$(get_properties_conf_value "vault.server.scheme" ${VAULT_SERVER_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_HOST}" ]; then
    VAULT_SERVER_HOST=$(get_properties_conf_value "vault.server.host" ${VAULT_SERVER_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_VERSION}" ]; then
    VAULT_SERVER_VERSION=$(get_properties_conf_value "vault.server.version" ${VAULT_SERVER_CONF_FILE})
  fi

fi

# reload existing vault-ssl-server-conf file if it exists and is not specified
if [ -z "${VAULT_SERVER_SSL_CONF_FILE}" ]; then
  DEFAULT_VAULT_SERVER_SSL_CONF_FILE="${VAULT_CONF_DIR}/ssl-server.properties"
  if [ -f "${DEFAULT_VAULT_SERVER_SSL_CONF_FILE}" ]; then
    echo "Since the vault-ssl-server-conf file is not specified, it uses the one in the ${VAULT_CONF_DIR} folder."
    VAULT_SERVER_SSL_CONF_FILE="${DEFAULT_VAULT_SERVER_SSL_CONF_FILE}"
  else
    echo "No vault-ssl-server-conf file found in the ${VAULT_CONF_DIR} folder."
  fi
fi

if [ -n "${VAULT_SERVER_SSL_CONF_FILE}" ]; then

  if [ ! -f "${VAULT_SERVER_SSL_CONF_FILE}" ]; then
    echo "Could not open ${VAULT_SERVER_SSL_CONF_FILE}!"
    exit 1
  fi

  echo "Loading vault-ssl-server-conf file (${VAULT_SERVER_SSL_CONF_FILE})"

  if [ -z "${VAULT_SERVER_TLS_CERT_FILE}" ]; then
    VAULT_SERVER_TLS_CERT_FILE=$(get_properties_conf_value "server.ssl.cert.pem" ${VAULT_SERVER_SSL_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_TLS_KEY_FILE}" ]; then
    VAULT_SERVER_TLS_KEY_FILE=$(get_properties_conf_value "server.ssl.key.pem" ${VAULT_SERVER_SSL_CONF_FILE})
  fi

  if [ -z "${VAULT_SERVER_TLS_KEY_PASS}" ]; then
    VAULT_SERVER_TLS_KEY_PASS=$(get_properties_conf_value "server.ssl.key-password" ${VAULT_SERVER_SSL_CONF_FILE})
  fi

fi

source "$SCRIPT_PATH"/common.sh # load commons
create_logs_dir
LOG_FILE="${VAULT_SERVER_LOG_DIR}"/$(basename "$0" .sh).log

VAULT_VER="${VAULT_SERVER_VERSION:-${DEFAULT_VAULT_VER}}"
OS=$(uname -s |  tr '[:upper:]' '[:lower:]')
ARCH=$(uname -m | sed 's/x86_64/amd64/')
VAULT_ZIP="vault_${VAULT_VER}_${OS}_${ARCH}.zip"

VAULT_CONF_FILE=${VAULT_CONF_DIR}/vault.hcl

VAULT_SERVER_HOST=${VAULT_SERVER_HOST:-"127.0.0.1"}
VAULT_SERVER_PORT=${VAULT_SERVER_PORT:-"8210"}
VAULT_SERVER_SCHEME=${VAULT_SERVER_SCHEME:-"http"}

VAULT_CLUSTER_PORT=${VAULT_CLUSTER_PORT:-${DEFAULT_VAULT_CLUSTER_PORT}}

export VAULT_ADDR="${VAULT_SERVER_SCHEME}://${VAULT_SERVER_HOST}:${VAULT_SERVER_PORT}"

if [ -z "${ACTION_START}" -a -z "${ACTION_INSTALL}" -a -z "${ACTION_STOP}" ]; then
  ACTION_START=true
fi

if [ "$ACTION_START" = true ]; then
  if [ ! -f "${VAULT_BIN}" ]; then
    ACTION_INSTALL=true
  elif [ "${VAULT_SERVER_VERSION}" != $(get_current_vault_version) ]; then
    ACTION_INSTALL=true
    rm -rf ${VAULT_DIR}/*
  fi
fi


# cleanup
create_vault_dirs

if [ "${ACTION_STOP}" = true ]; then
  stop_vault
fi

if [ "${ACTION_INSTALL}" = true ]; then
  install_vault
fi

if [ "${ACTION_START}" = true ]; then

  echo "Creating virtualenv in ${VENV_DIR} ..." | tee -a "$LOG_FILE"
  "${PYTHON_EXECUTABLE}" -m venv "${VENV_DIR}" | tee -a "$LOG_FILE"
  if [ "${PIPESTATUS[0]}" -ne "0" ]; then
    echo "Setting up virtualenv has failed! Exiting..." | tee -a "$LOG_FILE"
    exit 1
  fi
  source "${VENV_DIR}"/bin/activate
  "${VENV_DIR}/bin/python" -m pip install -r "${VAULT_RESOURCES_DIR}/requirements.txt" | tee -a "$LOG_FILE"
  deactivate

  if [ ! -f "${VAULT_SERVER_CONF_FILE}" ]; then
    create_vault_server_conf
  fi

  start_vault
  if [ "$RUN_IN_BACKGROUND" = true ]; then
    exit 0
  else
    echo "LOG_FILE: ${LOG_FILE}"
    tail -f ${LOG_FILE}
  fi
fi
