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

DEFAULT_SALT_DIR=/opt/so/saltstack/default

if [ -z $NOROOT ]; then
        # Check for prerequisites
        if [ "$(id -u)" -ne 0 ]; then
                echo "This script must be run using sudo!"
                exit 1
        fi
fi

# Ensure /usr/sbin is in path
if ! echo "$PATH" | grep -q "/usr/sbin"; then
  export PATH="$PATH:/usr/sbin"
fi

# Define a banner to separate sections
banner="========================================================================="

fleet_api() {
    local QUERYPATH=$1
    shift

    curl -sK /opt/so/conf/elasticsearch/curl.config -L "localhost:5601/api/fleet/${QUERYPATH}" "$@" --retry 3 --retry-delay 10 --fail 2>/dev/null
}

# Max number of concurrent Fleet write jobs (create/update). Override via env if needed.
MAX_FLEET_JOBS=${MAX_FLEET_JOBS:-10}

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

# Load every integration JSON in a directory into a single agent policy.
# The agent policy is fetched ONCE (not per file), and the create/update writes
# are dispatched as throttled background jobs.
#   $1 AGENT_POLICY     - the agent policy id/name to load integrations into
#   $2 DIR              - directory of integration *.json files
#   $3 LABEL           - human-readable label for log output
#   $4 SKIP_CREATE_NAME - (optional) integration name to skip when creating (still updated if present)
# Returns 1 if any integration failed to create/update.
elastic_fleet_load_integrations_dir() {
    local AGENT_POLICY=$1
    local DIR=$2
    local LABEL=$3
    local SKIP_CREATE_NAME=$4
    local POLICY_JSON FAIL_FILE OUT_DIR INTEGRATION NAME ID i

    FAIL_FILE=$(mktemp)
    # Each job buffers its full output (header + API response) into its own file so the
    # parent can print them grouped and in submission order after concurrent writes finish.
    OUT_DIR=$(mktemp -d)
    i=0

    # Fetch the agent policy a single time; we look up integration ids locally below.
    POLICY_JSON=$(fleet_api "agent_policies/$AGENT_POLICY")

    for INTEGRATION in "$DIR"/*.json; do
        [ -e "$INTEGRATION" ] || continue
        NAME=$(jq -r .name "$INTEGRATION")
        ID=$(jq -r --arg n "$NAME" '.item.package_policies[]? | select(.name==$n) | .id' <<<"$POLICY_JSON")

        elastic_fleet_throttle
        {
            local RESP
            if [ -n "$ID" ]; then
                printf "\n\n%s - Updating integration %s\n" "$LABEL" "$NAME"
                if ! RESP=$(elastic_fleet_integration_update "$ID" "@$INTEGRATION"); then
                    flock 9; echo "update ${INTEGRATION##*/}" >&9
                fi
                printf '%s\n' "$RESP"
            elif [ -n "$SKIP_CREATE_NAME" ] && [ "$NAME" == "$SKIP_CREATE_NAME" ]; then
                printf "\n\n%s - Skipping creation of %s\n" "$LABEL" "$NAME"
            else
                printf "\n\n%s - Creating integration %s\n" "$LABEL" "$NAME"
                if ! RESP=$(elastic_fleet_integration_create "@$INTEGRATION"); then
                    flock 9; echo "create ${INTEGRATION##*/}" >&9
                fi
                printf '%s\n' "$RESP"
            fi
        } >"$OUT_DIR/$(printf '%03d' "$i")" 9>>"$FAIL_FILE" &
        i=$((i+1))
    done
    wait

    # Emit per-integration output grouped and in submission order (glob sorts numerically).
    cat "$OUT_DIR"/* 2>/dev/null
    rm -rf "$OUT_DIR"

    local rc=0
    if [ -s "$FAIL_FILE" ]; then
        printf "\n%s: failed integrations:\n" "$LABEL"
        cat "$FAIL_FILE"
        rc=1
    fi
    rm -f "$FAIL_FILE"
    return $rc
}

elastic_fleet_integration_check() {

    AGENT_POLICY=$1

    JSON_STRING=$2

    NAME=$(jq -r .name $JSON_STRING)

    INTEGRATION_ID=$(/usr/sbin/so-elastic-fleet-agent-policy-view "$AGENT_POLICY" | jq -r '.item.package_policies[] | select(.name=="'"$NAME"'") | .id')

}

elastic_fleet_integration_create() {

    JSON_STRING=$1

    # --retry-all-errors so transient 409 conflicts (concurrent writes to the same agent
    # policy) are retried; curl --retry alone does not retry 409.
    if ! fleet_api "package_policies" --retry-all-errors -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -XPOST -d "$JSON_STRING"; then
        return 1
    fi
}


elastic_fleet_integration_remove() {

    AGENT_POLICY=$1

    NAME=$2

    INTEGRATION_ID=$(/usr/sbin/so-elastic-fleet-agent-policy-view "$AGENT_POLICY" | jq -r '.item.package_policies[] | select(.name=="'"$NAME"'") | .id')

    JSON_STRING=$( jq -n \
                    --arg INTEGRATIONID "$INTEGRATION_ID" \
                    '{"packagePolicyIds":[$INTEGRATIONID]}'
                    )

    if ! fleet_api "package_policies/delete" -XPOST -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING"; then
        echo "Error: Unable to delete '$NAME' from '$AGENT_POLICY'"
        return 1
    fi
}

elastic_fleet_integration_update() {

    UPDATE_ID=$1

    JSON_STRING=$2

    # --retry-all-errors so transient 409 conflicts (concurrent writes to the same agent
    # policy) are retried; curl --retry alone does not retry 409.
    if ! fleet_api "package_policies/$UPDATE_ID" --retry-all-errors -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -XPUT -d "$JSON_STRING"; then
        return 1
    fi
}

elastic_fleet_integration_policy_upgrade() {

    INTEGRATION_ID=$1

    JSON_STRING=$( jq -n \
                    --arg INTEGRATIONID "$INTEGRATION_ID" \
                    '{"packagePolicyIds":[$INTEGRATIONID]}'
                    )

    if ! fleet_api "package_policies/upgrade" -XPOST -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING"; then
        return 1
    fi
}


elastic_fleet_package_version_check() {
    PACKAGE=$1

    if output=$(fleet_api "epm/packages/$PACKAGE"); then
        echo "$output" | jq -r '.item.version'
    else
        echo "Error: Failed to get current package version for '$PACKAGE'"
        return 1
    fi
}

elastic_fleet_package_latest_version_check() {
    PACKAGE=$1
    if output=$(fleet_api "epm/packages/$PACKAGE"); then
        if version=$(jq -e -r '.item.latestVersion' <<< $output); then
            echo "$version"
        fi
    else
        echo "Error: Failed to get latest version for '$PACKAGE'"
        return 1
    fi
}

elastic_fleet_package_install() {
    PKG=$1
    VERSION=$2
    if ! fleet_api "epm/packages/$PKG/$VERSION" -XPOST -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d '{"force":true}'; then
        return 1
    fi
}

elastic_fleet_bulk_package_install() {
    BULK_PKG_LIST=$1
    if ! fleet_api "epm/packages/_bulk" -XPOST -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d@$BULK_PKG_LIST; then
        return 1
    fi
}

elastic_fleet_get_package_list_by_type() {
    if ! output=$(fleet_api "epm/packages"); then
        return 1
    else
        is_integration=$(jq '[.items[] | select(.type=="integration") | .name ]' <<< "$output")
        is_input=$(jq '[.items[] | select(.type=="input") | .name ]' <<< "$output")
        is_content=$(jq '[.items[] | select(.type=="content") | .name ]' <<< "$output")
        jq -n --argjson is_integration "${is_integration:-[]}" \
              --argjson is_input "${is_input:-[]}" \
              --argjson is_content "${is_content:-[]}" \
              '{"integration": $is_integration,"input": $is_input, "content": $is_content}'
    fi
}
elastic_fleet_installed_packages_components() {
    package_type=${1,,}
    if [[ "$package_type" != "integration" && "$package_type" != "input" && "$package_type" != "content" ]]; then
        echo "Error: Invalid package type ${package_type}. Valid types are 'integration', 'input', or 'content'."
        return 1
    fi

    packages_by_type=$(elastic_fleet_get_package_list_by_type)
    packages=$(jq --arg package_type "$package_type" '.[$package_type]' <<< "$packages_by_type")

    if ! output=$(fleet_api "epm/packages/installed?perPage=500"); then
        return 1
    else
        jq -c --argjson packages "$packages" '[.items[] | select(.name | IN($packages[])) | {name: .name, dataStreams: .dataStreams}]' <<< "$output"
    fi
}

elastic_fleet_agent_policy_ids() {
    if output=$(fleet_api "agent_policies"); then
        echo "$output" | jq -r .items[].id
    else
        echo "Error: Failed to retrieve agent policies."
        return 1
    fi
}

elastic_fleet_integration_policy_names() {
    AGENT_POLICY=$1
    if output=$(fleet_api "agent_policies/$AGENT_POLICY"); then
        echo "$output" | jq -r .item.package_policies[].name
    else
        echo "Error: Failed to retrieve integrations for '$AGENT_POLICY'."
        return 1
    fi
}

elastic_fleet_integration_policy_package_name() {
    AGENT_POLICY=$1
    INTEGRATION=$2
    if output=$(fleet_api "agent_policies/$AGENT_POLICY"); then
        echo "$output" | jq -r --arg INTEGRATION "$INTEGRATION" '.item.package_policies[] | select(.name==$INTEGRATION)| .package.name'
    else
        echo "Error: Failed to retrieve package name for '$INTEGRATION' in '$AGENT_POLICY'."
        return 1
    fi
}

elastic_fleet_integration_policy_package_version() {
    AGENT_POLICY=$1
    INTEGRATION=$2

    if output=$(fleet_api "agent_policies/$AGENT_POLICY"); then
        if version=$(jq -e -r --arg INTEGRATION "$INTEGRATION" '.item.package_policies[] | select(.name==$INTEGRATION)| .package.version' <<< "$output"); then
            echo "$version"
        fi
    else
        echo "Error: Failed to retrieve integration version for '$INTEGRATION' in policy '$AGENT_POLICY'"
        return 1
    fi
}

elastic_fleet_integration_id() {
    AGENT_POLICY=$1
    INTEGRATION=$2
    if output=$(fleet_api "agent_policies/$AGENT_POLICY"); then
        echo "$output" | jq -r --arg INTEGRATION "$INTEGRATION" '.item.package_policies[] | select(.name==$INTEGRATION)| .id'
    else
        echo "Error: Failed to retrieve integration ID for '$INTEGRATION' in '$AGENT_POLICY'."
        return 1
    fi
}

elastic_fleet_integration_policy_dryrun_upgrade() {
    INTEGRATION_ID=$1
    if ! fleet_api "package_policies/upgrade/dryrun" -H "Content-Type: application/json" -H 'kbn-xsrf: true' -XPOST -d "{\"packagePolicyIds\":[\"$INTEGRATION_ID\"]}"; then
        echo "Error: Failed to complete dry run for '$INTEGRATION_ID'."
        return 1
    fi
}

elastic_fleet_policy_create() {

    NAME=$1
    DESC=$2
    FLEETSERVER=$3
    TIMEOUT=$4

    JSON_STRING=$( jq -n \
        --arg NAME "$NAME" \
        --arg DESC "$DESC" \
        --arg TIMEOUT $TIMEOUT \
        --arg FLEETSERVER "$FLEETSERVER" \
            '{"name": $NAME,"id":$NAME,"description":$DESC,"namespace":"default","monitoring_enabled":["logs"],"inactivity_timeout":$TIMEOUT,"has_fleet_server":$FLEETSERVER,"advanced_settings":{"agent_logging_level": "warning"}}'
        )
    # Create Fleet Policy
    if ! fleet_api "agent_policies" -XPOST -H 'kbn-xsrf: true' -H 'Content-Type: application/json' -d "$JSON_STRING"; then
        return 1
    fi

}
