From a1294ae508af9802523b9cd9568b31428b5eff41 Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Fri, 30 Oct 2015 10:43:55 -0700
Subject: [PATCH] ability to run commands on multiple hosts

---
 DOCS.md  | 21 ++++++++++++++++++++-
 main.go  | 21 +++++++++++++--------
 types.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 9 deletions(-)
 create mode 100644 types.go

diff --git a/DOCS.md b/DOCS.md
index 5011879..6efe438 100644
--- a/DOCS.md
+++ b/DOCS.md
@@ -5,7 +5,7 @@ Use the SSH plugin to execute commands on a remote server. The following paramet
 * `user` - user to log in as on the remote machine
 * `commands` - list of commands to execute
 
-The following is a sample SSH configuration in your .drone.yml file:
+Example configuration in your .drone.yml file:
 
 ```yaml
 deploy:
@@ -18,4 +18,23 @@ deploy:
       - echo world
 ```
 
+Example multi-host configuration in your .drone.yml file:
+
+```yaml
+deploy:
+  ssh:
+    host:
+     - foo.com
+     - bar.com
+    user: root
+    port: 22
+    commands:
+      - echo hello
+      - echo world
+```
+
+In the above example Drone executes the commands on multiple hosts sequentially. If the commands fail on a single host this plugin exits immediatly, and will not run your commands on the remaining hosts in the list.
+
+## Keys
+
 The plugin authenticates to your server using a per-repository SSH key generated by Drone. You can find the public key in your repository settings in Drone. You will need to copy / paste this key into your `~/.ssh/authorized_keys` file on your remote machine.
\ No newline at end of file
diff --git a/main.go b/main.go
index 8604952..ff0d4fb 100644
--- a/main.go
+++ b/main.go
@@ -17,7 +17,7 @@ type Params struct {
 	Commands []string `json:"commands"`
 	Login    string   `json:"user"`
 	Port     int      `json:"port"`
-	Host     string   `json:"host"`
+	Host     StrSlice `json:"host"`
 }
 
 func main() {
@@ -27,13 +27,15 @@ func main() {
 	plugin.Param("vargs", &v)
 	plugin.MustParse()
 
-	err := run(w.Keys, v)
-	if err != nil {
-		os.Exit(1)
+	for _, host := range v.Host.Slice() {
+		err := run(w.Keys, v, host)
+		if err != nil {
+			os.Exit(1)
+		}
 	}
 }
 
-func run(keys *plugin.Keypair, params *Params) error {
+func run(keys *plugin.Keypair, params *Params, host string) error {
 
 	// if no username is provided assume root
 	if len(params.Login) == 0 {
@@ -46,11 +48,14 @@ func run(keys *plugin.Keypair, params *Params) error {
 	}
 
 	// join the host and port if necessary
-	host := net.JoinHostPort(
-		params.Host,
+	addr := net.JoinHostPort(
+		host,
 		strconv.Itoa(params.Port),
 	)
 
+	// trace command used for debugging in the build logs
+	fmt.Printf("$ ssh %s@%s -p %d\n", params.Login, addr, params.Port)
+
 	signer, err := ssh.ParsePrivateKey([]byte(keys.Private))
 	if err != nil {
 		return fmt.Errorf("Error parsing private key. %s.", err)
@@ -61,7 +66,7 @@ func run(keys *plugin.Keypair, params *Params) error {
 		Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)},
 	}
 
-	client, err := ssh.Dial("tcp", host, config)
+	client, err := ssh.Dial("tcp", addr, config)
 	if err != nil {
 		return fmt.Errorf("Error dialing server. %s.", err)
 	}
diff --git a/types.go b/types.go
new file mode 100644
index 0000000..7831e0c
--- /dev/null
+++ b/types.go
@@ -0,0 +1,45 @@
+package main
+
+import "encoding/json"
+
+// StrSlice representes a string or an array of strings.
+// We need to override the json decoder to accept both options.
+type StrSlice struct {
+	parts []string
+}
+
+// UnmarshalJSON decodes the byte slice whether it's a string or an array of strings.
+// This method is needed to implement json.Unmarshaler.
+func (e *StrSlice) UnmarshalJSON(b []byte) error {
+	if len(b) == 0 {
+		return nil
+	}
+
+	p := make([]string, 0, 1)
+	if err := json.Unmarshal(b, &p); err != nil {
+		var s string
+		if err := json.Unmarshal(b, &s); err != nil {
+			return err
+		}
+		p = append(p, s)
+	}
+
+	e.parts = p
+	return nil
+}
+
+// Len returns the number of parts of the StrSlice.
+func (e *StrSlice) Len() int {
+	if e == nil {
+		return 0
+	}
+	return len(e.parts)
+}
+
+// Slice gets the parts of the StrSlice as a Slice of string.
+func (e *StrSlice) Slice() []string {
+	if e == nil {
+		return nil
+	}
+	return e.parts
+}