{%- if grains['role'] == 'so-master' or grains['role'] == 'so-eval' %}
{%- set ip = salt['pillar.get']('static:masterip', '') %}
{%- elif grains['role'] == 'so-node' %} 
{%- set ip = salt['pillar.get']('node: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="localhost"
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
}

register_agent() {
  # Adding agent and getting Id from manager
  echo ""
  echo "Adding agent:"
  echo "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 $PROTOCOL://$API_IP:$API_PORT/agents)
  echo -e $API_RESULT | grep -q "\"error\":0" 2>&1

  if [ "$?" != "0" ]; then
    echo -e $API_RESULT | sed -rn 's/.*"message":"(.+)".*/\1/p'
    exit 0
  fi
  # Get agent id and agent key 
  AGENT_ID=$(echo $API_RESULT | cut -d':' -f 4 | cut -d ',' -f 1)
  AGENT_KEY=$(echo $API_RESULT | cut -d':' -f 5 | cut -d '}' -f 1)

  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

  exit 0
}

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 $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 $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

# Default action -> try to register the agent
sleep 10s
STATUS=$(curl -s -k -u $USER:$PASSWORD $PROTOCOL://$API_IP:$API_PORT/agents/$AGENT_ID | jq .data.status | sed s'/"//g')
if [[ $STATUS == "Active" ]]; then 
  echo "Agent $AGENT_ID already registered!"
else
  register_agent
fi
#remove_agent