From 626365d72ed48917f05a5215077fd986fdb5946e Mon Sep 17 00:00:00 2001
From: 9e3fce0a <9e3fce0a@users.noreply.github.com>
Date: Wed, 2 Feb 2022 09:38:48 +0000
Subject: [PATCH] #198 Option to delay shutdown if there are active SSH
 sessions

---
 cmd/deploy.go                         | 51 +++++++++++++++------------
 internal/templates/templates.go       |  2 ++
 templates/generated_vars_and_funcs.sh |  7 +++-
 3 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/cmd/deploy.go b/cmd/deploy.go
index ef4e039..1b509c8 100644
--- a/cmd/deploy.go
+++ b/cmd/deploy.go
@@ -31,6 +31,7 @@ var (
 	coreConfigRepo, customConfigRepo                                          string
 	coreConfigRepoBranch, customConfigRepoBranch                              string
 	outputDir                                                                 string
+	instanceDebugDelayTermination                                             bool
 	// TODO: apv workaround - remove once alternative is built
 	apvRemote, apvBranch, apvRevision                                         string
 )
@@ -113,6 +114,9 @@ func deployInit() {
 
 	flags.BoolVar(&dryRun, "dry-run", false, "only generate the output files, but do not deploy with terraform.")
 
+	flags.BoolVar(&instanceDebugDelayTermination, "instance-debug-delay-termination", false, "delay instance shutdown/termination if there are active SSH sessions")
+	_ = viper.BindPFlag("instance-debug-delay-termination", flags.Lookup("instance-debug-delay-termination"))
+
 	// TODO: apv workaround - remove once alternative is built
 	flags.StringVar(&apvRemote, "apv-remote", "", "remote that contains android-prepare-vendor repo (e.g. https://github.com/example/)")
 	_ = viper.BindPFlag("apv-remote", flags.Lookup("apv-remote"))
@@ -228,29 +232,30 @@ var deployCmd = &cobra.Command{
 		log.Infof("all generated files will be placed in %v", configuredOutputDir)
 
 		templateConfig := &templates.Config{
-			Version:                stackVersion,
-			Name:                   viper.GetString("name"),
-			Region:                 viper.GetString("region"),
-			Device:                 viper.GetString("device"),
-			DeviceDetails:          supportedDevices.GetDeviceDetails(viper.GetString("device")),
-			Email:                  viper.GetString("email"),
-			InstanceType:           viper.GetString("instance-type"),
-			InstanceRegions:        viper.GetString("instance-regions"),
-			SkipPrice:              viper.GetString("skip-price"),
-			MaxPrice:               viper.GetString("max-price"),
-			SSHKey:                 viper.GetString("ssh-key"),
-			Schedule:               viper.GetString("schedule"),
-			ChromiumBuildDisabled:  viper.GetBool("chromium-build-disabled"),
-			ChromiumVersion:        viper.GetString("chromium-version"),
-			CoreConfigRepo:         viper.GetString("core-config-repo"),
-			CoreConfigRepoBranch:   viper.GetString("core-config-repo-branch"),
-			CustomConfigRepo:       viper.GetString("custom-config-repo"),
-			CustomConfigRepoBranch: viper.GetString("custom-config-repo-branch"),
-			ReleasesURL:            viper.GetString("releases-url"),
-			Cloud:                  viper.GetString("cloud"),
-			ApvRemote:              viper.GetString("apv-remote"),
-			ApvBranch:              viper.GetString("apv-branch"),
-			ApvRevision:            viper.GetString("apv-revision"),
+			Version:                       stackVersion,
+			Name:                          viper.GetString("name"),
+			Region:                        viper.GetString("region"),
+			Device:                        viper.GetString("device"),
+			DeviceDetails:                 supportedDevices.GetDeviceDetails(viper.GetString("device")),
+			Email:                         viper.GetString("email"),
+			InstanceType:                  viper.GetString("instance-type"),
+			InstanceRegions:               viper.GetString("instance-regions"),
+			SkipPrice:                     viper.GetString("skip-price"),
+			MaxPrice:                      viper.GetString("max-price"),
+			SSHKey:                        viper.GetString("ssh-key"),
+			Schedule:                      viper.GetString("schedule"),
+			ChromiumBuildDisabled:         viper.GetBool("chromium-build-disabled"),
+			ChromiumVersion:               viper.GetString("chromium-version"),
+			CoreConfigRepo:                viper.GetString("core-config-repo"),
+			CoreConfigRepoBranch:          viper.GetString("core-config-repo-branch"),
+			CustomConfigRepo:              viper.GetString("custom-config-repo"),
+			CustomConfigRepoBranch:        viper.GetString("custom-config-repo-branch"),
+			ReleasesURL:                   viper.GetString("releases-url"),
+			Cloud:                         viper.GetString("cloud"),
+			InstanceDebugDelayTermination: viper.GetBool("instance-debug-delay-termination"),
+			ApvRemote:                     viper.GetString("apv-remote"),
+			ApvBranch:                     viper.GetString("apv-branch"),
+			ApvRevision:                   viper.GetString("apv-revision"),
 		}
 
 		templateRenderer, err := templates.New(templateConfig, templatesFiles, configuredOutputDir)
diff --git a/internal/templates/templates.go b/internal/templates/templates.go
index 02ff7d4..6cb53da 100644
--- a/internal/templates/templates.go
+++ b/internal/templates/templates.go
@@ -93,6 +93,8 @@ type Config struct {
 	ReleasesURL string
 	// Cloud specifies which cloud to build on (only aws supported right now)
 	Cloud string
+	// Delay instance shutdown/termination if there are active SSH sessions
+	InstanceDebugDelayTermination bool
 	// TODO: apv workaround - remove once alternative is built
 	// ApvRemote is the git remote that contains an android-prepare-vendor repo
 	ApvRemote string
diff --git a/templates/generated_vars_and_funcs.sh b/templates/generated_vars_and_funcs.sh
index 2153deb..9b29bd5 100644
--- a/templates/generated_vars_and_funcs.sh
+++ b/templates/generated_vars_and_funcs.sh
@@ -79,6 +79,11 @@ cleanup() {
   if [ $rv -ne 0 ]; then
     notify "RattlesnakeOS Build FAILED" 1
   fi
+  <% if .InstanceDebugDelayTermination -%>
+  while pgrep -u ubuntu sshd > /dev/null 2>&1; do
+    sleep 300
+  done
+  <%- end %>
   sudo shutdown -h now
   <%- else %>
   echo "todo"
@@ -142,4 +147,4 @@ delete_build_artifact() {
   <%- else %>
   echo "todo"
   <%- end %>
-}
\ No newline at end of file
+}