{%- if grains['role'] in ['so-manager', 'so-eval', 'so-managersearch', 'so-standalone'] %}
  {%- set ip = salt['pillar.get']('global:managerip', '') %}
{%- elif grains['role'] == 'so-node' or grains['role'] == 'so-heavynode' %}
  {%- set ip = salt['pillar.get']('elasticsearch:mainip', '') %}
{%- elif grains['role'] == 'so-sensor' %}
  {%- set ip = salt['pillar.get']('sensor:mainip', '') %}
{%- endif %}
#!/bin/bash

###
#  Shell script for registering agents automatically with the API
#  Copyright (C) 2017 Wazuh, Inc. All rights reserved.
#  Wazuh.com
#
#  This program is a free software; you can redistribute it
#  and/or modify it under the terms of the GNU General Public
#  License (version 2) as published by the FSF - Free Software
#  Foundation.
###
#
#  12/11/2018
#  This script has been modified by Security Onion Solutions
#  - Added Agent IP variable and option
###

# Connection variables
API_IP="{{ ip }}"
API_PORT="55000"
PROTOCOL="https"
USER="foo"
PASSWORD="bar"
AGENT_NAME=$(hostname)
AGENT_IP="{{ip}}"
AGENT_ID=001

display_help() {
cat <<HELP_USAGE

    $0  [-h] [-f|--force] [-q|--quiet] [agent]

   -h             Show this message.
   -f|--force     Force agent removal (if already registered)
                  The agent will be re-regitered with a new ID
   -s|--silent    Surpress the output while removing the agent
   agent          Agent name (if missing we will use the output
                  of the hostname command)
HELP_USAGE
}

cleanup_creds() {
  /usr/sbin/so-wazuh-user-remove $USER
}

register_agent() {
  # Adding agent and getting Id from manager
  echo ""
  echo "Adding agent:"
  echo "Executing: curl -s -u $USER:**** -k -X POST -d 'name=$AGENT_NAME&ip=$AGENT_IP' $PROTOCOL://$API_IP:$API_PORT/agents"
  API_RESULT=$(curl -s -u $USER:"$PASSWORD" -k -X POST -d 'name='$AGENT_NAME'&ip='$AGENT_IP -L $PROTOCOL://$API_IP:$API_PORT/agents)
  # Get agent id and key
  AGENT_ID=$(echo "$API_RESULT" | jq -er ".data.id")
  GOT_ID=$?
  AGENT_KEY=$(echo "$API_RESULT" | jq -er ".data.key")
  GOT_KEY=$?

  if [[ -z "$AGENT_ID" || -z "$AGENT_KEY" || $GOT_ID -ne 0 || $GOT_KEY -ne 0 ]]; then
    echo "Failed Result: $API_RESULT"
    return 1
  else
    echo "Agent '$AGENT_NAME' with ID '$AGENT_ID' added."
    echo "Key for agent '$AGENT_ID' received."

    # Importing key
    echo ""
    echo "Importing authentication key:"
    echo "y" | /var/ossec/bin/manage_agents -i '$AGENT_KEY'

    # Restarting agent
    echo ""
    echo "Restarting:"
    echo ""
    /var/ossec/bin/ossec-control restart
    return 0
  fi
}

wait_for_manager() {
  echo "Waiting for Wazuh manager to become ready..."

  maxAttempts=$1
  attempts=0
  while [[ $attempts -lt $maxAttempts ]]; do
    attempts=$((attempts+1))
    AGENTS_OUTPUT=$(curl -s -u $USER:"$PASSWORD" -k -X GET -L $PROTOCOL://$API_IP:$API_PORT/agents)
    MANAGER_STATUS=$(echo "$AGENTS_OUTPUT" | jq -r ".data.items[0].status")
    if [ "$MANAGER_STATUS" == "Active" ]; then
      echo "Wazuh manager is active, ready to proceed."
      return 0
    else
      echo "Received non-Active status response: "
      echo "$AGENTS_OUTPUT"
      echo
      echo "Manager is not ready after attempt $attempts of $maxAttempts, sleeping for 30 seconds."
      sleep 30
    fi
  done
  return 1
}

remove_agent() {
  echo "Found: $AGENT_ID"
  echo "Removing previous registration for '$AGENT_NAME' using ID: $AGENT_ID ..."
  # curl -u foo:bar -k -X DELETE "https://127.0.0.1:55000/agents/001
  REMOVE_AGENT=$(curl -s -u $USER:"$PASSWORD" -k -X DELETE -L $PROTOCOL://$API_IP:$API_PORT/agents/$AGENT_ID)
  echo -e $REMOVE_AGENT
}

get_agent_id() {
  echo ""
  echo "Checking for Agent ID..."
  AGENT_ID=$(curl -s -u $USER:"$PASSWORD" -k -X GET -L $PROTOCOL://$API_IP:$API_PORT/agents/name/$AGENT_NAME | rev | cut -d: -f1 | rev | grep -o '".*"' | tr -d '"')
}

# MAIN
# ENTRY POINT

while getopts ':hfsi:' OPTION; do
  case "$OPTION" in
    h)
      display_help
      exit 0
      ;;
    f|--force)
      FORCE=true
      ;;
    i|--ip)
      AGENT_IP=${OPTARG}
      ;;
    s|--silent)
      SILENT=true
      ;;
  esac
done
# reset $1, $2 .... as normal argument after the flag
shift $(($OPTIND - 1))

# if no arguments are passed in after the flags, we assign the hostname value to the AGENT_NAME
#AGENT_NAME=${1:-$(hostname)}

#get_agent_id

# check the return value. If we get an integer back then the agent is already registered. Anything else -> agent is not registered
#  if ! [ "$AGENT_ID" -eq "$AGENT_ID" ] 2> /dev/null ; then
#   echo "Starting registration process ..."
#   :
#  elif [[ "$FORCE" = true && "$SILENT" = "true" ]] ; then
#   remove_agent > /dev/null 2>&1
#  else
#    if [[ "$FORCE" = true ]] ; then
#      remove_agent
#    fi
#  fi

if [ -f /opt/so/conf/wazuh/initial_agent_registration.log ]; then
  echo "Agent $AGENT_ID already registered!"
  exit 0
else
  retries=30
  if wait_for_manager $retries; then
    if register_agent; then
      cleanup_creds
      echo "Initial agent $AGENT_ID with IP $AGENT_IP registered on $DATE." > /opt/so/conf/wazuh/initial_agent_registration.log
      exit 0
    else
      echo "ERROR: Failed to register agent"
    fi
  else
    echo "ERROR: Wazuh manager did not become ready after $retries attempts; unable to proceed with registration"
  fi
fi

exit 1