From 6c4946f4e27152b5f65701d651f657a75aeda2e5 Mon Sep 17 00:00:00 2001
From: Jason Ertel <jason.ertel@securityonionsolutions.com>
Date: Thu, 28 May 2020 10:20:39 -0400
Subject: [PATCH] Provide option to skip reboot after setup completes

---
 setup/automation/pm_standalone_defaults | 1 +
 setup/so-functions                      | 3 +--
 setup/so-setup                          | 9 ++++++---
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/setup/automation/pm_standalone_defaults b/setup/automation/pm_standalone_defaults
index 74ba8323f..ae4554a3f 100644
--- a/setup/automation/pm_standalone_defaults
+++ b/setup/automation/pm_standalone_defaults
@@ -67,6 +67,7 @@ PLAYBOOK=1
 REDIRECTINFO=IP
 RULESETUP=ETOPEN
 # SHARDCOUNT=
+SKIP_REBOOT=1
 SOREMOTEPASS1=onionuser
 SOREMOTEPASS2=onionuser
 STRELKA=1
diff --git a/setup/so-functions b/setup/so-functions
index 6707e6841..52aee37d6 100755
--- a/setup/so-functions
+++ b/setup/so-functions
@@ -193,7 +193,7 @@ check_admin_pass() {
 	check_pass_match "$ADMINPASS1" "$ADMINPASS2" "APMATCH"
 }
 
-check_hive_init_then_reboot() {
+check_hive_init() {
 
 	wait_for_file /opt/so/state/thehive.txt 20 5
 	local return_val=$?
@@ -203,7 +203,6 @@ check_hive_init_then_reboot() {
 
 	docker stop so-thehive
 	docker rm so-thehive
-	shutdown -r now
 }
 
 check_network_manager_conf() {
diff --git a/setup/so-setup b/setup/so-setup
index 0d0022feb..24079bb94 100755
--- a/setup/so-setup
+++ b/setup/so-setup
@@ -37,6 +37,8 @@ while [[ $# -gt 0 ]]; do
 			export ALLOW_ROLE="${arg#*=}";;
 		"--allow-cidr="* )
 			export ALLOW_CIDR="${arg#*=}";;
+		"--skip-reboot" )
+			export SKIP_REBOOT=1
 		* )
 			if [[ "$arg" == "--"* ]]; then
 				echo "Invalid option"
@@ -563,11 +565,12 @@ if [[ "$success" = 0 ]]; then
 		so-allow -$ALLOW_ROLE >> $setup_log 2>&1
 	fi
 	if [[ $THEHIVE == 1 ]]; then
-		check_hive_init_then_reboot
-	else
-		shutdown -r now
+		check_hive_init
 	fi
 else
 	whiptail_setup_failed
+fi
+
+if [[ -z $SKIP_REBOOT ]]; then
 	shutdown -r now
 fi