#!/bin/bash
# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one
# or more contributor license agreements. Licensed under the Elastic License 2.0 as shown at
# https://securityonion.net/license; you may not use this file except in compliance with the
# Elastic License 2.0.

. /usr/sbin/so-common

MAX_JOBS=${MAX_ILM_JOBS:-10}
ILM_STATEFILE=/opt/so/state/esilm.txt
rm -f "$ILM_STATEFILE"

check_elasticsearch_responsive

ILM_TMP_DIR=$(mktemp -d)
trap 'rm -rf "$ILM_TMP_DIR"' EXIT
# Lock used to serialize block writes so concurrent jobs never interleave their output.
ILM_OUTPUT_LOCK="$ILM_TMP_DIR/output.lock"
ILM_FAIL_FILE="$ILM_TMP_DIR/failed"
touch "$ILM_FAIL_FILE"

# Policies are loaded concurrently (up to MAX_JOBS at a time) for speed. Each policy's block is
# printed the moment its curl returns, so output appears in COMPLETION ORDER, not the order
# policies are defined in configuration.
echo "Loading ILM policies concurrently; output below appears in completion order, not configuration order."
echo

put_policy() {
  local desc="$1" policyname="$2" data="$3" rc=1 attempt out http_code body result reason tmpdata
  tmpdata=$(mktemp "$ILM_TMP_DIR/${policyname}.XXXXXX.json")
  printf '%s' "$data" >"$tmpdata"

  for attempt in 1 2 3; do
    out=$(so-elasticsearch-query "_ilm/policy/${policyname}" -d@"${tmpdata}" -XPUT -w '\n%{http_code}' 2>/dev/null)
    http_code=${out##*$'\n'} body=${out%$'\n'*}

    if [[ "$http_code" == 200 ]] && jq -e '.acknowledged == true' <<<"$body" >/dev/null 2>&1; then
      rc=0
      break
    fi
    # 429 = ES rejected/throttled (may succeed later); other 4xx will not
    [[ "$http_code" == 4?? && "$http_code" != 429 ]] && break
    (( attempt < 3 )) && sleep 3
  done

  if (( rc == 0 )); then
    result='{"acknowledged":true}'
  else
    result=${body:-"HTTP ${http_code:-000}"}
    # flatten tabs/newlines so the TAB fail-file stays one record
    reason=$(jq -r '(if (.error|type) == "object" then (.error.reason // .error.type) elif (.error|type) == "string" then .error else empty end) // empty | tostring | gsub("[\t\r\n]+"; " ")' <<<"$body" 2>/dev/null)
    reason="HTTP ${http_code:-000}${reason:+: ${reason}}"
  fi

  # curl above ran in parallel; serialize just this block write so concurrent jobs never interleave.
  {
    flock 200
    printf 'Setting up %s policy...\n%s\n\n' "${desc}" "${result}"
    if (( rc != 0 )); then
      printf '%s\t%s\n' "${policyname}" "${reason}" >>"$ILM_FAIL_FILE"
    fi
  } 200>>"${ILM_OUTPUT_LOCK}"

  return "$rc"
}

# Block until fewer than MAX_JOBS background curls are running.
throttle() {
  while (( $(jobs -rp | wc -l) >= MAX_JOBS )); do
    wait -n || true
  done
}

{%- from 'elasticsearch/template.map.jinja' import ES_INDEX_SETTINGS %}
{%- if GLOBALS.role != "so-heavynode" %}
{%-   from 'elasticsearch/template.map.jinja' import ALL_ADDON_SETTINGS %}
{%- endif %}

{%- for index, settings in ES_INDEX_SETTINGS.items() %}
{%-   if settings.policy is defined %}
{%-     if index == 'so-logs-detections.alerts' %}
  throttle
  put_policy "so-logs-detections.alerts-so" "{{ index }}-so" '{ "policy": {{ settings.policy | tojson(true) }} }' &
{%-     elif index == 'so-logs-soc' %}
  throttle
  put_policy "so-soc-logs" "so-soc-logs" '{ "policy": {{ settings.policy | tojson(true) }} }' &
  throttle
  put_policy "{{ index }}-logs" "{{ index }}-logs" '{ "policy": {{ settings.policy | tojson(true) }} }' &
{%-     else %}
  throttle
  put_policy "{{ index }}-logs" "{{ index }}-logs" '{ "policy": {{ settings.policy | tojson(true) }} }' &
{%-     endif %}
{%-   endif %}
{%- endfor %}
{%- if GLOBALS.role != "so-heavynode" %}
{%-   for index, settings in ALL_ADDON_SETTINGS.items() %}
{%-     if settings.policy is defined %}
  throttle
  put_policy "{{ index }}-logs" "{{ index }}-logs" '{ "policy": {{ settings.policy | tojson(true) }} }' &
{%-     endif %}
{%-   endfor %}
{%- endif %}

wait || true

if [[ -s "$ILM_FAIL_FILE" ]]; then
  echo "ERROR: Failed to load ILM policy(s):"
  while IFS=$'\t' read -r POLICY REASON; do
    if [[ -n "$REASON" ]]; then
      echo " - ${POLICY}: ${REASON}"
    else
      echo " - ${POLICY}"
    fi
  done < "$ILM_FAIL_FILE"
  exit 1
else
  echo "Successfully loaded all ILM policies."
  touch "$ILM_STATEFILE"
fi
