From ae7c0a26be3d39e0f543ee7a05fc3278b31fd0b7 Mon Sep 17 00:00:00 2001
From: m0duspwnens <josh.patterson@securityonionsolutions.com>
Date: Wed, 6 Jan 2021 18:46:21 -0500
Subject: [PATCH] add a quiet mode to so-status for automation testing

---
 salt/common/tools/sbin/so-status | 107 +++++++++++++++++++++++--------
 1 file changed, 80 insertions(+), 27 deletions(-)

diff --git a/salt/common/tools/sbin/so-status b/salt/common/tools/sbin/so-status
index 8dd607bd6..dedd11d3e 100755
--- a/salt/common/tools/sbin/so-status
+++ b/salt/common/tools/sbin/so-status
@@ -20,7 +20,21 @@ if ! [ "$(id -u)" = 0 ]; then
    exit 1
 fi
 
+display_help() {
+cat <<HELP_USAGE
+
+    $0  [-h] [-q|--quiet]
+
+   -h             Show this message.
+   -q|--quiet     Suppress the output and only return a
+		  single status code for overall status
+		  0:Ok, 1:Error, 2:Starting/Pending
+HELP_USAGE
+}
+
 # Constants
+QUIET=false
+EXITCODE=0
 SYSTEM_START_TIME=$(date -d "$(</proc/uptime awk '{print $1}') seconds ago" +%s)
 # file populated by salt.lasthighstate state at end of successful highstate run
 LAST_HIGHSTATE_END=$([ -e "/opt/so/log/salt/lasthighstate" ] && date -r /opt/so/log/salt/lasthighstate +%s || echo 0)
@@ -81,7 +95,7 @@ compare_lists() {
 
 create_expected_container_list() {
 
-    mapfile -t expected_container_list < <(sort -u /opt/so/conf/so-status/so-status.conf | tr -d "#")    
+    mapfile -t expected_container_list < <(sort -u /opt/so/conf/so-status/so-status.conf | tr -d "#")
 
 }
 
@@ -111,43 +125,43 @@ populate_container_lists() {
 }
 
 parse_status() {
-    local container_state=${1}
-    local service_name=${2}
+    local service_name=${1}
+    local container_state=${2}
 
     for state in "${GOOD_STATUSES[@]}"; do
-        [[ $container_state = "$state" ]] && printf $SUCCESS_STRING && return 0
+        [[ $container_state = "$state" ]] && [[ $QUIET = "false" ]] && printf $SUCCESS_STRING && return 0 || [[ $container_state = "$state" ]] && return 0
     done
 
     for state in "${BAD_STATUSES[@]}"; do
-        [[ " ${DISABLED_CONTAINERS[@]} " =~ " ${service_name} " ]] && printf $DISABLED_STRING && return 0
+        [[ " ${DISABLED_CONTAINERS[@]} " =~ " ${service_name} " ]] && [[ $QUIET = "false" ]] && printf $DISABLED_STRING && return 0 || [[ " ${DISABLED_CONTAINERS[@]} " =~ " ${service_name} " ]] && return 0
     done
 
     # if a highstate has finished running since the system has started
     # then the containers should be running so let's check the status
     if [ $LAST_HIGHSTATE_END -ge $SYSTEM_START_TIME ]; then
 
-        [[ $container_state = "missing" ]] && printf $MISSING_STRING && return 1
+        [[ $container_state = "missing" ]] && [[ $QUIET = "false" ]] && printf $MISSING_STRING && return 1 || [[ $container_state = "missing" ]] && [[ "$EXITCODE" -lt 2 ]] && EXITCODE=1  && return 1
 
         for state in "${PENDING_STATUSES[@]}"; do
-            [[ $container_state = "$state" ]] && printf $PENDING_STRING && return 0
+            [[ $container_state = "$state" ]] && [[ $QUIET = "false" ]] && printf $PENDING_STRING && return 0
         done
 
         # This is technically not needed since the default is error state
         for state in "${BAD_STATUSES[@]}"; do
-            [[ $container_state = "$state" ]] &&  printf $ERROR_STRING && return 1
+            [[ $container_state = "$state" ]] && [[ $QUIET = "false" ]] && printf $ERROR_STRING && return 1 || [[ $container_state = "$state" ]] && [[ "$EXITCODE" -lt 2 ]] && EXITCODE=1  && return 1
         done
 
-        printf $ERROR_STRING && return 1
-    
+        [[ $QUIET = "false" ]] && printf $ERROR_STRING && return 1 || [[ "$EXITCODE" -lt 2 ]] && EXITCODE=1 && return 1
+
     # if a highstate has not run since system start time, but a highstate is currently running
     # then show that the containers are STARTING
     elif [[ "$HIGHSTATE_RUNNING" == 0 ]]; then
-        printf $STARTING_STRING && return 0
+        [[ $QUIET = "false" ]] && printf $STARTING_STRING && return 2 || EXITCODE=2 && return 2
 
     # if a highstate has not finished running since system startup and isn't currently running
     # then just show that the containers are WAIT_START; waiting to be started
     else
-        printf $WAIT_START_STRING && return 1
+        [[ $QUIET = "false" ]] && printf $WAIT_START_STRING && return 2 || EXITCODE=2 && return 2
 
     fi
 }
@@ -156,18 +170,22 @@ parse_status() {
 
 print_line() {
     local service_name=${1}
-    local service_state="$( parse_status ${2} ${1} )"
+    local service_state="$( parse_status ${1} ${2} )"
     local columns=$(tput cols)
     local state_color="\e[0m"
 
     local PADDING_CONSTANT=15
 
-    if [[ $service_state = "$ERROR_STRING" ]] || [[ $service_state = "$MISSING_STRING" ]] || [[ $service_state = "$WAIT_START_STRING" ]]; then
+    if [[ $service_state = "$ERROR_STRING" ]] || [[ $service_state = "$MISSING_STRING" ]]; then
         state_color="\e[1;31m"
+        if [[ "$EXITCODE" -eq 0 ]]; then
+            EXITCODE=1
+        fi
     elif [[ $service_state = "$SUCCESS_STRING" ]]; then
         state_color="\e[1;32m"
-    elif [[ $service_state = "$PENDING_STRING" ]] || [[ $service_state = "$DISABLED_STRING" ]] || [[ $service_state = "$STARTING_STRING" ]]; then
+    elif [[ $service_state = "$PENDING_STRING" ]] || [[ $service_state = "$DISABLED_STRING" ]] || [[ $service_state = "$STARTING_STRING" ]] || [[ $service_state = "$WAIT_START_STRING" ]]; then
         state_color="\e[1;33m"
+        EXITCODE=2
     fi
 
     printf "    $service_name "
@@ -181,7 +199,15 @@ print_line() {
 
 non_term_print_line() {
     local service_name=${1}
-    local service_state="$( parse_status ${2} ${1} )"
+    local service_state="$( parse_status ${1} ${2} )"
+
+    if [[ $service_state = "$ERROR_STRING" ]] || [[ $service_state = "$MISSING_STRING" ]]; then
+        if [[ "$EXITCODE" -eq 0 ]]; then
+            EXITCODE=1
+        fi
+    elif [[ $service_state = "$PENDING_STRING" ]] || [[ $service_state = "$DISABLED_STRING" ]] || [[ $service_state = "$STARTING_STRING" ]] || [[ $service_state = "$WAIT_START_STRING" ]]; then
+        EXITCODE=2
+    fi
 
     printf "    $service_name "
     for i in $(seq 0 $(( 35 - ${#service_name} - ${#service_state} ))); do
@@ -218,37 +244,64 @@ main() {
       done
 
       printf "\n"
-    
+
     # else if running from a terminal
     else
 
-      local focus_color="\e[1;34m"
-      printf "\n"
-      printf "${focus_color}%b\e[0m" "Checking Docker status\n\n"
+      if [ "$QUIET" = true ]; then
+          print_or_parse="parse_status"
+      else
+          print_or_parse="print_line"
+
+          local focus_color="\e[1;34m"
+          printf "\n"
+          printf "${focus_color}%b\e[0m" "Checking Docker status\n\n"
+      fi
 
       systemctl is-active --quiet docker
       if [[ $? = 0 ]]; then
-          print_line "Docker" "running"
+          ${print_or_parse} "Docker" "running"
       else
-          print_line "Docker" "exited"
+          ${print_or_parse} "Docker" "exited"
       fi
 
       populate_container_lists
 
-      printf "\n"
-      printf "${focus_color}%b\e[0m" "Checking container statuses\n\n"
+      if [ "$QUIET" = false ]; then
+          printf "\n"
+          printf "${focus_color}%b\e[0m" "Checking container statuses\n\n"
+      fi
 
       local num_containers=${#container_name_list[@]}
 
       for i in $(seq 0 $(($num_containers - 1 ))); do
-          print_line ${container_name_list[$i]} ${container_state_list[$i]}
+          ${print_or_parse} ${container_name_list[$i]} ${container_state_list[$i]}
       done
 
-      printf "\n"
+      if [ "$QUIET" = false ]; then
+          printf "\n"
+      fi
     fi
 }
 
 # {% endraw %}
 
+while getopts ':hq' OPTION; do
+  case "$OPTION" in
+    h)
+      display_help
+      exit 0
+      ;;
+    q)
+      QUIET=true
+      ;;
+    \?)
+      display_help
+      exit 0
+      ;;
+  esac
+done
 
-main
\ No newline at end of file
+main
+
+exit $EXITCODE
\ No newline at end of file