Handle flush file and INI together

main.go

@@ -55,42 +55,70 @@ type IAMPolicy struct {
 	Statement []Statement `json:"Statement"`
 }
 
-func setCSMConfig() {
-	cfgfile, err := homedir.Expand("~/.aws/config")
-	if err != nil {
-		return
-	}
+func setCSMConfigAndFileFlush() {
+	// set ini
+	if *setiniFlag {
+		cfgfile, err := homedir.Expand("~/.aws/config")
+		if err != nil {
+			return
+		}
 
-	cfg, err := ini.Load(cfgfile)
-	if err != nil {
-		return
-	}
+		cfg, err := ini.Load(cfgfile)
+		if err != nil {
+			return
+		}
 
-	if *profileFlag == "default" {
-		cfg.Section("default").Key("csm_enabled").SetValue("true")
-	} else {
-		cfg.Section(fmt.Sprintf("profile %s", *profileFlag)).Key("csm_enabled").SetValue("true")
-	}
+		if *profileFlag == "default" {
+			cfg.Section("default").Key("csm_enabled").SetValue("true")
+		} else {
+			cfg.Section(fmt.Sprintf("profile %s", *profileFlag)).Key("csm_enabled").SetValue("true")
+		}
 
-	cfg.SaveTo(cfgfile)
+		cfg.SaveTo(cfgfile)
+	}
 
-	// listen for exit and cleanup
+	// listen for exit, cleanup and flush
 	sigc := make(chan os.Signal, 1)
 	signal.Notify(sigc,
 		syscall.SIGHUP,
 		syscall.SIGINT,
 		syscall.SIGTERM,
 		syscall.SIGQUIT)
 	go func() {
-		<-sigc
-		if *profileFlag == "default" {
-			cfg.Section("default").DeleteKey("csm_enabled")
-		} else {
-			cfg.Section(fmt.Sprintf("profile %s", *profileFlag)).DeleteKey("csm_enabled")
-		}
-		cfg.SaveTo(cfgfile)
+		for s := range sigc {
+			// flush to file
+			if *outputFileFlag != "" {
+				err := ioutil.WriteFile(*outputFileFlag, getPolicyDocument(), 0644)
+				if err != nil {
+					log.Fatalf("Error writing policy to %s", *outputFileFlag)
+				}
+			}
+
+			if s == syscall.SIGINT || s == syscall.SIGTERM || s == syscall.SIGQUIT {
+				// revert ini
+				cfgfile, err := homedir.Expand("~/.aws/config") // need to redeclare
+				if err != nil {
+					os.Exit(1)
+				}
+
+				cfg, err := ini.Load(cfgfile)
+				if err != nil {
+					os.Exit(1)
+				}
+
+				if *setiniFlag {
+					if *profileFlag == "default" {
+						cfg.Section("default").DeleteKey("csm_enabled")
+					} else {
+						cfg.Section(fmt.Sprintf("profile %s", *profileFlag)).DeleteKey("csm_enabled")
+					}
+					cfg.SaveTo(cfgfile)
+				}
 
-		os.Exit(0)
+				// exit
+				os.Exit(0)
+			}
+		}
 	}()
 }
 
@@ -321,38 +349,11 @@ func setTerminalRefresh() {
 	}()
 }
 
-func setFileFlush() {
-	if *outputFileFlag == "" {
-		log.Fatal("No file specified")
-	}
-
-	sigc := make(chan os.Signal, 1)
-	signal.Notify(sigc,
-		syscall.SIGHUP,
-		syscall.SIGINT,
-		syscall.SIGTERM)
-	go func() {
-		for s := range sigc {
-			err := ioutil.WriteFile(*outputFileFlag, getPolicyDocument(), 0644)
-			if err != nil {
-				log.Fatalf("Error writing policy to %s", *outputFileFlag)
-			}
-			if s == syscall.SIGINT || s == syscall.SIGTERM {
-				os.Exit(0)
-			}
-		}
-	}()
-}
-
 func main() {
 	flag.Parse()
 
-	if *setiniFlag {
-		setCSMConfig()
-	}
-	if *outputFileFlag != "" {
-		setFileFlush()
-	}
+	setCSMConfigAndFileFlush()
+
 	if *terminalRefreshSecsFlag != 0 {
 		setTerminalRefresh()
 	}