#!/bin/bash
# Atomicorp API: Server CLI
# Copyright Atomicorp 2025
# All rights reserved

# Global variables
VERSION=1.1.0
SCANID=$(date +%s)
DEBUG=0
LOG_FILE="/var/awp/log/atomicorp-api.log"

function log_event() {
    if [ "$DEBUG" -ge 1 ]; then
        if [ ! -d "/var/awp/log" ]; then
            mkdir -p "/var/awp/log"
        fi
        echo "$(date) atomicorp-api: $1" >> "$LOG_FILE"
    fi
}


function show_help() {
	echo
	echo "Atomicorp API CLI"
	echo "Version: ${VERSION}"
	echo "Usage: $0 -l|-lg|-a|-g group|-i id -i <package>|-u <package> |-x <package>"
	echo
	echo "  example: $0 -a -u atomicorp"
	echo
	echo 
	echo "Command line paramenters"
	echo
	echo "  List agents or groups"
	echo "    -l				List Agents"
	echo "    -lg 				List Groups"
	echo
	echo "  Target requires one of the following"
	echo "    -a 				All agents"
	echo "    -g <group> 			All agents in group <group>"
	echo "    -i <id>			Specified Agent ID"
	echo 
	echo "  Action requires one of the following"
	echo "    -init <package>		Configure package (first time setup)"
	echo "		auditd 			- configure/update/start auditd"
	echo "		clamav 			- configure/update/start clamd"
	echo "		fapolicyd 		- configure/update/start fapolicyd"
	echo "    -install <package>		Install package where package is:"
	echo "		auditd 			- installs auditd"
	echo "		clamav 			- installs clamav"
	echo "		fapolicyd 		- installs fapolicyd"
	echo "    -restart <package>		Restart package"
	echo "		clamav 			- restart clamd"
	echo "		ossec-hids		- restart ossec"
	echo "    -scan <module> <option>	Scan module"
	echo "		clamav 			- scan <path>"
	echo "    -update <package>		Update package where package is:"
	echo "		clamav 			- update clamav sigantures"
	echo "		ossec-hids		- update atomic ossec"
	echo "    -isolate <state>      Isolate agent (requires -i flag)"
	echo "      enable          - enable network isolation"
	echo "      disable         - disable network isolation"
	echo
	echo "  Optional"
	echo "    -h					Show this help"
	echo "    -r					Realtime flag, shows responses in real time"
	echo "    -t <timeout>			Timeout in seconds (default 30)"
	echo
	exit 1


}

function tail_alerts() {
    AGENT_ID=$1
    SEARCH=$2
    SCANID=$3

	# if TIMEOUT is not set, set it to 30
	if [ -z "${TIMEOUT}" ]; then
		TIMEOUT=30
	fi
    timeout ${TIMEOUT} tail -f /var/ossec/logs/alerts/alerts.json  | while read line; do
        # if json field "predecoder.program_name" is "atomicorp-api" and json field "agent.id" is "${AGENT_ID}" then print the fields "agent.name", "agent.id" and "data.atomicorp.msg"
        if [[ $(echo $line | jq -r '.predecoder.program_name') == "atomicorp-api" ]] && [[ $(echo $line | jq -r '.agent.id') == ${AGENT_ID} ]] && [[ $(echo $line | jq -r '.data.atomicorp.scanid') == ${SCANID} ]]; then
			echo -n "  "
            echo $(echo $line | jq -r '.agent.name') $(echo $line | jq -r '.agent.id') $(echo $line | jq -r '.data.atomicorp.msg') $(echo $line | jq -r '.data.atomicorp.scanid')
		fi
	done
}

# Create a function that runs the tail_alerts function in the background, and then kills it after the value $TIMEOUT is reached
function run_tail_alerts() {
	# if TIMEOUT is not set, set it to 30
	if [ -z "${TIMEOUT}" ]; then
		TIMEOUT=30
	fi
    AGENT_ID=$1
    SEARCH=$2
    SCANID=$3
	# tell us what the SEARCH string is
	echo "  Searching for ${SEARCH}"
    tail_alerts ${AGENT_ID} ${SEARCH} ${SCANID} &
    tail_pid=$!
    sleep ${TIMEOUT}
    kill $tail_pid
}


function install_package() {
	ID=$1
	PACKAGE=$2

	# if REALTIME is set then we run the run_tail_alerts function in the background and pass the agent id and the search string
	if [ -n "${REALTIME}" ]; then
		# print that the action is being run in realtime
		echo
		echo "Running in realtime mode: "	
		run_tail_alerts ${ID} install_package:success ${SCANID} &
	fi

    if [[ ${PACKAGE} =~ auditd* ]]; then
		 log_event "install:auditd id:${ID} scanid:${SCANID}"
		 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "install:auditd scanid:${SCANID}" >/dev/null
	elif [[ ${PACKAGE} =~ clamav* ]]; then
		 log_event "install:clamav id:${ID} scanid:${SCANID}"
		 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "install:clamav scanid:${SCANID}" >/dev/null
	elif [[ ${PACKAGE} =~ fapolicyd* ]]; then
		 log_event "install:fapolicyd id:${ID} scanid:${SCANID}"
		 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "install:fapolicyd scanid:${SCANID}" >/dev/null
	else
		 log_event "install:ERROR:${PACKAGE} installation not supported"
		echo "Error: ${PACKAGE} installation not supported"
		exit 1
	fi

	# if REALTIME is set then we wait for the run_tail_alerts function to finish
	if [ -n "${REALTIME}" ]; then
		wait
		# print that the action is finished
		echo
		echo "Finished running in realtime mode"
	fi	

}

function update_package() {
        ID=$1
        PACKAGE=$2
        if [[ ${PACKAGE} =~ ossec-hids* ]] || [[ ${PACKAGE} == "user-check-ossec" ]]; then
                 log_event "upgrade:ossec-hids id:${ID}"
                 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "upgrade:user-check-ossec" >/dev/null
        elif [[ ${PACKAGE} =~ clamav* ]] || [[ ${PACKAGE} == "user-check-clam" ]]; then
                 log_event "upgrade:clamav id:${ID}"
                 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "upgrade:user-check-clam" >/dev/null
        elif [[ ${PACKAGE} == "user-check-all" ]]; then
                 log_event "upgrade:all id:${ID}"
                 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "upgrade:user-check-all" >/dev/null
        else
                log_event "upgrade:ERROR:${PACKAGE} installation not supported"
                echo "Error: ${PACKAGE} installation not supported"
                exit 1
        fi
}


function init_package() {
        ID=$1
        PACKAGE=$2
        if [[ ${PACKAGE} =~ auditd* ]]; then
                 log_event "init:auditd id:${ID}"
                 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "init:auditd" >/dev/null
        elif [[ ${PACKAGE} =~ clamav* ]]; then
                 log_event "init:clamav id:${ID}"
                 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "init:clamav" >/dev/null
        elif [[ ${PACKAGE} =~ fapolicyd* ]]; then
                 log_event "init:fapolicyd id:${ID}"
                 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "init:fapolicyd" >/dev/null
        else
                log_event "init:ERROR:${PACKAGE} initialization not supported"
                echo "Error: ${PACKAGE} initilization not supported"
                exit 1
        fi
}

# Create the restart_package function
function restart_package() {
		ID=$1
		PACKAGE=$2
		if [[ ${PACKAGE} =~ ossec-hids* ]]; then
				log_event "restart:ossec-hids id:${ID}"
				 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "restart:ossec-hids" >/dev/null
		elif [[ ${PACKAGE} =~ clamav* ]]; then
				log_event "restart:clamav id:${ID}"
				 /var/ossec/bin/agent_control -f atomicorp-api0 -u ${ID} -b "restart:clamav" >/dev/null
		else
				log_event "restart:ERROR:${PACKAGE} installation not supported"
				echo "Error: ${PACKAGE} installation not supported"
				exit 1
		fi
}

# create the scan_package function
function scan_package() {
		ID=$1
		PACKAGE=$2
		if [[ ${PACKAGE} =~ clamav* ]]; then
			log_event "scan:clamav id:${ID} path:${SCANPATH}"
			/var/awp/bin/malware-scan -i ${ID} -s "${SCANPATH}" >/dev/null
		else
			log_event "scan:ERROR:${PACKAGE} installation not supported"
			echo "Error: ${PACKAGE} installation not supported"
			exit 1
		fi
}

# create the isolate_agent function
function isolate_agent() {
    ID=$1
    STATE=$2
    
    # Verify we're using -i flag
    if [ -z "$ID" ] || [ -n "$ALL" ] || [ -n "$GROUP" ]; then
        echo "Error: -isolate can only be used with the -i flag"
        exit 1
    fi

    case "$STATE" in
        "enable")
            log_event "isolate:enable id:${ID}"
            /var/ossec/bin/agent_control -f atomic-isolate0 -u ${ID} -b "isolate:enable" >/dev/null
            ;;
        "disable")
            log_event "isolate:disable id:${ID}"
            /var/ossec/bin/agent_control -f atomic-isolate0 -u ${ID} -b "isolate:disable" >/dev/null
            ;;
        *)
            log_event "isolate:ERROR:Invalid isolation state:${STATE}"
            echo "Error: Invalid isolation state. Use 'enable' or 'disable'"
            exit 1
            ;;
    esac
}

if [ ! $1 ]; then
	show_help
fi


while [ $# -ge 1 ]; do
   case "$1" in
	# Scan every node
	-a) 
		ALL=1
		;;

	# Scan all in a group
	-g)
		GROUP=""

		case "$2" in
			"")
				;;
			-*)
				;;
			*)
				shift
				GROUP="$1"
				;;
		esac
		if [ -z "${GROUP}" ]; then
			echo "No group specified with -g option"
			exit 1
		fi

		;;
	# Scan a specific node
      	-i) 
		ID=""
		case "$2" in 
			"")
                ;;
            -*)
                ;;
            *)
                shift
                ID="$1"
                ;;
        esac
        if [ -z "${ID}" ]; then
                echo "No agent specified with -i option"
                exit 1
        fi

        ;;

	# realtime flag
	-r)
		REALTIME=1
		;;
	# timeout flag
	-t)
		TIMEOUT=""
		case "$2" in
			"")
				;;
			-*)
				;;
			*)
				shift
				TIMEOUT="$1"
				;;
		esac
		if [ -z "${TIMEOUT}" ]; then
			echo "No timeout specified with -t option"
			exit 1
		fi
		;;
	# List agents
   	-l) 
		LIST=1
		;;
	# List groups
   	-lg) 
		LIST_GROUP=1
		;;
	# Install
   	-install) 
	    PACKAGE=""
        case "$2" in
            "")
                    ;;
            -*)
                    ;;
            *)
                    shift
                    PACKAGE="$1"
                    ;;
        esac
        if [ -z "${PACKAGE}" ]; then
                echo "No Package specified with -install option"
                exit 1
        fi

		INSTALL=1
        ;;

	# Update a target
	-update)
        PACKAGE=""
        case "$2" in
            "")
                ;;
            -*)
                ;;
            *)
                shift
                PACKAGE="$1"
                ;;
        esac
        if [ -z "${PACKAGE}" ]; then
                echo "No Package specified with -update option"
                exit 1
        fi
		UPDATE=1
        ;;
	# -scan accepts two arguments, the package name and a path. If no path is specified exit
	-scan)
		PACKAGE=""
		case "$2" in
			"")
				;;
			-*)
				;;
			*)
				shift
				PACKAGE="$1"
				;;
		esac
		if [ -z "${PACKAGE}" ]; then
			echo "No Package specified with -scan option"
			exit 1
		fi
		SCANPATH=""
		case "$2" in
			"")
				;;
			-*)
				;;
			*)
				shift
				SCANPATH="$1"
				;;
		esac
		if [ -z "${SCANPATH}" ]; then
			echo "No path specified with -scan option"
			exit 1
		fi
		SCAN=1
		;;

    # Add the isolate option
    -isolate)
        STATE=""
        case "$2" in
            "")
                ;;
            -*)
                ;;
            *)
                shift
                STATE="$1"
                ;;
        esac
        if [ -z "${STATE}" ]; then
            echo "No state specified with -isolate option (use enable or disable)"
            exit 1
        fi
        ISOLATE=1
        ;;

	# Restart a target with -restart
	-restart)
	    PACKAGE=""
	    case "$2" in
	    	"")
	    			;;
	    	-*)
	    			;;
	    	*)
	    			shift
	    			PACKAGE="$1"
	    			;;
	    esac
	    if [ -z "${PACKAGE}" ]; then
	    	echo "No Package specified with -restart option"
	    	exit 1
	    fi
	    RESTART=1
	    ;;
	# Initialize
	-init)
        PACKAGE=""
        case "$2" in
                "")
                        ;;
                -*)
                        ;;
                *)
                        shift
                        PACKAGE="$1"
                        ;;
        esac
        if [ -z "${PACKAGE}" ]; then
                echo "No Package specified with -init option"
                exit 1
        fi
        INIT=1
        ;;

        *) # display Help
      		show_help
      		exit 1;;
   esac
   shift
done

####################
# Main
####################

#if [[ -z $ALL && -z $GROUP && -z $ID && -z $INIT && -z $UPDATE && -z $LIST ]]; then
#	show_help
#fi

# Build our target list
if [[ $ALL ]]; then
	# Get list of all online agents
	for list in $(/var/ossec/bin/agent_control -lc -s |grep Active |awk -F, '{print $1}');  do
		targets="$targets $list"
	done
elif [[ $GROUP ]]; then
	for list in $(/var/ossec/bin/agent_groups -g $GROUP -l | awk '/ID/ {print $2}') ; do	
		targets="$targets $list"
	done
elif [[ $LIST ]]; then
	echo
	fmt="%-5s%-48s%-10s%-16s\n"
	printf "$fmt" "ID" "Name" "Version" "Status" 
	for id in $(/var/ossec/bin/agent_control -l -s |egrep -v "^000|^na" |awk -F, '{print $1}'); do
		INFO=$(/var/ossec/bin/agent_control -i $id -s)
		NAME=$(echo $INFO |awk -F, '{print $2}')
		STATUS=$(echo $INFO |awk -F, '{print $4}')
		AGENT_VERSION=$(echo $INFO |awk -F, '{print $6}' |awk '{print $2}')
		if [[ $AGENT_VERSION == "" ]]; then
			AGENT_VERSION=null		
		fi
		printf "$fmt" $id $NAME $AGENT_VERSION "${STATUS}"
	done
	echo
	exit
elif [[ $LIST_GROUP ]]; then
	echo
	for group_name in $(/var/ossec/bin/agent_groups -l |egrep -v :| awk '{ print $1}' ); do
		/var/ossec/bin/agent_groups -g $group_name -l
	done
	echo
	exit
elif [[ $ID ]]; then
	targets=$ID
else
	echo 
	echo "  Error: No targets declared"
	echo
	exit 1
	
fi


if [[ $INSTALL ]] ; then
	for agent in $targets; do
		echo "Installing ${PACKAGE} on host: $agent"
		install_package $agent ${PACKAGE}
	done
elif [[ $ISOLATE ]]; then
    if [ -n "$ALL" ] || [ -n "$GROUP" ]; then
        echo "Error: -isolate can only be used with the -i flag"
        exit 1
    fi
    for agent in $targets; do
        echo "Setting isolation state to ${STATE} for agent: $agent"
        isolate_agent $agent ${STATE}
    done
elif [[ $UPDATE ]]; then
	for agent in $targets; do
		echo "Updating ${PACKAGE} on host: $agent"
		update_package $agent ${PACKAGE}
	done
elif [[ $RESTART ]]; then
	for agent in $targets; do
		echo "Restarting ${PACKAGE} on host: $agent"
		restart_package $agent ${PACKAGE}
	done
elif [[ $INIT ]]; then
	for agent in $targets; do
		echo "Initializing ${PACKAGE} on host: $agent"
		init_package $agent ${PACKAGE} 
	done
elif [[ $SCAN ]]; then
	for agent in $targets; do
		echo "Scanning module ${PACKAGE} on host: $agent with the target: (${SCANPATH})"
		scan_package $agent ${PACKAGE}
	done
else
	echo "  Error: no valid commands passed"
	show_help
fi
