From fdab66761481c27ea66451b3bf7203fc3fb1dea5 Mon Sep 17 00:00:00 2001
From: ChrisMcKenzie <chris@kbsurfer.com>
Date: Wed, 4 Nov 2015 14:21:01 -0800
Subject: [PATCH] added support for preCommand; added error rollback and
 cleanup support

---
 commands/agent.go     |  5 ++++-
 installer/install.go  | 19 ++++++++++++++++++-
 installer/tar.go      |  2 ++
 package.sh            | 12 ++++++++----
 service/dispatcher.go | 10 +++++++++-
 service/service.go    |  1 +
 6 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/commands/agent.go b/commands/agent.go
index a3ac786..3028a24 100644
--- a/commands/agent.go
+++ b/commands/agent.go
@@ -34,7 +34,10 @@ func agent(c *cobra.Command, args []string) {
 	shutdownCh := make(chan struct{})
 
 	for _, s := range services {
-		service.NewDispatcher(s, t, shutdownCh)
+		_, err := service.NewDispatcher(s, t, shutdownCh)
+		if err != nil {
+			log.Fatal(err)
+		}
 	}
 
 	sigs := make(chan os.Signal, 1)
diff --git a/installer/install.go b/installer/install.go
index f21c39b..d33c07e 100644
--- a/installer/install.go
+++ b/installer/install.go
@@ -1,9 +1,26 @@
 package installer
 
-import "io"
+import (
+	"io"
+	"os"
+	"path/filepath"
+)
 
 // Installer is an interface that allows different methods of writing
 // the given io.Reader to disk.
 type Installer interface {
 	Install(dest string, r io.Reader) (int, error)
 }
+
+func moveOld(dest string) error {
+	return os.Rename(dest, filepath.Join(dest, ".old"))
+}
+
+func cleanup(dest string, err error) error {
+	oldPath := filepath.Join(dest, ".old")
+	if err != nil {
+		return os.Rename(oldPath, dest)
+	}
+
+	return os.RemoveAll(oldPath)
+}
diff --git a/installer/tar.go b/installer/tar.go
index 8cf5f62..47b31e3 100644
--- a/installer/tar.go
+++ b/installer/tar.go
@@ -18,6 +18,7 @@ var ErrNilReader = errors.New("Install: must have a non-nil Reader")
 type TarInstaller struct{}
 
 func (i TarInstaller) Install(dest string, fr io.Reader) (count int, err error) {
+	moveOld(dest)
 	if fr == nil {
 		return count, ErrNilReader
 	}
@@ -65,5 +66,6 @@ func (i TarInstaller) Install(dest string, fr io.Reader) (count int, err error)
 		count++
 	}
 
+	defer cleanup(dest, err)
 	return
 }
diff --git a/package.sh b/package.sh
index af3640e..b2c5941 100755
--- a/package.sh
+++ b/package.sh
@@ -1,9 +1,13 @@
 mkdir -p packaging/output
 
+arch=$1
+os=$2
+version=$3
+
 set -x
-GOARCH=$1 GOOS=$2 go build -o packaging/root/usr/local/bin/dropship main.go
+GOARCH=$arch GOOS=$os go build -o packaging/root/usr/local/bin/dropship main.go
 
-fpm -s dir -t deb -n dropship -v "$1" -p packaging/output/dropship.$2-$1.deb \
+fpm -s dir -t deb -n dropship -v "$version" -p packaging/output/dropship-$version.$os-$arch.deb \
   --deb-priority optional --category admin \
   --force \
   --after-install packaging/scripts/postinstall.deb \
@@ -13,8 +17,8 @@ fpm -s dir -t deb -n dropship -v "$1" -p packaging/output/dropship.$2-$1.deb \
   --url https://github.com/chrismckenzie/dropship \
   --description "Dropship automatically keeps you software up to date" \
   -m "Chris McKenzie <chris@chrismckenzie.io>" \
-  -a $1 \
+  -a $arch \
   --config-files /etc/dropship.d/dropship.hcl \
   packaging/root/=/
 
-cp packaging/output/dropship.$2-$1.deb .
+cp packaging/output/dropship-$version.$os-$arch.deb .
diff --git a/service/dispatcher.go b/service/dispatcher.go
index cd3e7a9..fc4baad 100644
--- a/service/dispatcher.go
+++ b/service/dispatcher.go
@@ -88,6 +88,14 @@ func (w *Dispatcher) Work() {
 			defer l.Release()
 		}
 
+		if w.config.PreCommand != "" {
+			res, err := executeCommand(w.config.PreCommand)
+			if err != nil {
+				log.Printf("[ERR]: Unable to execute preComment. %v", err)
+			}
+			log.Printf("[INF]: preCommand executed successfully. %v", res)
+		}
+
 		log.Printf("[INF]: Installing update for %s...", w.config.Name)
 		fr, meta, err := u.Download(opts)
 		if err != nil {
@@ -109,7 +117,7 @@ func (w *Dispatcher) Work() {
 		if w.config.PostCommand != "" {
 			res, err := executeCommand(w.config.PostCommand)
 			if err != nil {
-				log.Printf("[ERR]: Unable to execute postComment. %v", err)
+				log.Printf("[ERR]: Unable to execute postCommand. %v", err)
 			}
 			log.Printf("[INF]: postCommand executed successfully. %v", res)
 		}
diff --git a/service/service.go b/service/service.go
index 33e4f31..f476da0 100644
--- a/service/service.go
+++ b/service/service.go
@@ -11,6 +11,7 @@ type Config struct {
 	Name          string   `hcl:",key"`
 	CheckInterval string   `hcl:"checkInterval"`
 	PostCommand   string   `hcl:postCommand`
+	PreCommand    string   `hcl:preCommand`
 	Sequential    bool     `hcl:"sequentialUpdates"`
 	Artifact      Artifact `hcl:"artifact,expand"`
 }