#!/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}

# Lock used to serialize block writes so concurrent jobs never interleave their output.
ILM_OUTPUT_LOCK=$(mktemp)
ILM_FAIL_FILE=$(mktemp)
trap 'rm -f "$ILM_OUTPUT_LOCK" "$ILM_FAIL_FILE"' EXIT

# 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" result rc=0
  if ! result=$(curl -K /opt/so/conf/elasticsearch/curl.config -s -k -L --fail \
    -X PUT "https://localhost:9200/_ilm/policy/${policyname}" \
    -H 'Content-Type: application/json' -d"${data}" 2>&1); then
    rc=1
  elif ! jq -e '.acknowledged == true' <<<"$result" >/dev/null 2>&1; then
    rc=1
  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\n' "${policyname}" >>"$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 read -r POLICY; do
    echo " - $POLICY"
  done < "$ILM_FAIL_FILE"
  exit 1
else
  echo "Successfully loaded all ILM policies."
fi
