#!/bin/bash
#
## Copyright 2024 Google LLC
##
## 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.
#
#
# Installs the BeyondCorp Remote Agent

set -euo pipefail

readonly BEYONDCORP_REPOSITORIES='gcr.io/appconnector-external-release'
readonly REMOTE_AGENT_REPOSITORY="${BEYONDCORP_REPOSITORIES}/appconnector_remote_cp"
readonly REMOTE_AGENT_TAG='appconnector_dp_rollout_20240302_rc00'
readonly REMOTE_AGENT_INSTALLER_IMAGE="${REMOTE_AGENT_REPOSITORY}:${REMOTE_AGENT_TAG}"
readonly REMOTE_AGENT_RUNNER_TARGET_IMAGE="${REMOTE_AGENT_REPOSITORY}:${REMOTE_AGENT_TAG}"
readonly REMOTE_AGENT_RUNNER_STABLE_IMAGE="${REMOTE_AGENT_REPOSITORY}:${REMOTE_AGENT_TAG}"
readonly REMOTE_AGENT_INSTALLER_ENTRYPOINT='/applink_control_runtime/bin/install'
readonly REMOTE_AGENT_CONTAINER='bce-control-runtime'
readonly BEYONDCORP_CONTAINERS=("${REMOTE_AGENT_CONTAINER}" 'bce-connector' 'bce-logagent')
readonly BEYONDCORP_SERVICE='beyondcorp'
readonly BEYONDCORP_DIR='/var/beyondcorp'

readonly COLOR_GREEN='\033[32m'
readonly COLOR_RED='\033[31m'
readonly COLOR_RESET='\033[0m'

BEYONDCORP_USER='beyondcorp'
UNINSTALL='false'

# logger
INFO() { echo "INFO: ${1}"; }
WARN() { echo "WARN: ${1}"; }
ERROR() { echo "ERROR: ${1}"; }
FATAL() { echo "FATAL: ${1}"; exit 1; }

repeat() {
  for ((i = 0; i < "${2}"; i++)); do
    echo -n "${1}"
  done
  echo
}

title() {
  local len="${#1}"
  repeat "#" "$((len + 4))"
  echo "# ${1} #"
  repeat "#" "$((len + 4))"
}

# 1. flags
usage() {
  cat << USAGE
  Configure and start beyondcorp remote agent installer
  Usage: $0 [-u <beyondcorp_user>] [-r]
  -h                       Print this help message.
  -u <beyondcorp_user>     The user the beyondcorp remote agent will run as.
                           Default is "${BEYONDCORP_USER}" if not specified.
  -r <uninstall>           Uninstall beyondcorp remote agent.
                           No other flags are required.
USAGE
}

install_complete_message() {
  local message
  message="$(cat << EOM
###############################################################################
# BeyondCorp Remote Agent has been successfully installed and started.        #
# Please run the following command to finish enrolling the remote agent:      #
# $ bce-connctl init                                                          #
#                                                                             #
# Other BeyondCorp service commands:                                          #
# --------------------------------------------------------------------------- #
# sudo systemctl stop beyondcorp & sudo systemctl disable beyondcorp | stop   #
# sudo systemctl start beyondcorp & sudo systemctl enable beyondcorp | start  #
# sudo systemctl status beyondcorp                                   | status #
# --------------------------------------------------------------------------- #
#                                                                             #
# BeyondCorp utility scripts for verifying Remote Agent environment setup:    #
# --------------------------------------------------------------------------- #
# sudo /var/beyondcorp/scripts/run-diagnostics         | General diagnostics  #
# sudo /var/beyondcorp/scripts/run-pre-install-checks  | Pre-install checks   #
# sudo /var/beyondcorp/scripts/run-post-install-checks | Post-install checks  #
# --------------------------------------------------------------------------- #
###############################################################################
EOM
)"
  printf "%b%s%b\n" "${COLOR_GREEN}" "${message}" "${COLOR_RESET}"
}

forbid_empty_flag() {
  if [[ -z "${1}" ]]; then
    ERROR "Empty flag '${2}' is forbidden."
    echo
    usage >&2
    exit 1
  fi
}

parse_flags() {
  title 'Parsing flags...'
  while getopts 'hu:c:s:r' opt; do
    case "${opt}" in
      h)
        usage
        exit 0
        ;;
      u) BEYONDCORP_USER="${OPTARG}" ;;
      r) UNINSTALL='true' ;;
      ?)
        usage >&2
        exit 1
        ;;
    esac
  done
  if "${UNINSTALL}"; then
    return 0
  fi
  forbid_empty_flag "${BEYONDCORP_USER}" "<beyondcorp_user>"
}

# 2. checks
check_dependency() {
  INFO "Checking dependency '${1}'"
  if ! ${2}; then
    FATAL "Missing mandatory command: '${1}'. ${3}."
  fi
  echo
}

precondition_checks() {
  title 'Checking dependencies...'
  check_dependency "docker" "docker --version" "Please install Docker Engine on this machine (https://docs.docker.com/engine/install/)."

  # Enabling root privileges for the $USER
  if [[ "${USER}" != 'root' ]] && ! (id -nG "${USER}" | grep -qw "docker") ; then
    INFO "Appending user '${USER}' to group 'docker'."
    if ! sudo usermod -aG "docker" "${USER}"; then
      FATAL "Failed to append user '${USER}' to group 'docker'."
    fi
    printf "%b%s%b\n" "${COLOR_RED}" "Updated group membership for '${USER}'. Please log out and log back in for the changes to be re-evaluated and run the installation script again." "${COLOR_RESET}"
    exit
  fi

  check_dependency "systemctl" "systemctl --version" "Please install systemctl on this machine (https://www.man7.org/linux/man-pages/man1/systemctl.1.html)."
}

installation_checks() {
  title 'Checking BeyondCorp Remote Agent environment...'
  sudo "${BEYONDCORP_DIR}/scripts/run-pre-install-checks"
}

# 3. user
add_user() {
  title 'Adding user...'
  if id "${BEYONDCORP_USER}" &> /dev/null; then
    INFO "User '${BEYONDCORP_USER}' already exists. Continue."
  else
    INFO "Adding user '${BEYONDCORP_USER}'."
    if ! sudo useradd "${BEYONDCORP_USER}"; then
      FATAL "Failed to add user '${BEYONDCORP_USER}'."
    fi
  fi

  if id -nG "${BEYONDCORP_USER}" | grep -qw "docker"; then
    INFO "User '${BEYONDCORP_USER}' belongs to group 'docker'. Continue."
  else
    INFO "Appending user '${BEYONDCORP_USER}' to group 'docker'."
    if ! sudo usermod -aG "docker" "${BEYONDCORP_USER}"; then
      FATAL "Failed to append user '${BEYONDCORP_USER}' to group 'docker'."
    fi
  fi

  INFO "Creating home directory for user '${BEYONDCORP_USER}'."
  if ! sudo mkhomedir_helper "${BEYONDCORP_USER}"; then
    FATAL "Failed to create home directory for user '${BEYONDCORP_USER}'."
  fi

  docker_dir="/home/${BEYONDCORP_USER}/.docker"
  if sudo test -d "${docker_dir}"; then
    INFO "Directory '${docker_dir}' already exists. Continue."
  else
    INFO "Creating '${docker_dir}' directory."
    if ! sudo mkdir "${docker_dir}"; then
      FATAL "Failed to create '${docker_dir}' directory for user '${BEYONDCORP_USER}'."
    fi
  fi

  INFO "Updating permissions for '${docker_dir}' directory."
  if ! sudo chown -R "${BEYONDCORP_USER}:${BEYONDCORP_USER}" "${docker_dir}"; then
    FATAL "Failed to update permissions for '${docker_dir}' directory."
  fi
}

# Before starting the uninstall, find the installed beyondcorp user. It will be
# use by delete_user().
installed_beyondcorp_user=''
get_user() {
  local user arr
  # load user from the running container.
  IFS=":" read -r -a arr <<< "$(docker inspect --format '{{.Config.User}}' "${REMOTE_AGENT_CONTAINER}" 2> /dev/null)"
  if [[ "${#arr[@]}" -ne 0 ]]; then
    if user="$(id -nu "${arr[0]}" 2> /dev/null)"; then
      installed_beyondcorp_user="${user}"
      return 0
    fi
  fi
  # load user from the systemd service.
  IFS="=" read -r -a arr <<< "$(systemctl show "${BEYONDCORP_SERVICE}" -p UID)"
  if [[ "${#arr[@]}" -eq 2 ]] && [[ "${arr[1]}" =~ ^[0-9]+$ ]]; then
    if user="$(id -nu "${arr[1]}" 2> /dev/null)"; then
      installed_beyondcorp_user="${user}"
      return 0
    fi
  fi
  # load user from the created directory.
  if user="$(stat -c '%U' "${BEYONDCORP_DIR}" 2> /dev/null)"; then
    installed_beyondcorp_user="${user}"
    return 0
  fi
}

delete_user() {
  # Don't delete custom user.
  if [[ "${installed_beyondcorp_user}" != 'beyondcorp' ]]; then
    return 0
  fi

  title 'Deleting user...'
  if [[ -z "${installed_beyondcorp_user}" ]]; then
    WARN "Failed to find the installed beyondcorp user. Continue."
    return 0
  fi
  if ! id "${installed_beyondcorp_user}" &> /dev/null; then
    WARN "User '${installed_beyondcorp_user}' doesn't exist. Continue."
    return 0
  fi
  INFO "Deleting user '${installed_beyondcorp_user}'."
  if ! sudo userdel -r "${installed_beyondcorp_user}"; then
    WARN "Failed to delete user '${installed_beyondcorp_user}'. Continue."
  fi
}

# 4. docker
docker_flags=(
  -v /:/mounted_host_root
  -v "${HOME}":/mounted_host_home
  --entrypoint="${REMOTE_AGENT_INSTALLER_ENTRYPOINT}"
)

clean_up_previous_installations() {
  WARN 'Any previous installations of the BeyondCorp Remote Agent will be removed.'
  read -p 'Do you want to continue (y/N)? ' USER_INPUT
  if [[ "${USER_INPUT}" == [yY] || "${USER_INPUT}" == [yY][eE][sS] ]]; then
    title 'Cleaning up previous BeyondCorp installations...'
    ./"${0}" -r
  else
    INFO 'Stopping the BeyondCorp Remote Agent installer.'
    exit 1
  fi
}

set_docker_proxy() {
  set +u
  title 'Checking proxy settings...'

  if [[ -z "${NO_PROXY}" && -z "${no_proxy}" && -z "${HTTP_PROXY}" && \
        -z "${http_proxy}" && -z "${HTTPS_PROXY}" && -z "${https_proxy}"  ]]; then
    INFO 'No proxy settings detected. Continue.'
    set -u
    return 0
  elif [[ (-n "${http_proxy}" || -n "${HTTP_PROXY}") && -z "${HTTPS_PROXY}" && -z "${https_proxy}" ]]; then
    FATAL 'HTTP proxy is set but HTTPS proxy is not. Please set both HTTP and HTTPS proxies. Stopping the BeyondCorp Remote Agent installer.'
  elif [[ (-n "${https_proxy}" || -n "${HTTPS_PROXY}") && -z "${HTTP_PROXY}" && -z "${http_proxy}" ]]; then
    FATAL 'HTTPS proxy is set but HTTP proxy is not. Please set both HTTP and HTTPS proxies. Stopping the BeyondCorp Remote Agent installer.'
  fi

  if [[ -n "${HTTPS_PROXY}" && ! "${HTTPS_PROXY}" =~ ^http:// ]]; then
    FATAL 'Only the http:// schema is supported for BeyondCorp proxies. Please use the http:// schema. Stopping the BeyondCorp Remote Agent installer.'
  elif [[ -n "${https_proxy}" && ! "${https_proxy}" =~ ^http:// ]]; then
    FATAL 'Only the http:// schema is supported for BeyondCorp proxies. Please use the http:// schema. Stopping the BeyondCorp Remote Agent installer.'
  fi

  docker_config_file="/home/${BEYONDCORP_USER}/.docker/config.json"

  # Let the user set docker proxy config if we are not using the default
  # beyondcorp user.
  if [[ "${BEYONDCORP_USER}" != 'beyondcorp' ]]; then
    WARN "Proxy settings detected and using '${BEYONDCORP_USER}' as beyondcorp_user."
    read -p "Please confirm that '${docker_config_file}' is updated with proxy settings. See https://docs.docker.com/network/proxy/#configure-the-docker-client for details. (y/N)" USER_INPUT
    if [[ "${USER_INPUT}" == [yY] || "${USER_INPUT}" == [yY][eE][sS] ]]; then
      INFO "User confirmed '${docker_config_file}' is updated with proxy settings."
      set -u
      return 0
    else
      INFO "Please update '${docker_config_file}' with proxy settings. Stopping the BeyondCorp Remote Agent installer."
      exit 1
    fi
  fi

  INFO "Proxy settings detected, creating '${docker_config_file}'"

  proxies=()

  if [[ -n "${HTTP_PROXY}" ]]; then
    proxies+=("   \"httpProxy\": \"${HTTP_PROXY}\"")
  elif [[ -n "${http_proxy}" ]]; then
    proxies+=("   \"httpProxy\": \"${http_proxy}\"")
  fi

  if [[ -n "${HTTPS_PROXY}" ]]; then
    proxies+=("   \"httpsProxy\": \"${HTTPS_PROXY}\"")
  elif [[ -n "${https_proxy}" ]]; then
    proxies+=("   \"httpsProxy\": \"${https_proxy}\"")
  fi

  if [[ -n "${NO_PROXY}" ]]; then
    proxies+=("   \"noProxy\": \"${NO_PROXY}\"")
  elif [[ -n "${no_proxy}" ]]; then
    proxies+=("   \"noProxy\": \"${no_proxy}\"")
  fi

  printf -v default_proxies '%s,\n' "${proxies[@]}"
  default_proxies="${default_proxies%,$'\n'}"

  sudo dd of="${docker_config_file}" << EOF
{
"proxies": {
  "default": {
${default_proxies}
  }
 }
}
EOF

  INFO "Updating permissions for '${docker_config_file}'."
  if ! sudo chown "${BEYONDCORP_USER}:${BEYONDCORP_USER}" "${docker_config_file}"; then
    FATAL "Failed to update permissions for '${docker_config_file}'"
  fi

  set -u
}

install() {
  title 'Installing BeyondCorp Remote Agent...'
  installer_flags=(
    --target_image="${REMOTE_AGENT_RUNNER_TARGET_IMAGE}"
    --stable_image="${REMOTE_AGENT_RUNNER_STABLE_IMAGE}"
    --shell="$(ps -p $PPID -o fname --no-headers)"
    --admin_uid="$(id -u)"
    --admin_gid="$(id -g)"
    --beyondcorp_uid="$(id -u "${BEYONDCORP_USER}")"
    --beyondcorp_gid="$(id -g "${BEYONDCORP_USER}")"
  )
  # Pull before run. Otherwise, docker reuse the local one with the same tag
  # which may not be the version we want.
  docker pull "${REMOTE_AGENT_INSTALLER_IMAGE}"
  if ! sudo -u "${BEYONDCORP_USER}" docker run --rm "${docker_flags[@]}" "${REMOTE_AGENT_INSTALLER_IMAGE}" "${installer_flags[@]}"; then
    FATAL "Failed to install BeyondCorp Remote Agent."
  fi
  # Remove the installer.
  if [[ "${REMOTE_AGENT_INSTALLER_IMAGE}" != "${REMOTE_AGENT_RUNNER_TARGET_IMAGE}" && "${REMOTE_AGENT_INSTALLER_IMAGE}" != "${REMOTE_AGENT_RUNNER_STABLE_IMAGE}" ]]; then
    if ! docker rmi -f "${REMOTE_AGENT_INSTALLER_IMAGE}"; then
      WARN "Failed to remove installer docker image '${image}'. Continue."
    fi
  fi
}

uninstall() {
  title 'Uninstalling BeyondCorp Remote Agent...'
  installer_flags=(
    --uninstall="true"
    # If user removed ~/.bce_alias before uninstall, we rely on `shell` to
    # locate the shell rc.
    --shell="$(ps -p $PPID -o fname --no-headers)"
  )
  # uninstall directories and files.
  if ! docker run --rm "${docker_flags[@]}" "${REMOTE_AGENT_INSTALLER_IMAGE}" "${installer_flags[@]}"; then
    WARN "Failed to remove installed directories and files. Continue."
  fi
  # remove containers
  for container in "${BEYONDCORP_CONTAINERS[@]}"; do
    if [[ "$(docker ps --filter "name=${container}" --format="{{.Names}}" | wc -l)" -eq 0 ]]; then
      continue
    fi
    INFO "Removing docker container '${container}'."
    if ! docker rm -f "${container}"; then
      WARN "Failed to remove docker container '${container}'. Continue."
    fi
  done
  # remove images
  mapfile -t images < <(docker images --format='{{if eq .Tag "<none>"}}{{.ID}}{{else}}{{.Repository}}:{{.Tag}}{{end}}' "${BEYONDCORP_REPOSITORIES}/*")
  for image in "${images[@]}"; do
    INFO "Removing docker image '${image}'."
    if ! docker rmi -f "${image}"; then
      WARN "Failed to remove docker image '${image}'. Continue."
    fi
  done
}

# 5. systemd service
start_service() {
  title 'Starting systemd service...'
  INFO "Enabling systemd service '${BEYONDCORP_SERVICE}'."
  if ! sudo systemctl enable "${BEYONDCORP_SERVICE}"; then
    FATAL "Failed to enable systemd service '${BEYONDCORP_SERVICE}'."
  fi
  INFO "Starting systemd service '${BEYONDCORP_SERVICE}'."
  if ! sudo systemctl start "${BEYONDCORP_SERVICE}"; then
    FATAL "Failed to start systemd service '${BEYONDCORP_SERVICE}'."
  fi
}

stop_service() {
  title 'Stopping systemd service...'
  INFO "Stopping systemd service '${BEYONDCORP_SERVICE}'."
  if ! sudo systemctl stop "${BEYONDCORP_SERVICE}"; then
    WARN "Failed to stop systemd service '${BEYONDCORP_SERVICE}'. Continue."
  fi
  INFO "Disabling systemd service '${BEYONDCORP_SERVICE}'."
  if ! sudo systemctl disable "${BEYONDCORP_SERVICE}"; then
    WARN "Failed to disable systemd service '${BEYONDCORP_SERVICE}'. Continue."
  fi
}

# main
parse_flags "$@"
if "${UNINSTALL}"; then
  get_user
  stop_service
  uninstall
  delete_user
  title 'BeyondCorp Remote Agent has been successfully stopped and uninstalled.'
  exit 0
fi
precondition_checks
clean_up_previous_installations
add_user
set_docker_proxy
install
installation_checks
start_service
install_complete_message