#!/bin/sh

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT

set -e
set -u

readonly AGENTDIR="/opt/aws/amazon-cloudwatch-agent"
readonly CMDDIR="${AGENTDIR}/bin"
readonly CONFDIR="${AGENTDIR}/etc"
readonly CWA_RESTART_FILE="${CONFDIR}/restart"
readonly VERSION_FILE="${CMDDIR}/CWAGENT_VERSION"
readonly WORKLOAD_DISCOVERY="${CMDDIR}/workload-discovery"

# The systemd and upstart scripts assume exactly this .toml file name
readonly TOML="${CONFDIR}/amazon-cloudwatch-agent.toml"
readonly OTEL_YAML="${CONFDIR}/amazon-cloudwatch-agent.yaml"
readonly JSON="${CONFDIR}/amazon-cloudwatch-agent.json"
readonly JSON_DIR="${CONFDIR}/amazon-cloudwatch-agent.d"
readonly CV_LOG_FILE="${AGENTDIR}/logs/configuration-validation.log"
readonly COMMON_CONIG="${CONFDIR}/common-config.toml"
readonly ENV_CONFIG="${CONFDIR}/env-config.json"

readonly CWA_NAME='amazon-cloudwatch-agent'
readonly ALL_CONFIG='all'

SYSTEMD='false'

UsageString="


        usage:  amazon-cloudwatch-agent-ctl -a
                stop|start|status|status-with-workloads|fetch-config|append-config|remove-config|set-log-level
                [-m ec2|onPremise|onPrem|auto]
                [-c default|all|ssm:<parameter-store-name>|file:<file-path>]
                [-s]
                [-d]
                [-l INFO|DEBUG|WARN|ERROR|OFF]

        e.g.
        1. apply a SSM parameter store config on EC2 instance and restart the agent afterwards:
            amazon-cloudwatch-agent-ctl -a fetch-config -m ec2 -c ssm:AmazonCloudWatch-Config.json -s
        2. append a local json config file on onPremise host and restart the agent afterwards:
            amazon-cloudwatch-agent-ctl -a append-config -m onPremise -c file:/tmp/config.json -s
        3. query agent status:
            amazon-cloudwatch-agent-ctl -a status
        4. apply a SSM parameter store config with dual-stack endpoints:
            amazon-cloudwatch-agent-ctl -d -a fetch-config -m ec2 -c ssm:AmazonCloudWatch-Config.json -s

        -a: action
            stop:                                   stop the agent process.
            start:                                  start the agent process.
            status:                                 get the status of the agent process.
            status-with-workloads:                  get the status of the agent process and runs workload discovery.
            fetch-config:                           apply config for agent, followed by -c. Target config can be based on location (ssm parameter store name, file name), or 'default'.
            append-config:                          append json config with the existing json configs if any, followed by -c. Target config can be based on the location (ssm parameter store name, file name), or 'default'.
            remove-config:                          remove config for agent, followed by -c. Target config can be based on the location (ssm parameter store name, file name), or 'all'.
            set-log-level:                          sets the log level, followed by -l to provide the level in all caps.

        -m: mode
            ec2:                                    indicate this is on ec2 host.
            onPremise, onPrem:                      indicate this is on onPremise host.
            auto:                                   use ec2 metadata to determine the environment, may not be accurate if ec2 metadata is not available for some reason on EC2.

        -c: amazon-cloudwatch-agent configuration
            default:                                default configuration for quick trial.
            ssm:<parameter-store-name>:             ssm parameter store name.
            file:<file-path>:                       file path on the host.
            all:                                    all existing configs. Only apply to remove-config action.

        -s: optionally restart after configuring the agent configuration
            this parameter is used for 'fetch-config', 'append-config', 'remove-config' action only.

        -d: enable dual-stack endpoints for AWS API calls during config download
            this parameter is used for 'fetch-config', 'append-config' actions only.

        -l: log level to set the agent to INFO, DEBUG, WARN, ERROR, or OFF
            this parameter is used for 'set-log-level' only.

"

start_all() {
     mode="${1:-}"

     echo ""
     echo "****** processing amazon-cloudwatch-agent ******"
     agent_start "${CWA_NAME}" "${mode}"
}

agent_start() {
     agent_name="${1:-}"
     mode="${2:-}"

     if [ "$(runstatus ${agent_name})" = 'running' ]; then
          echo "${agent_name} has already been started"
          return 0
     fi

     if [ "${agent_name}" = "${CWA_NAME}" ] && [ ! -f "${TOML}" ]; then
          echo "${CWA_NAME} is not configured. Applying amazon-cloudwatch-agent default configuration."
          cwa_config 'default' 'false' "${mode}" 'default' 'false'
     fi

     if [ "${SYSTEMD}" = 'true' ]; then
          systemctl daemon-reload || return
          systemctl enable "${agent_name}.service" || return
          systemctl restart "${agent_name}.service" || return
     else
          start "${agent_name}" || return
          sleep 1
     fi
}

stop_all() {

     echo ""
     echo "****** processing amazon-cloudwatch-agent ******"
     agent_stop_and_disable "${CWA_NAME}"
}

agent_stop_and_disable() {
     agent_name="${1:-}"

     agent_stop "${agent_name}"
}

agent_stop() {
     agent_name="${1:-}"

     if [ "$(runstatus ${agent_name})" = 'stopped' ]; then
          echo "${agent_name} has already been stopped"
          return 0
     fi

     if [ "${SYSTEMD}" = 'true' ]; then
          systemctl stop "${agent_name}.service" || return
     else
          stop "${agent_name}" || return
     fi
}

# support for restart during upgrade via SSM packages
prep_restart_all() {
     agent_prep_restart "${CWA_NAME}" "${CWA_RESTART_FILE}"
}

agent_prep_restart() {
     agent_name="${1:-}"
     restart_file="${2:-}"
     if [ "$(runstatus ${agent_name})" = 'running' ]; then
          touch "${restart_file}"
     fi
}

# support for restart during upgrade via SSM packages
cond_restart_all() {
     agent_cond_restart "${CWA_NAME}" "${CWA_RESTART_FILE}"
}

agent_cond_restart() {
     agent_name="${1:-}"
     restart_file="${2:-}"
     if [ -f "${restart_file}" ]; then
          agent_start "${agent_name}" "${mode}"
          rm -f "${restart_file}"
     fi
}

preun_all() {
     agent_preun "${CWA_NAME}"
}

agent_preun() {
     agent_name="${1:-}"
     agent_stop "${agent_name}"
     if [ "${SYSTEMD}" = 'true' ]; then
          systemctl disable "${agent_name}.service" || return
          systemctl daemon-reload || return
          systemctl reset-failed || return
     fi
}

cwa_status() {
     include_workloads="${1:-false}"

     cwa_config_status='configured'
     if [ ! -f "${TOML}" ]; then
          cwa_config_status='not configured'
     fi

     version="$(cat ${VERSION_FILE})"

     echo "{"
     echo "  \"status\": \"$(runstatus ${CWA_NAME})\","
     echo "  \"starttime\": \"$(get_starttime_fmt ${CWA_NAME})\","
     echo "  \"configstatus\": \"${cwa_config_status}\","
     if [ "${include_workloads}" = "true" ]; then
          workloads_json="[]"
          if [ -x "${WORKLOAD_DISCOVERY}" ]; then
               workloads_json="$(${WORKLOAD_DISCOVERY} 2>/dev/null || echo '[]')"
               if [ -z "${workloads_json}" ]; then
                    workloads_json="[]"
               fi
          fi
          echo "  \"version\": \"${version}\","
          echo "  \"workloads\": ${workloads_json}"
     else
          echo "  \"version\": \"${version}\""
     fi

     echo "}"
}

get_starttime_fmt() {
     agent_name="${1:-}"

     agentPid=''
     if [ "${SYSTEMD}" = 'true' ]; then
          agentPid="$(systemctl show -p MainPID "${agent_name}.service" | sed s/MainPID=//)"
     else
          agentPid="$(initctl status "${agent_name}" | sed -n s/^.*process\ //p)"
     fi

     starttime_fmt=''
     if [ "${agentPid}" ] && [ "${agentPid}" -ne "0" ]; then
          starttime="$(TZ=UTC ps -o lstart= "${agentPid}")"
          starttime_fmt="$(TZ=UTC date -Isec -d "${starttime}")"
     fi

     echo "${starttime_fmt}"
}

runstatus() {
     agent_name="${1:-}"

     running=false
     if [ "${SYSTEMD}" = 'true' ]; then
          set +e
          if systemctl is-active "${agent_name}.service" 1>/dev/null; then
               running='true'
          fi
          set -e
     else
          if [ "$(initctl status "${agent_name}" | grep -c running)" = 1 ]; then
               running='true'
          fi
     fi

     if [ "${running}" = 'true' ]; then
          echo "running"
     else
          echo "stopped"
     fi
}

config_all() {
     cwa_config_location="${1:-}"
     restart="${2:-}"
     mode="${3:-}"
     multi_config="${4:-}"
     use_dualstack="${5:-}"

     if [ -z "${cwa_config_location}" ]; then
          cwa_config_location='default'
     fi

     mkdir -p "${CONFDIR}"

     if [ -n "${cwa_config_location}" ]; then
          echo "****** processing amazon-cloudwatch-agent ******"
          cwa_config "${cwa_config_location}" "${restart}" "${mode}" "${multi_config}" "${use_dualstack}"
     fi
}

cwa_config() {
     cwa_config_location="${1:-}"
     restart="${2:-}"
     param_mode="${3:-}"
     multi_config="${4:-}"
     use_dualstack="${5:-}"

     if [ "${cwa_config_location}" = "${ALL_CONFIG}" ] && [ "${multi_config}" != 'remove' ]; then
          echo "ignore cwa configuration \"${ALL_CONFIG}\" as it is only supported by action \"remove-config\""
          return
     fi

     if [ "${cwa_config_location}" = "${ALL_CONFIG}" ]; then
          rm -rf "${JSON_DIR}"/*
     else
          if [ "${use_dualstack}" = 'true' ]; then
               runDownloaderCommand=$("${CMDDIR}/config-downloader" --output-dir "${JSON_DIR}" --download-source "${cwa_config_location}" --mode ${param_mode} --config "${COMMON_CONIG}" --multi-config ${multi_config} --dualstack)
          else
               runDownloaderCommand=$("${CMDDIR}/config-downloader" --output-dir "${JSON_DIR}" --download-source "${cwa_config_location}" --mode ${param_mode} --config "${COMMON_CONIG}" --multi-config ${multi_config})
          fi
          echo "${runDownloaderCommand}" || return
     fi

     if [ ! "$(ls ${JSON_DIR})" ]; then
          echo "all amazon-cloudwatch-agent configurations have been removed"
          rm -f "${TOML}"
          rm -f "${OTEL_YAML}"
     else
          echo "Start configuration validation..."
          runTranslatorCommand=$("${CMDDIR}/config-translator" --input "${JSON}" --input-dir "${JSON_DIR}" --output "${TOML}" --mode ${param_mode} --config "${COMMON_CONIG}" --multi-config ${multi_config})
          echo "${runTranslatorCommand}" || return

          runAgentSchemaTestCommand="${CMDDIR}/amazon-cloudwatch-agent -schematest -config ${TOML}"
          echo "${runAgentSchemaTestCommand}"
          # We will redirect the verbose error message out
          if ! ${runAgentSchemaTestCommand} >${CV_LOG_FILE} 2>&1; then
               echo "Configuration validation second phase failed"
               echo "======== Error Log ========"
               cat ${CV_LOG_FILE}
               exit 1
          fi
          echo "Configuration validation second phase succeeded"
          echo "Configuration validation succeeded"

          chmod ug+rw "${TOML}"
          if [ -f "${OTEL_YAML}" ]; then
               chmod ug+rw "${OTEL_YAML}"
          fi

          # for translator:
          #       default:    only process .tmp files
          #       append:     process both existing files and .tmp files
          #       remove:     only process existing files
          # At this point, all json configs have been validated
          # multi_config:
          #       default:    delete non .tmp file, rename .tmp file
          #       append:     rename .tmp file
          #       remove:     no-op
          if [ "${multi_config}" = 'default' ]; then
               rm -f "${JSON}"
               for file in "${JSON_DIR}"/*; do
                    base="${JSON_DIR}/$(basename "${file}" .tmp)"
                    if [ "${file}" = "${base}" ]; then
                         rm -f "${file}"
                    else
                         mv -f "${file}" "${base}"
                    fi
               done
          elif [ "${multi_config}" = 'append' ]; then
               for file in "${JSON_DIR}"/*.tmp; do
                    mv -f "${file}" "${JSON_DIR}/$(basename "${file}" .tmp)"
               done
          fi
     fi

     if [ "${restart}" = 'true' ]; then
          agent_stop_and_disable "${CWA_NAME}"
          agent_start "${CWA_NAME}" "${param_mode}"
     fi
}

set_log_level_all() {
     log_level="${1:-}"
     case "${log_level}" in
     INFO) ;;

     DEBUG) ;;

     ERROR) ;;

     WARN) ;;

     OFF) ;;

     *)
          echo "Invalid log level: ${log_level} ${UsageString}" >&2
          exit 1
          ;;
     esac

     runEnvConfigCommand=$("${CMDDIR}/amazon-cloudwatch-agent" -setenv CWAGENT_LOG_LEVEL=${log_level} -envconfig "${ENV_CONFIG}")
     echo "${runEnvConfigCommand}" || return
     echo "Set CWAGENT_LOG_LEVEL to ${log_level}"
}

main() {
     action=''
     cwa_config_location=''
     restart='false'
     mode='ec2'
     use_dualstack='false'

     # detect which init system is in use
     if [ "$(/sbin/init --version 2>/dev/null | grep -c upstart)" = 1 ]; then
          SYSTEMD='false'
     elif [ "$(systemctl | grep -c -E '\-\.mount\s')" = 1 ]; then
          SYSTEMD='true'
     elif [ -f /etc/init.d/cron ] && [ ! -h /etc/init.d/cron ]; then
          echo "sysv-init is not supported" >&2
          exit 1
     else
          echo "unknown init system" >&2
          exit 1
     fi

     OPTIND=1
     while getopts ":hsa:c:m:l:d" opt; do
          case "${opt}" in
          h)
               echo "${UsageString}"
               exit 0
               ;;
          s) restart='true' ;;
          a) action="${OPTARG}" ;;
          c) cwa_config_location="${OPTARG}" ;;
          m) mode="${OPTARG}" ;;
          l) log_level="${OPTARG}" ;;
          d) use_dualstack='true' ;;
          \?)
               echo "Invalid option: -${OPTARG} ${UsageString}" >&2
               ;;
          :)
               echo "Option -${OPTARG} requires an argument ${UsageString}" >&2
               exit 1
               ;;
          esac
     done
     shift "$((OPTIND - 1))"

     case "${mode}" in
     ec2) ;;
     onPremise) ;;
     onPrem) ;;
     auto) ;;
     *)
          echo "Invalid mode: ${mode} ${UsageString}" >&2
          exit 1
          ;;
     esac

     case "${action}" in
     stop) stop_all ;;
     start) start_all "${mode}" ;;
     fetch-config) config_all "${cwa_config_location}" "${restart}" "${mode}" 'default' "${use_dualstack}" ;;
     append-config) config_all "${cwa_config_location}" "${restart}" "${mode}" 'append' "${use_dualstack}" ;;
     remove-config) config_all "${cwa_config_location}" "${restart}" "${mode}" 'remove' ;;
     status) cwa_status "false" ;;
     status-with-workloads) cwa_status "true" ;;
          # helpers for ssm package scripts to workaround fact that it can't determine if invocation is due to
          # upgrade or install
     prep-restart) prep_restart_all ;;
     cond-restart) cond_restart_all ;;
          # helper for rpm+deb uninstallation hooks, not expected to be called manually
     preun) preun_all ;;
     set-log-level) set_log_level_all "${log_level}" ;;
     *)
          echo "Invalid action: ${action} ${UsageString}" >&2
          exit 1
          ;;
     esac
     exit 0
}

main "$@"