#!/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=10

# Lock used to serialize block writes so concurrent jobs never interleave their output.
ILM_OUTPUT_LOCK=$(mktemp)
trap 'rm -f "$ILM_OUTPUT_LOCK"' 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
  result=$(curl -K /opt/so/conf/elasticsearch/curl.config -s -k -L \
    -X PUT "https://localhost:9200/_ilm/policy/${policyname}" \
    -H 'Content-Type: application/json' -d"${data}")
  # 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}"
  } 200>>"${ILM_OUTPUT_LOCK}"
}

# Block until fewer than MAX_JOBS background curls are running.
throttle() {
  while (( $(jobs -rp | wc -l) >= MAX_JOBS )); do
    wait -n
  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
