#!/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; you may not use
# this file except in compliance with the Elastic License 2.0.
{%- import_yaml 'elasticfleet/defaults.yaml' as ELASTICFLEETDEFAULTS %}
{% set SUB = salt['pillar.get']('elasticfleet:config:subscription_integrations', default=false) %}
{% set AUTO_UPGRADE_INTEGRATIONS = salt['pillar.get']('elasticfleet:config:auto_upgrade_integrations', default=false) %}
{%- set SUPPORTED_PACKAGES = salt['pillar.get']('elasticfleet:packages', default=ELASTICFLEETDEFAULTS.elasticfleet.packages, merge=True) %}

. /usr/sbin/so-common
. /usr/sbin/so-elastic-fleet-common

# Check that /opt/so/state/estemplates.txt exists to signal that Elasticsearch
#  has completed its first run of core-only integrations/indices/components/ilm
STATE_FILE_SUCCESS=/opt/so/state/estemplates.txt
INSTALLED_PACKAGE_LIST=/tmp/esfleet_installed_packages.json
BULK_INSTALL_PACKAGE_LIST=/tmp/esfleet_bulk_install.json
BULK_INSTALL_OUTPUT=/opt/so/state/esfleet_bulk_install_results.json
INTEGRATION_PACKAGE_COMPONENTS=/opt/so/state/esfleet_package_components.json
INPUT_PACKAGE_COMPONENTS=/opt/so/state/esfleet_input_package_components.json
CONTENT_PACKAGE_COMPONENTS=/opt/so/state/esfleet_content_package_components.json
COMPONENT_TEMPLATES=/opt/so/state/esfleet_component_templates.json

PENDING_UPDATE=false

# Integrations which are included in the package registry, but excluded from automatic installation via this script.
#   Requiring some level of manual Elastic Stack configuration before installation
EXCLUDED_INTEGRATIONS=('apm')

IFS=$'\n'
agent_policies=$(elastic_fleet_agent_policy_ids)
if [ $? -ne 0 ]; then
    echo "Error: Failed to retrieve agent policies."
    exit 1
fi

default_packages=({% for pkg in SUPPORTED_PACKAGES %}"{{ pkg }}"{% if not loop.last %} {% endif %}{% endfor %})

in_use_integrations=()

# Fetch each agent policy once; its package_policies[] already contain both the integration name
#  and the .package.name, so extract all non-default package names locally in a single jq instead
#  of re-fetching the same policy per integration.
default_packages_json=$(printf '%s\n' "${default_packages[@]}" | jq -R . | jq -s '.')
for AGENT_POLICY in $agent_policies; do

    if ! policy_json=$(fleet_api "agent_policies/$AGENT_POLICY"); then
        # skip the agent policy if we can't get required info, let salt retry. Integrations loaded by this script are non-default integrations.
        echo "Skipping $AGENT_POLICY.. "
        continue
    fi
    # non-default integrations that are in-use in any policy
    while IFS= read -r PACKAGE_NAME; do
        [ -n "$PACKAGE_NAME" ] && in_use_integrations+=("$PACKAGE_NAME")
    done < <(jq -r --argjson defaults "$default_packages_json" \
        '.item.package_policies[].package.name | select(. as $n | ($defaults | index($n)) | not)' \
        <<<"$policy_json")
done

if [[ -f $STATE_FILE_SUCCESS  ]]; then
    if retry 3 1 "curl -s -K /opt/so/conf/elasticsearch/curl.config --output /dev/null --silent --head --fail localhost:5601/api/fleet/epm/packages"; then
        # Package_list contains all integrations beta / non-beta.
        latest_package_list=$(/usr/sbin/so-elastic-fleet-package-list)
        echo '{ "packages" : []}' > $BULK_INSTALL_PACKAGE_LIST
        rm -f $INSTALLED_PACKAGE_LIST
        echo $latest_package_list | jq '{packages: [.items[] | {name: .name, latest_version: .version, installed_version: .installationInfo.version, subscription: .conditions.elastic.subscription }]}' >> $INSTALLED_PACKAGE_LIST

        # Build the bulk install list and the per-package status messages with two jq passes
        #  instead of a per-package bash loop. The old loop forked ~10 processes per package
        #  (5 jq + awk/bc for the version compare) and re-parsed/rewrote a growing JSON file on
        #  every add (O(n^2)). Selection and messages below are identical to that logic.
        SUB={% if SUB %}true{% else %}false{% endif %}
        AUTOUP={% if AUTO_UPGRADE_INTEGRATIONS %}true{% else %}false{% endif %}
        EXCLUDED_JSON=$(printf '%s\n' "${EXCLUDED_INTEGRATIONS[@]}" | jq -R 'select(length>0)' | jq -s '.')
        INUSE_JSON=$(printf '%s\n' "${in_use_integrations[@]}" | jq -R 'select(length>0)' | jq -s 'unique')

        # vnum replicates the previous version_conversion (%d%03d%03d of the first three dotted
        #  fields); needs() replicates the excluded/subscription/installed/upgrade/in-use logic.
        JQ_DECISION='
def vnum:
  [ (split(".")|.[0:3][] | gsub("[^0-9].*";"") | (if .=="" then "0" else . end) | tonumber) ]
  | (.[0]//0)*1000000 + (.[1]//0)*1000 + (.[2]//0);
def needs($sub;$autoup;$excluded;$inuse):
  .name as $n
  | ($n | IN($excluded[]) | not)
  and ( $sub or (.subscription==null or .subscription=="basic" or .subscription=="") )
  and ( (.installed_version==null or .installed_version=="")
        or ( ((.latest_version|vnum) > (.installed_version|vnum))
             and ( $autoup or ($n | IN($inuse[]) | not) ) ) );'

        JQ_ARGS=(--argjson sub "$SUB" --argjson autoup "$AUTOUP" --argjson excluded "$EXCLUDED_JSON" --argjson inuse "$INUSE_JSON")

        # (a) Per-package status messages (parity with the previous echo output).
        jq -r "${JQ_ARGS[@]}" "$JQ_DECISION"'
          .packages[]
          | .name as $n
          | if ($n|IN($excluded[])) then "Skipping \($n)..."
            elif (($sub|not) and (.subscription!=null and .subscription!="basic" and .subscription!="")) then
                 "\($n) integration requires an Elastic license of \(.subscription) or greater... skipping"
            elif (.installed_version==null or .installed_version=="") then
                 "\($n) is not installed... Adding to next update."
            elif ((.latest_version|vnum) > (.installed_version|vnum)) then
                 (if ($autoup or ($n|IN($inuse[])|not))
                  then "\($n) is at version \(.installed_version) latest version is \(.latest_version)... Adding to next update."
                  else "skipping available upgrade for in use integration - \($n)." end)
            else empty end
        ' "$INSTALLED_PACKAGE_LIST"

        # (b) The bulk install list, built in a single pass.
        jq "${JQ_ARGS[@]}" "$JQ_DECISION"'
          {packages: [ .packages[] | select(needs($sub;$autoup;$excluded;$inuse)) | {name, version: .latest_version} ]}
        ' "$INSTALLED_PACKAGE_LIST" > "$BULK_INSTALL_PACKAGE_LIST"

        if jq -e '.packages | length > 0' "$BULK_INSTALL_PACKAGE_LIST" >/dev/null; then
            PENDING_UPDATE=true
        fi

        if [ "$PENDING_UPDATE" = true ]; then
            # Run chunked install of packages
            echo "" > $BULK_INSTALL_OUTPUT
            pkg_group=1
            pkg_filename="${BULK_INSTALL_PACKAGE_LIST%.json}"

            jq -c '.packages | _nwise(25)' $BULK_INSTALL_PACKAGE_LIST | while read -r line; do
                echo "$line" | jq '{ "packages": . }' > "${pkg_filename}_${pkg_group}.json"
                pkg_group=$((pkg_group + 1))
            done

            for file in "${pkg_filename}_"*.json; do
                [ -e "$file" ] || continue
                if ! elastic_fleet_bulk_package_install $file >> $BULK_INSTALL_OUTPUT; then
                    # integrations loaded my this script are non-essential and shouldn't cause exit, skip them for now next highstate run can retry
                    echo "Failed to complete a chunk of bulk package installs -- $file "
                    continue
                fi
            done
            # cleanup any temp files for chunked package install
            rm -f ${pkg_filename}_*.json $BULK_INSTALL_PACKAGE_LIST
        else
            echo "Elastic integrations don't appear to need installation/updating..."
        fi
        # Write out file for generating index/component/ilm templates, keeping each package type separate
        for package_type in "INTEGRATION" "INPUT" "CONTENT"; do
            if latest_installed_package_list=$(elastic_fleet_installed_packages_components "$package_type"); then
                outfile="${package_type}_PACKAGE_COMPONENTS"
                echo $latest_installed_package_list > "${!outfile}"
            fi
        done
        if retry 3 1 "so-elasticsearch-query / --fail --output /dev/null"; then
            # Refresh installed component template list
            latest_component_templates_list=$(so-elasticsearch-query _component_template | jq '.component_templates[] | .name' | jq -s '.')
            echo $latest_component_templates_list > $COMPONENT_TEMPLATES
        fi

    else
        # This is the installation of add-on integrations and upgrade of existing integrations. Exiting without error, next highstate will attempt to re-run.
        echo "Elastic Fleet does not appear to be responding... Exiting... "
        exit 0
    fi
else
    # This message will appear when an update to core integration is made and this script is run at the same time as
    #  elasticsearch.enabled -> detects change to core index settings -> deletes estemplates.txt
    echo "Elasticsearch may not be fully configured yet or is currently updating core index settings."
    exit 0
fi