From 97d926b432d1aa53f412830396f85504b495307f Mon Sep 17 00:00:00 2001
From: Leo Antunes <leo@costela.net>
Date: Wed, 22 Aug 2018 00:19:42 +0200
Subject: [PATCH] add cleanup on exit

---
 main.go | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/main.go b/main.go
index a81d366..89b6c72 100644
--- a/main.go
+++ b/main.go
@@ -18,7 +18,10 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 package main
 
 import (
+	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 
 	docker "docker.io/go-docker"
 
@@ -48,6 +51,13 @@ func main() {
 
 	log.SetLevel(logLevelMap[strings.ToLower(config.LogLevel)])
 
+	quitSig := make(chan os.Signal)
+	signal.Notify(quitSig, syscall.SIGTERM, syscall.SIGINT)
+	go func() {
+		<-quitSig
+		cleanup()
+	}()
+
 	client, err := docker.NewEnvClient()
 	if err != nil {
 		log.Fatalf("error starting docker client: %s", err)
@@ -60,3 +70,9 @@ func main() {
 		syncAndListenForEvents(client)
 	}
 }
+
+func cleanup() {
+	log.Info("cleaning up hosts file")
+	writeToEtcHosts(ipsToNamesMap{})
+	os.Exit(0)
+}