From ea1ab75072bfc7fb74d79886b3d1b9582c6fad5d Mon Sep 17 00:00:00 2001
From: Jason Ertel <jason.ertel@securityonionsolutions.com>
Date: Wed, 13 Jan 2021 12:42:41 -0500
Subject: [PATCH] Refactored so-common node type checks for improved
 readability; Updated so-tcpreplay to support distributed grids

---
 salt/common/tools/sbin/so-common    | 32 ++++++++++++++++++++++-------
 salt/common/tools/sbin/so-tcpreplay | 28 ++++++++++++++++++-------
 setup/so-setup                      |  2 +-
 3 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/salt/common/tools/sbin/so-common b/salt/common/tools/sbin/so-common
index 0c18c4482..881be83ca 100755
--- a/salt/common/tools/sbin/so-common
+++ b/salt/common/tools/sbin/so-common
@@ -111,9 +111,7 @@ set_version() {
 }
 
 require_manager() {
-	# Check to see if this is a manager
-	MANAGERCHECK=$(cat /etc/salt/grains | grep role | awk '{print $2}')
-	if [ $MANAGERCHECK == 'so-eval' ] || [ $MANAGERCHECK == 'so-manager' ] || [ $MANAGERCHECK == 'so-managersearch' ] || [ $MANAGERCHECK == 'so-standalone' ] || [ $MANAGERCHECK == 'so-helix' ] || [ $MANAGERCHECK == 'so-import' ]; then
+	if is_manager; then
 		echo "This is a manager, We can proceed."
 	else
 		echo "Please run this command on the manager; the manager controls the grid."
@@ -121,12 +119,32 @@ require_manager() {
 	fi
 }
 
+is_manager() {
+	# Check to see if this is a manager node
+	role=$(lookup_role)
+	is_single_node_grid && return 0
+	[ $role == 'manager' ] && return 0
+	[ $role == 'managersearch' ] && return 0
+	[ $role == 'helix' ] && return 0
+	return 1
+}
+
+is_sensor() {
+	# Check to see if this is a sensor (forward) node
+	role=$(lookup_role)
+	is_single_node_grid && return 0
+	[ $role == 'sensor' ] && return 0
+	[ $role == 'heavynode' ] && return 0
+	[ $role == 'helix' ] && return 0
+	return 1
+}
+
 is_single_node_grid() {
 	role=$(lookup_role)
-	if [ "$role" != "eval" ] && [ "$role" != "standalone" ] && [ "$role" != "import" ]; then
-		return 1
-	fi
-	return 0
+	[ $role == 'eval' ] && return 0
+	[ $role == 'standalone' ] && return 0
+	[ $role == 'import' ] && return 0
+	return 1
 }
 
 fail() {
diff --git a/salt/common/tools/sbin/so-tcpreplay b/salt/common/tools/sbin/so-tcpreplay
index fa992bdd8..8b81c32cf 100755
--- a/salt/common/tools/sbin/so-tcpreplay
+++ b/salt/common/tools/sbin/so-tcpreplay
@@ -47,13 +47,27 @@ if ! docker ps | grep -q so-tcpreplay; then
   echo "Replay functionality not enabled; attempting to enable now (may require Internet access)..."
   echo
 
-  TRUSTED_CONTAINERS=("so-tcpreplay")
-  mkdir -p /opt/so/log/tcpreplay
-  update_docker_containers "tcpreplay" "" "" "/opt/so/log/tcpreplay/init.log"
-  so-tcpreplay-start || fail "Unable to initialize tcpreplay"
+  if is_manager; then
+	  TRUSTED_CONTAINERS=("so-tcpreplay")
+	  mkdir -p /opt/so/log/tcpreplay
+	  update_docker_containers "tcpreplay" "" "" "/opt/so/log/tcpreplay/init.log"
+	elif is_sensor; then
+		if ! is_manager; then
+			echo "Attempting to start replay container. If this fails then you may need to run this command on the manager first."
+		fi
+  	so-tcpreplay-start || fail "Unable to initialize tcpreplay"
+  else
+  	echo "Unable to enable replay functionality on this node type."
+  fi
 fi
 
-echo "Replaying PCAP(s) at ${REPLAYSPEED} Mbps on interface ${REPLAYIFACE}..."
-docker exec so-tcpreplay /usr/bin/bash -c "/usr/local/bin/tcpreplay -i ${REPLAYIFACE} -M${REPLAYSPEED} $@"
+if is_sensor; then
+	echo "Replaying PCAP(s) at ${REPLAYSPEED} Mbps on interface ${REPLAYIFACE}..."
+	docker exec so-tcpreplay /usr/bin/bash -c "/usr/local/bin/tcpreplay -i ${REPLAYIFACE} -M${REPLAYSPEED} $@"
 
-echo "Replay completed. Warnings shown above are typically expected."
+	echo "Replay completed. Warnings shown above are typically expected."
+elif is_manager; then
+	echo "The sensor nodes in this grid can now replay traffic."
+else
+	echo "Unable to replay traffic since this node is not a sensor node."
+fi
diff --git a/setup/so-setup b/setup/so-setup
index b4b0fd6ed..2cee0dc6a 100755
--- a/setup/so-setup
+++ b/setup/so-setup
@@ -852,7 +852,7 @@ if [[ -n $SO_ERROR ]]; then
 else 
 	echo "Successfully completed setup! Continuing with post-installation steps" >> $setup_log 2>&1
 	{
-		[[ -n "$TESTING" && $is_sensor ]] && logCmd so-test
+		[[ -n "$TESTING" ]] && logCmd so-test
 
 		export percentage=95 # set to last percentage used in previous subshell
 		if [[ -n $ALLOW_ROLE && -n $ALLOW_CIDR ]]; then