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

# Ensure script is run as root
if [ "$EUID" -ne 0 ]; then
  echo -e "\033[1;31mERROR: Please run this script as root (using sudo)\033[0m" >&2
  exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PYTHON_SCRIPT="$SCRIPT_DIR/update-certs.py"

if command -v python3 >/dev/null 2>&1 && [ -f "$PYTHON_SCRIPT" ]; then
  exec python3 "$PYTHON_SCRIPT" "$@"
fi

# ==========================================
# COLOR & FORMATTING DEFINITIONS
# ==========================================
RED='\033[0;31m'
BOLD_RED='\033[1;31m'
GREEN='\033[0;32m'
BOLD_GREEN='\033[1;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
BOLD_BLUE='\033[1;34m'
CYAN='\033[0;36m'
BOLD_CYAN='\033[1;36m'
DIM='\033[2m'
RESET='\033[0m'

log_section() {
  echo -e "\n${BOLD_CYAN}====================================================${RESET}"
  echo -e " ${BOLD_BLUE}$1${RESET}"
  echo -e "${BOLD_CYAN}====================================================${RESET}"
}

log_step() {
  echo -e " ${BOLD_CYAN}[+]${RESET} $1"
}

log_sub() {
  echo -e "    ${CYAN}↳${RESET} $1"
}

log_success() {
  echo -e "    ${BOLD_GREEN}✔${RESET} $1"
}

log_warn() {
  echo -e "    ${YELLOW}⚠${RESET} $1"
}

log_error() {
  echo -e "    ${BOLD_RED}✘${RESET} $1"
}

spin_until() {
  local msg="$1"
  local check_cmd="$2"
  local max_attempts="${3:-30}"
  local delay="${4:-2}"
  
  local spin_chars=('⠋' '⠙' '⠹' '⠸' '⠼' '⠴' '⠦' '⠧' '⠇' '⠏')
  local spin_i=0

  for ((i=1; i<=max_attempts; i++)); do
    if eval "$check_cmd" >/dev/null 2>&1; then
      echo -e "\r    ${BOLD_GREEN}✔${RESET} ${msg} ${BOLD_GREEN}(Ready!)${RESET}                                  "
      return 0
    fi

    for ((s=0; s<delay*10; s++)); do
      local char="${spin_chars[$spin_i]}"
      spin_i=$(( (spin_i + 1) % ${#spin_chars[@]} ))
      echo -ne "\r    ${YELLOW}${char}${RESET} ${msg} ${DIM}(attempt ${i}/${max_attempts})${RESET}   "
      sleep 0.1
    done
  done

  echo -e "\r    ${BOLD_RED}✘${RESET} ${msg} ${BOLD_RED}(Timed out after ${max_attempts}s)${RESET}                                  "
  return 1
}

# ==========================================
# CONFIGURATION SECTION
# ==========================================
CERT_DIR="${CERT_DIR:-/etc/elasticsearch/certs}"

ES_HTTP_P12_TARGET="${ES_HTTP_P12_TARGET:-/etc/elasticsearch/certs/http.p12}"
ES_TRANSPORT_P12_TARGET="${ES_TRANSPORT_P12_TARGET:-/etc/elasticsearch/certs/transport.p12}"
ES_HTTP_CA_TARGET="${ES_HTTP_CA_TARGET:-/etc/elasticsearch/certs/http_ca.crt}"

KIBANA_CRT_TARGET="${KIBANA_CRT_TARGET:-/etc/kibana/certs/kibana.crt}"
KIBANA_KEY_TARGET="${KIBANA_KEY_TARGET:-/etc/kibana/certs/kibana.key}"
KIBANA_CA_TARGET="${KIBANA_CA_TARGET:-/var/lib/kibana/ca_1721297493415.crt}"

KIBANA_CONFIG="${KIBANA_CONFIG:-/etc/kibana/kibana.yml}"

FLEET_SERVER_ES_URL="${FLEET_SERVER_ES_URL:-https://10.1.10.3:9200}"
FLEET_SERVER_URL="${FLEET_SERVER_URL:-https://10.1.10.3:8220}"
FLEET_SERVER_POLICY="${FLEET_SERVER_POLICY:-fleet-server-policy}"
FLEET_SERVER_PORT="${FLEET_SERVER_PORT:-8220}"

SAN_DNS="${SAN_DNS:-localhost cmsrv2024 siem2024}"
SAN_IPS="${SAN_IPS:-127.0.0.1 10.1.10.3 192.168.196.95}"

write_instances_yml() {
  local target_path="$1"
  cat << EOF > "$target_path"
instances:
  - name: "elasticsearch"
    dns:
EOF
  for d in $SAN_DNS; do
    echo "      - \"$d\"" >> "$target_path"
  done
  echo "    ip:" >> "$target_path"
  for ip in $SAN_IPS; do
    echo "      - \"$ip\"" >> "$target_path"
  done
  cat << EOF >> "$target_path"
  - name: "kibana"
    dns:
EOF
  for d in $SAN_DNS; do
    echo "      - \"$d\"" >> "$target_path"
  done
  echo "    ip:" >> "$target_path"
  for ip in $SAN_IPS; do
    echo "      - \"$ip\"" >> "$target_path"
  done
}

# ==========================================
# EXECUTION START
# ==========================================
CA_PATH="$CERT_DIR/elastic-stack-ca.p12"
ES_CERTUTIL="/usr/share/elasticsearch/bin/elasticsearch-certutil"

log_section "Elastic Stack Certificate Renewal & Service Sync"

log_step "Ensuring certificate directory exists..."
mkdir -p "$CERT_DIR"
log_sub "Target directory: $CERT_DIR"

log_step "Checking Certificate Authority (CA) status..."
CA_VALID=false
if [ -f "$CA_PATH" ]; then
  if openssl pkcs12 -in "$CA_PATH" -passin pass: -nokeys &>/dev/null; then
    log_success "Existing passwordless CA found."
    CA_VALID=true
  else
    log_warn "CA exists but is password-protected."
  fi
else
  log_sub "No existing CA found."
fi

if [ "$CA_VALID" = false ]; then
  log_step "Regenerating fresh passwordless CA..."
  rm -f "$CA_PATH"
  "$ES_CERTUTIL" ca --out "$CA_PATH" --pass "" >/dev/null 2>&1
  log_success "New CA generated successfully."
fi

log_step "Extracting HTTP CA certificate..."
openssl pkcs12 -in "$CA_PATH" -passin pass: -nokeys -out "$CERT_DIR/http_ca.crt" >/dev/null 2>&1
log_success "Extracted HTTP CA certificate to $CERT_DIR/http_ca.crt"

log_step "Generating instances configuration..."
write_instances_yml "$CERT_DIR/instances.yml"
log_sub "instances.yml written."

log_step "Generating node certificates..."
rm -f "$CERT_DIR/certs.zip"
"$ES_CERTUTIL" cert --ca "$CA_PATH" --ca-pass "" --in "$CERT_DIR/instances.yml" --out "$CERT_DIR/certs.zip" --pass "" >/dev/null 2>&1
log_success "Node certificates generated into certs.zip"

log_step "Extracting certificate archive..."
python3 -c "import zipfile; zipfile.ZipFile('$CERT_DIR/certs.zip').extractall('$CERT_DIR')"
log_sub "Certificates extracted."

log_step "Processing and organizing certificate files..."
mv "$CERT_DIR/elasticsearch/elasticsearch.p12" "$CERT_DIR/http.p12"
cp "$CERT_DIR/http.p12" "$CERT_DIR/transport.p12"

openssl pkcs12 -in "$CERT_DIR/kibana/kibana.p12" -clcerts -nokeys -out "$CERT_DIR/kibana.crt" -passin pass: >/dev/null 2>&1
openssl pkcs12 -in "$CERT_DIR/kibana/kibana.p12" -nocerts -nodes -out "$CERT_DIR/kibana.key.tmp" -passin pass: >/dev/null 2>&1
openssl pkey -in "$CERT_DIR/kibana.key.tmp" -out "$CERT_DIR/kibana.key" >/dev/null 2>&1

rm -rf "$CERT_DIR/elasticsearch" "$CERT_DIR/kibana" "$CERT_DIR/certs.zip" "$CERT_DIR/instances.yml" "$CERT_DIR/kibana.key.tmp"
log_success "Certificates converted to required formats."

log_step "Clearing legacy secure keystore passwords..."
/usr/share/elasticsearch/bin/elasticsearch-keystore remove xpack.security.http.ssl.keystore.secure_password >/dev/null 2>&1 || true
/usr/share/elasticsearch/bin/elasticsearch-keystore remove xpack.security.transport.ssl.keystore.secure_password >/dev/null 2>&1 || true
/usr/share/elasticsearch/bin/elasticsearch-keystore remove xpack.security.transport.ssl.truststore.secure_password >/dev/null 2>&1 || true
log_sub "Keystore parameters updated."

log_step "Applying strict file ownership and permissions..."
chown root:elasticsearch "$CERT_DIR"
chmod 755 "$CERT_DIR"

chown root:elasticsearch "$CERT_DIR/elastic-stack-ca.p12" "$CERT_DIR/http.p12" "$CERT_DIR/transport.p12"
chmod 640 "$CERT_DIR/elastic-stack-ca.p12" "$CERT_DIR/http.p12" "$CERT_DIR/transport.p12"

chown root:elasticsearch "$CERT_DIR/http_ca.crt"
chmod 644 "$CERT_DIR/http_ca.crt"

chown root:kibana "$CERT_DIR/kibana.crt" "$CERT_DIR/kibana.key"
chmod 644 "$CERT_DIR/kibana.crt"
chmod 640 "$CERT_DIR/kibana.key"

if [ -f "/etc/elasticsearch/service_tokens" ]; then
  chown root:elasticsearch "/etc/elasticsearch/service_tokens" || true
  chmod 660 "/etc/elasticsearch/service_tokens" || true
fi

usermod -aG elasticsearch kibana || true
log_success "Permissions set."

log_step "Updating symlinks to target service locations..."
symlink_cert() {
  local src="$1"
  local dest="$2"
  if [ "$src" != "$dest" ]; then
    mkdir -p "$(dirname "$dest")"
    rm -f "$dest"
    ln -sf "$src" "$dest"
    log_sub "$src ➜ $dest"
  fi
}

symlink_cert "$CERT_DIR/http.p12" "$ES_HTTP_P12_TARGET"
symlink_cert "$CERT_DIR/transport.p12" "$ES_TRANSPORT_P12_TARGET"
symlink_cert "$CERT_DIR/http_ca.crt" "$ES_HTTP_CA_TARGET"
symlink_cert "$CERT_DIR/kibana.crt" "$KIBANA_CRT_TARGET"
symlink_cert "$CERT_DIR/kibana.key" "$KIBANA_KEY_TARGET"
symlink_cert "$CERT_DIR/http_ca.crt" "$KIBANA_CA_TARGET"

chown root:elasticsearch "$ES_HTTP_P12_TARGET" "$ES_TRANSPORT_P12_TARGET" "$ES_HTTP_CA_TARGET" || true
chown -R root:kibana "$(dirname "$KIBANA_CRT_TARGET")" || true
chown kibana:kibana "$KIBANA_CA_TARGET" || true
log_success "Symlinks updated."

log_step "Computing new CA fingerprint & updating Kibana..."
NEW_FINGERPRINT=$(openssl x509 -in "$CERT_DIR/http_ca.crt" -noout -sha256 -fingerprint | cut -d'=' -f2 | tr -d ':' | tr '[:upper:]' '[:lower:]')
log_sub "New CA fingerprint: ${BOLD_GREEN}$NEW_FINGERPRINT${RESET}"

python3 -c "
import re
with open('$KIBANA_CONFIG', 'r') as f:
    content = f.read()
content = re.sub(r'ca_trusted_fingerprint:\s*[a-fA-F0-9]+', 'ca_trusted_fingerprint: $NEW_FINGERPRINT', content)
with open('$KIBANA_CONFIG', 'w') as f:
    f.write(content)
"
log_success "Kibana config updated."

log_section "Restarting Core Stack Services"

log_step "Restarting Elasticsearch service..."
systemctl restart elasticsearch

if ! spin_until "Waiting for Elasticsearch to start (port 9200)..." "curl -k -s https://localhost:9200" 30 6; then
  log_warn "Elasticsearch logs excerpt:"
  journalctl -u elasticsearch -n 15 --no-pager
fi

log_step "Restarting Kibana service..."
systemctl restart kibana

if ! spin_until "Waiting for Kibana to start (port 5601)..." "curl -k -s https://localhost:5601"; then
  log_warn "Kibana logs excerpt:"
  journalctl -u kibana -n 15 --no-pager
fi

log_section "Re-enrolling Elastic Agent (Fleet Server)"

log_step "Collecting pre-enrollment Elastic Agent IDs for cleanup..."
ES_SA_TOKEN=""
if [ -f "$KIBANA_CONFIG" ]; then
  ES_SA_TOKEN=$(grep -E '^elasticsearch.serviceAccountToken:' "$KIBANA_CONFIG" | awk '{print $2}' | tr -d '"' || true)
fi

PRE_AGENT_IDS=()
HOST_NAME=$(hostname)
if [ -n "$ES_SA_TOKEN" ]; then
  RAW_IDS=$(curl -k -s -H "Authorization: Bearer $ES_SA_TOKEN" -H 'Content-Type: application/json' -X POST 'https://localhost:9200/.fleet-agents/_search' -d "{\"query\":{\"term\":{\"local_metadata.host.hostname\":\"$HOST_NAME\"}}}" 2>/dev/null | python3 -c "import sys, json; data=json.load(sys.stdin); print(' '.join([h['_id'] for h in data.get('hits',{}).get('hits',[])]))" 2>/dev/null || true)
  for id in $RAW_IDS; do
    PRE_AGENT_IDS+=("$id")
  done
  log_sub "Captured ${#PRE_AGENT_IDS[@]} pre-existing agent ID(s) for host ${BOLD_CYAN}$HOST_NAME${RESET}."
fi

log_step "Preparing Fleet Server service token before touching Agent..."
TOKEN_NAME="cert-exchange-$(date +%s)"
FLEET_SERVICE_TOKEN=""

for attempt in {1..12}; do
  FLEET_SERVICE_TOKEN=$(/usr/share/elasticsearch/bin/elasticsearch-service-tokens create elastic/fleet-server "$TOKEN_NAME" 2>/dev/null | awk '{print $NF}' || true)
  if [ -n "$FLEET_SERVICE_TOKEN" ]; then
    break
  fi
  log_sub "Attempt $attempt/12: Elasticsearch initializing token store, waiting 5s..."
  sleep 5
done

if [ -z "$FLEET_SERVICE_TOKEN" ]; then
  log_error "Failed to generate Fleet Server service token after 12 attempts."
  exit 1
fi
log_success "Generated service token."

if [ -f "/etc/elasticsearch/service_tokens" ]; then
  chown root:elasticsearch "/etc/elasticsearch/service_tokens" || true
  chmod 660 "/etc/elasticsearch/service_tokens" || true
fi

log_step "Stopping Elastic Agent service..."
systemctl stop elastic-agent || true
log_success "Elastic Agent stopped."

log_step "Re-enrolling Elastic Agent (Fleet Server)..."
/opt/Elastic/Agent/elastic-agent enroll --force \
  --url="$FLEET_SERVER_URL" \
  --fleet-server-es="$FLEET_SERVER_ES_URL" \
  --fleet-server-es-ca="$ES_HTTP_CA_TARGET" \
  --fleet-server-service-token="$FLEET_SERVICE_TOKEN" \
  --fleet-server-policy="$FLEET_SERVER_POLICY" \
  --fleet-server-es-ca-trusted-fingerprint="$NEW_FINGERPRINT" \
  --fleet-server-port="$FLEET_SERVER_PORT" \
  --insecure >/dev/null 2>&1
log_success "Elastic Agent re-enrolled with new CA fingerprint."

log_step "Starting Elastic Agent service..."
systemctl start elastic-agent

spin_until "Waiting for Fleet Server endpoint (port 8220)..." "curl -k -s https://localhost:8220" || true

if [ ${#PRE_AGENT_IDS[@]} -gt 0 ] && [ -n "$ES_SA_TOKEN" ]; then
  log_step "Cleaning up former Agent entry/entries from Kibana Fleet UI..."
  NEW_AGENT_ID=$(curl -k -s -H "Authorization: Bearer $ES_SA_TOKEN" -H 'Content-Type: application/json' -X POST 'https://localhost:9200/.fleet-agents/_search' -d "{\"query\":{\"term\":{\"local_metadata.host.hostname\":\"$HOST_NAME\"}},\"sort\":[{\"last_checkin\":{\"order\":\"desc\"}}]}" 2>/dev/null | python3 -c "import sys, json; data=json.load(sys.stdin); hits=data.get('hits',{}).get('hits',[]); print(hits[0]['_id'] if hits else '')" 2>/dev/null || true)

  for old_id in "${PRE_AGENT_IDS[@]}"; do
    if [ "$old_id" != "$NEW_AGENT_ID" ]; then
      curl -k -s -H "Authorization: Bearer $ES_SA_TOKEN" -X DELETE "https://localhost:9200/.fleet-agents/_doc/$old_id" >/dev/null 2>&1 || true
      log_sub "Removed obsolete agent ID: ${DIM}$old_id${RESET}"
    fi
  done
  log_success "Duplicate/stale agent entries purged from Fleet UI."
fi

log_section "Final Service & Certificate Verification"

log_step "Elasticsearch (10.1.10.3:9200):"
if openssl s_client -connect 10.1.10.3:9200 -showcerts </dev/null 2>/dev/null | openssl x509 -noout -dates -subject -issuer; then
  log_success "Elasticsearch SSL endpoint verified."
else
  log_error "Failed to verify Elasticsearch SSL endpoint."
fi

log_step "Kibana (10.1.10.3:5601):"
if openssl s_client -connect 10.1.10.3:5601 -showcerts </dev/null 2>/dev/null | openssl x509 -noout -dates -subject -issuer; then
  log_success "Kibana SSL endpoint verified."
else
  log_error "Failed to verify Kibana SSL endpoint."
fi

log_step "Fleet Server (10.1.10.3:8220):"
if openssl s_client -connect 10.1.10.3:8220 -showcerts </dev/null 2>/dev/null | openssl x509 -noout -dates -subject -issuer; then
  log_success "Fleet Server SSL endpoint verified."
else
  log_error "Failed to verify Fleet Server SSL endpoint."
fi

echo -e "\n${BOLD_GREEN}====================================================${RESET}"
echo -e " ${BOLD_GREEN}✔ All certificates renewed and services active!${RESET}"
echo -e "${BOLD_GREEN}====================================================${RESET}\n"