From c5eff1d89eee1aa6203cd305acec6d49f4f03483 Mon Sep 17 00:00:00 2001
From: William Wernert <william.wernert@securityonionsolutions.com>
Date: Thu, 9 Jul 2020 14:47:55 -0400
Subject: [PATCH] [feat][WIP] Add option to run `so-allow -a <ip/cidr>` during
 setup

---
 salt/common/tools/sbin/so-allow | 26 +++++++++++++-------------
 setup/so-setup                  |  4 +++-
 setup/so-whiptail               | 26 +++++++++++++++++++++++---
 3 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/salt/common/tools/sbin/so-allow b/salt/common/tools/sbin/so-allow
index 71c8f8c9b..c6d3d6bf0 100755
--- a/salt/common/tools/sbin/so-allow
+++ b/salt/common/tools/sbin/so-allow
@@ -89,7 +89,7 @@ if [ "$SKIP" -eq 0 ]; then
     echo "[p] - Wazuh API - port 55000/tcp"
     echo "[r] - Wazuh registration service - 1515/tcp"
     echo ""
-    echo "Please enter your selection (a - analyst, b - beats, o - osquery, w - wazuh):"
+    echo "Please enter your selection:"
     read -r ROLE
     echo "Enter a single ip address or range to allow (example: 10.10.10.10 or 10.10.0.0/16):"
     read -r IP
@@ -127,16 +127,16 @@ salt-call state.apply firewall queue=True
 if grep -q -R "wazuh: 1" $local_salt_dir/pillar/*; then
     # If analyst, add to Wazuh AR whitelist
     if [ "$FULLROLE" == "analyst" ]; then
-          WAZUH_MGR_CFG="/opt/so/wazuh/etc/ossec.conf"
-          if ! grep -q "<white_list>$IP</white_list>" $WAZUH_MGR_CFG ; then
-                  DATE=$(date)
-                  sed -i 's/<\/ossec_config>//' $WAZUH_MGR_CFG
-                  sed -i '/^$/N;/^\n$/D' $WAZUH_MGR_CFG
-                  echo -e "<!--Address $IP added by /usr/sbin/so-allow on \"$DATE\"-->\n  <global>\n    <white_list>$IP</white_list>\n  </global>\n</ossec_config>" >> $WAZUH_MGR_CFG
-                  echo "Added whitelist entry for $IP in $WAZUH_MGR_CFG."
-		  echo
-                  echo "Restarting OSSEC Server..."
-                  /usr/sbin/so-wazuh-restart
-          fi
-     fi
+        WAZUH_MGR_CFG="/opt/so/wazuh/etc/ossec.conf"
+        if ! grep -q "<white_list>$IP</white_list>" $WAZUH_MGR_CFG ; then
+            DATE=$(date)
+            sed -i 's/<\/ossec_config>//' $WAZUH_MGR_CFG
+            sed -i '/^$/N;/^\n$/D' $WAZUH_MGR_CFG
+            echo -e "<!--Address $IP added by /usr/sbin/so-allow on \"$DATE\"-->\n  <global>\n    <white_list>$IP</white_list>\n  </global>\n</ossec_config>" >> $WAZUH_MGR_CFG
+            echo "Added whitelist entry for $IP in $WAZUH_MGR_CFG."
+        echo
+            echo "Restarting OSSEC Server..."
+            /usr/sbin/so-wazuh-restart
+        fi
+    fi
 fi
diff --git a/setup/so-setup b/setup/so-setup
index 3f6d42380..0f28f1244 100755
--- a/setup/so-setup
+++ b/setup/so-setup
@@ -331,6 +331,8 @@ else
 	FLEETNODEPASSWD1=$WEBPASSWD1
 fi
 
+if [[ $is_master ]]; then whiptail_so_allow; fi
+
 whiptail_make_changes
 
 if [[ -n "$TURBO" ]]; then
@@ -620,7 +622,7 @@ if [[ -n $ALLOW_ROLE && -n $ALLOW_CIDR ]]; then
 	IP=$ALLOW_CIDR so-allow -$ALLOW_ROLE >> $setup_log 2>&1
 fi
 
-if [[ $success != 0 || $SO_ERROR == 1 ]]; then whiptail_setup_failed
+if [[ $success != 0 || -n $SO_ERROR ]]; then whiptail_setup_failed
 else 
 	whiptail_setup_complete
 	if [[ $THEHIVE == 1 ]]; then check_hive_init; fi
diff --git a/setup/so-whiptail b/setup/so-whiptail
index e165ba351..0b34d3cb8 100755
--- a/setup/so-whiptail
+++ b/setup/so-whiptail
@@ -435,9 +435,7 @@ whiptail_homenet_sensor() {
 		whiptail_check_exitstatus $exitstatus
 
 		export HNSENSOR
-
 	fi
-
 }
 
 whiptail_install_type() {
@@ -1057,6 +1055,29 @@ whiptail_shard_count() {
 
 }
 
+whiptail_so_allow() {
+
+	[ -n "$TESTING" ] && return
+
+	whiptail --title "Security Onion Setup" \
+	--yesno "Do you want to run so-allow to allow access to the web tools?" \
+	8 75
+
+	local exitstatus=$?
+
+	if [[ $exitstatus == 0 ]]; then
+		ALLOW_CIDR=$(whiptail --title "Security Onion Setup" \
+		--inputbox "Enter a single ip address or range (in CIDR notation) to allow" \
+		10 75 125 3>&1 1>&2 2>&3)
+		local exitstatus=$?
+		
+		export ALLOW_ROLE='a'
+		export ALLOW_CIDR
+	fi
+
+	whiptail_check_exitstatus $exitstatus
+}
+
 whiptail_strelka_rules() {
 
         [ -n "$TESTING" ] && return
@@ -1073,7 +1094,6 @@ whiptail_strelka_rules() {
                 export STRELKARULES
 
         fi
-
 }
 
 whiptail_suricata_pins() {