Refactor command bits into its own file

Signed-off-by: Blaine Gardner <[email protected]>
BlaineEXE · Sep 16, 2018 · 9eddf15 · 9eddf15
1 parent 2fb0bd3
commit 9eddf15
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 94 deletions.
diff --git a/cmd/octopus/command.go b/cmd/octopus/command.go
@@ -0,0 +1,89 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"strings"
+
+	"golang.org/x/crypto/ssh"
+)
+
+func newCommandConfig(identityFile string) (*ssh.ClientConfig, error) {
+	key, err := ioutil.ReadFile(identityFile)
+	if err != nil {
+		return nil, fmt.Errorf("unable to read private key: %v", err)
+	}
+
+	signer, err := ssh.ParsePrivateKey(key)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse private key: %v", err)
+	}
+
+	config := &ssh.ClientConfig{
+		User: "root",
+		Auth: []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	return config, nil
+}
+
+func runCommand(host, command string, config *ssh.ClientConfig, out chan<- tentacle) {
+	runFailedText := "run command failed"
+	t := tentacle{
+		// fallback hostname includes the raw host (e.g., IP) for some ability to identify the host
+		hostname: fmt.Sprintf("%s: could not get hostname", host),
+		stdout:   new(bytes.Buffer),
+		err:      fmt.Errorf("%s: unable to get more detail", runFailedText), // fallback error
+	}
+	defer func() { out <- t }()
+
+	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:22", host), config)
+	if err != nil {
+		t.err = fmt.Errorf("%v: %v", t.err, err)
+		return
+	}
+
+	// Get the host's hostname for easier identification
+	hch := make(chan string)
+	go func() {
+		b := new(bytes.Buffer)
+		err := doRunCommand("hostname", client, b)
+		if err != nil {
+			// command run wasn't a failure if there were problems getting the hostname
+			hch <- t.hostname // just output whatever hostname was set as the fallback on err
+		} else {
+			hch <- strings.TrimRight(b.String(), "\n")
+		}
+		close(hch)
+	}()
+
+	t.err = doRunCommand(command, client, t.stdout)
+	if t.err != nil {
+		t.err = fmt.Errorf("%s: %v", runFailedText, t.err)
+	}
+
+	t.hostname = <-hch
+	return
+}
+
+func doRunCommand(command string, client *ssh.Client, stdoutBuffer *bytes.Buffer) error {
+	session, err := client.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	stderr := new(bytes.Buffer)
+	session.Stdout = stdoutBuffer
+	session.Stderr = stderr
+
+	err = session.Run(command)
+	if err != nil {
+		return fmt.Errorf("%v:\n\n%s", err, strings.TrimRight(stderr.String(), "\n"))
+	}
+	return err
+}
diff --git a/cmd/octopus/octopus.go b/cmd/octopus/octopus.go
@@ -6,16 +6,9 @@ import (
 	"bytes"
 	"flag"
 	"fmt"
-	"io/ioutil"
 	"log"
 	"os"
 	"strings"
-
-	"golang.org/x/crypto/ssh"
-)
-
-const (
-	noHostNameText = "! could not get hostname !"
 )
 
 // Each of octopus's tentacles is a remote connection to a host executing the command
@@ -38,22 +31,9 @@ func main() {
 		os.Exit(1)
 	}
 
-	key, err := ioutil.ReadFile(*identityFile)
+	config, err := newCommandConfig(*identityFile)
 	if err != nil {
-		log.Fatalf("unable to read private key: %v", err)
-	}
-
-	signer, err := ssh.ParsePrivateKey(key)
-	if err != nil {
-		log.Fatalf("unable to parse private key: %v", err)
-	}
-
-	config := &ssh.ClientConfig{
-		User: "root",
-		Auth: []ssh.AuthMethod{
-			ssh.PublicKeys(signer),
-		},
-		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		log.Fatalf("could not generate command config: %v", err)
 	}
 
 	hosts := []string{"10.86.1.87", "10.86.1.103"}
@@ -76,81 +56,9 @@ func main() {
 	os.Exit(numErrors)
 }
 
-func runCommand(host, command string, config *ssh.ClientConfig, out chan<- tentacle) {
-	t := tentacle{
-		host:     host,
-		hostname: "",
-		err:      fmt.Errorf("run command failed"),
-	}
-	defer func() { out <- t }()
-
-	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:22", t.host), config)
-	if err != nil {
-		t.err = fmt.Errorf("%v: %v", t.err, err)
-		return
-		// log.Fatal("Failed to dial: ", err)
-	}
-
-	hn := make(chan tentacle)
-	go getHostname(client, hn)
-
-	session, err := client.NewSession()
-	if err != nil {
-		t.err = fmt.Errorf("%v: %v", t.err, err)
-		// log.Fatal("Failed to create session: ", err)
-	}
-	defer session.Close()
-
-	t.stdout = new(bytes.Buffer)
-	stderr := new(bytes.Buffer)
-	session.Stdout = t.stdout
-	session.Stderr = stderr
-	err = session.Run(command)
-
-	if err != nil {
-		t.err = fmt.Errorf("%v: %v\n\n%s\n\n%s", t.err, err, strings.TrimRight(stderr.String(), "\n"), "")
-	} else {
-		t.err = nil
-	}
-
-	tn := <-hn
-	t.hostname = tn.hostname
-	return
-}
-
-func getHostname(client *ssh.Client, out chan<- tentacle) {
-	defer close(out)
-
-	t := tentacle{
-		hostname: noHostNameText,
-	}
-	defer func() { out <- t }()
-
-	session, err := client.NewSession()
-	if err != nil {
-		t.err = fmt.Errorf("%v: %v", t.err, err)
-		return
-	}
-	defer session.Close()
-
-	t.stdout = new(bytes.Buffer)
-	stderr := new(bytes.Buffer)
-	session.Stdout = t.stdout
-	session.Stderr = stderr
-	err = session.Run("hostname")
-	if err == nil {
-		t.hostname = strings.TrimRight(t.stdout.String(), "\n")
-	}
-
-	return
-}
-
 func (t *tentacle) print() error {
 	fmt.Println("-----")
 	fmt.Println(t.hostname)
-	if t.hostname == noHostNameText {
-		fmt.Println(t.host)
-	}
 	fmt.Printf("-----\n\n")
 	o := strings.TrimRight(t.stdout.String(), "\n")
 	if o != "" {