#!/usr/bin/env bash
set -euo pipefail

if [[ "${EUID}" -ne 0 ]]; then
  echo "[ERROR] please run as root"
  exit 1
fi

if [[ ! -f /etc/os-release ]]; then
  echo "[ERROR] /etc/os-release not found"
  exit 1
fi

# shellcheck disable=SC1091
source /etc/os-release
os_id="$(echo "${ID:-}" | tr '[:upper:]' '[:lower:]')"
os_ver="${VERSION_ID:-}"
is_ubuntu22=false
is_amzn2023=false

if [[ "${os_id}" == "ubuntu" && "${os_ver}" == 22.04* ]]; then
  is_ubuntu22=true
fi
if [[ "${os_id}" == "amzn" && "${os_ver}" == 2023* ]]; then
  is_amzn2023=true
fi

if [[ "${is_ubuntu22}" != "true" && "${is_amzn2023}" != "true" ]]; then
  echo "[ERROR] only Ubuntu 22.04 or Amazon Linux 2023 is supported. current: ID=${ID:-unknown}, VERSION_ID=${VERSION_ID:-unknown}"
  exit 1
fi

if ! command -v openssl >/dev/null 2>&1 || ! command -v curl >/dev/null 2>&1; then
  if [[ "${is_ubuntu22}" == "true" ]]; then
    apt-get update -y
    DEBIAN_FRONTEND=noninteractive apt-get install -y openssl curl ca-certificates
  else
    dnf makecache -y
    dnf install -y openssl curl ca-certificates
  fi
fi

CH_HTTPS_PORT="${CH_HTTPS_PORT:-8443}"
CH_LISTEN_HOST="${CH_LISTEN_HOST:-::}"
CH_CERT_CN="${CH_CERT_CN:-$(hostname -f 2>/dev/null || hostname)}"
CH_CERT_DNS="${CH_CERT_DNS:-}"
CH_CERT_IP="${CH_CERT_IP:-}"
CH_CERT_DAYS="${CH_CERT_DAYS:-825}"
CH_GENERATE_CA="${CH_GENERATE_CA:-false}"

SRC_CERT="${SRC_CERT:-}"
SRC_KEY="${SRC_KEY:-}"
SRC_CA="${SRC_CA:-}"

CH_DIR="/etc/clickhouse-server"
CH_CONFIG_D_DIR="${CH_DIR}/config.d"
PKI_DIR="${CH_DIR}/pki"
SERVER_CERT="${CH_DIR}/server.crt"
SERVER_KEY="${CH_DIR}/server.key"
CA_CERT="${CH_DIR}/ca.crt"
OVERRIDE_FILE="${CH_CONFIG_D_DIR}/waf-https.xml"

mkdir -p "${CH_CONFIG_D_DIR}" "${PKI_DIR}"

split_csv() {
  local raw="$1"
  if [[ -z "${raw}" ]]; then
    return 0
  fi
  IFS=',' read -r -a arr <<<"${raw}"
  for item in "${arr[@]}"; do
    item="$(echo "${item}" | xargs)"
    if [[ -n "${item}" ]]; then
      echo "${item}"
    fi
  done
}

build_san_line() {
  local san_entries=()
  while IFS= read -r dns_item; do
    san_entries+=("DNS:${dns_item}")
  done < <(split_csv "${CH_CERT_DNS}")
  while IFS= read -r ip_item; do
    san_entries+=("IP:${ip_item}")
  done < <(split_csv "${CH_CERT_IP}")

  if [[ ${#san_entries[@]} -eq 0 ]]; then
    san_entries+=("DNS:${CH_CERT_CN}")
  fi

  local san_line
  san_line="$(IFS=,; echo "${san_entries[*]}")"
  echo "${san_line}"
}

generate_self_signed_cert() {
  echo "[INFO] generating self-signed server certificate (crt+key only) ..."
  local server_key="${PKI_DIR}/server.key"
  local server_csr="${PKI_DIR}/server.csr"
  local server_crt="${PKI_DIR}/server.crt"
  local ext_file="${PKI_DIR}/server.ext"
  local san_line
  san_line="$(build_san_line)"

  openssl genrsa -out "${server_key}" 2048
  openssl req -new -key "${server_key}" -out "${server_csr}" -subj "/CN=${CH_CERT_CN}"

  cat >"${ext_file}" <<EOF
subjectAltName=${san_line}
keyUsage=digitalSignature,keyEncipherment
extendedKeyUsage=serverAuth
EOF

  openssl x509 -req -in "${server_csr}" -signkey "${server_key}" \
    -out "${server_crt}" -days "${CH_CERT_DAYS}" -sha256 -extfile "${ext_file}"

  cp -f "${server_crt}" "${SERVER_CERT}"
  cp -f "${server_key}" "${SERVER_KEY}"
  rm -f "${CA_CERT}"
}

generate_cert_with_ca() {
  echo "[INFO] generating local CA and server certificate ..."
  local ca_key="${PKI_DIR}/ca.key"
  local ca_crt="${PKI_DIR}/ca.crt"
  local server_key="${PKI_DIR}/server.key"
  local server_csr="${PKI_DIR}/server.csr"
  local server_crt="${PKI_DIR}/server.crt"
  local ext_file="${PKI_DIR}/server.ext"
  local san_line
  san_line="$(build_san_line)"

  openssl genrsa -out "${ca_key}" 4096
  openssl req -x509 -new -nodes -key "${ca_key}" -sha256 -days 3650 \
    -out "${ca_crt}" -subj "/CN=ClickHouse Local CA"

  openssl genrsa -out "${server_key}" 2048
  openssl req -new -key "${server_key}" -out "${server_csr}" -subj "/CN=${CH_CERT_CN}"

  cat >"${ext_file}" <<EOF
subjectAltName=${san_line}
keyUsage=digitalSignature,keyEncipherment
extendedKeyUsage=serverAuth
EOF

  openssl x509 -req -in "${server_csr}" -CA "${ca_crt}" -CAkey "${ca_key}" -CAcreateserial \
    -out "${server_crt}" -days "${CH_CERT_DAYS}" -sha256 -extfile "${ext_file}"

  cp -f "${server_crt}" "${SERVER_CERT}"
  cp -f "${server_key}" "${SERVER_KEY}"
  cp -f "${ca_crt}" "${CA_CERT}"
}

if [[ -n "${SRC_CERT}" || -n "${SRC_KEY}" ]]; then
  if [[ -z "${SRC_CERT}" || -z "${SRC_KEY}" ]]; then
    echo "[ERROR] SRC_CERT and SRC_KEY must be provided together"
    exit 1
  fi
  echo "[INFO] using provided certificate files ..."
  cp -f "${SRC_CERT}" "${SERVER_CERT}"
  cp -f "${SRC_KEY}" "${SERVER_KEY}"
  if [[ -n "${SRC_CA}" ]]; then
    cp -f "${SRC_CA}" "${CA_CERT}"
  else
    rm -f "${CA_CERT}"
  fi
else
  case "$(echo "${CH_GENERATE_CA}" | tr '[:upper:]' '[:lower:]')" in
    1|true|yes|on)
      generate_cert_with_ca
      ;;
    *)
      generate_self_signed_cert
      ;;
  esac
fi

chown clickhouse:clickhouse "${SERVER_CERT}" "${SERVER_KEY}" || true
chmod 0644 "${SERVER_CERT}"
chmod 0640 "${SERVER_KEY}"
if [[ -f "${CA_CERT}" ]]; then
  chown clickhouse:clickhouse "${CA_CERT}" || true
  chmod 0644 "${CA_CERT}"
fi

echo "[INFO] writing ClickHouse HTTPS override config ..."
cat >"${OVERRIDE_FILE}" <<EOF
<clickhouse>
    <https_port>${CH_HTTPS_PORT}</https_port>
    <listen_host>${CH_LISTEN_HOST}</listen_host>
    <openSSL>
        <server>
            <certificateFile>${SERVER_CERT}</certificateFile>
            <privateKeyFile>${SERVER_KEY}</privateKeyFile>
            <verificationMode>none</verificationMode>
            <loadDefaultCAFile>true</loadDefaultCAFile>
            <cacheSessions>true</cacheSessions>
            <disableProtocols>sslv2,sslv3</disableProtocols>
            <preferServerCiphers>true</preferServerCiphers>
            <invalidCertificateHandler>
                <name>RejectCertificateHandler</name>
            </invalidCertificateHandler>
        </server>
    </openSSL>
</clickhouse>
EOF

echo "[INFO] restarting clickhouse-server ..."
systemctl restart clickhouse-server
sleep 2

echo "[INFO] service status ..."
systemctl --no-pager -l status clickhouse-server | sed -n '1,15p'

echo "[INFO] verifying HTTPS endpoint ..."
curl -sk "https://127.0.0.1:${CH_HTTPS_PORT}/?query=SELECT%201" || true
echo

echo "[OK] ClickHouse HTTPS setup finished"
echo "     HTTPS port : ${CH_HTTPS_PORT}"
echo "     cert file  : ${SERVER_CERT}"
echo "     key file   : ${SERVER_KEY}"
if [[ -f "${CA_CERT}" ]]; then
  echo "     CA file    : ${CA_CERT}"
  echo "     import this CA file into API/Fluent Bit hosts if tls.verify=On"
fi