Add support for loading a hosts file

Signed-off-by: Blaine Gardner <[email protected]>

cmd/octopus/hostsfile.go

@@ -0,0 +1,63 @@
+package main
+
+import (
+	"bufio"
+	"fmt"
+	"os"
+	"os/exec"
+	"regexp"
+	"strings"
+)
+
+const (
+	defaultHostsFile = "_node-list"
+)
+
+func getAddrsFromHostsFile(hostGroups []string, hostsFile string) ([]string, error) {
+	f, err := os.Open(hostsFile)
+	if err != nil {
+		return []string{}, fmt.Errorf("could not load hosts file %s: %v", hostsFile, err)
+	}
+
+	fileGroups, err := getAllGroupsInFile(f)
+	if err != nil {
+		return []string{}, fmt.Errorf("error parsing hosts file %s: %v", hostsFile, err)
+	}
+
+	// Make a '${<group>}' argument for each group
+	gVars := []string{}
+	for _, g := range hostGroups {
+		if _, ok := fileGroups[g]; !ok {
+			return []string{}, fmt.Errorf("host group %s not found in hosts file %s", g, hostsFile)
+		}
+		gVars = append(gVars, fmt.Sprintf("${%s}", g))
+	}
+
+	// Source the hosts file, and echo all the groups without newlines
+	cmd := exec.Command("/bin/bash", "-c",
+		fmt.Sprintf("source %s ; echo %s", hostsFile, strings.Join(gVars, " ")))
+	o, err := cmd.CombinedOutput()
+	// convert to string which has exactly one newline
+	os := strings.TrimRight(string(o), "\n")
+	if err != nil {
+		return []string{}, fmt.Errorf("could not get groups %v from %s: %v\n%s", hostGroups, hostsFile, err, os)
+	}
+
+	addrs := strings.Split(os, " ")
+	return addrs, nil
+}
+
+func getAllGroupsInFile(f *os.File) (map[string]bool, error) {
+	scanner := bufio.NewScanner(f)
+	fileGroups := map[string]bool{}
+	// Regex to match Bash variable definition of a host group. Matches: <varname>="
+	// <varname> can be any bash variable; the double quote is required
+	varRegex, _ := regexp.Compile("^([a-zA-Z_][a-zA-Z0-9_]+)=\"")
+	for scanner.Scan() {
+		l := strings.TrimLeft(scanner.Text(), " \t")
+		if m := varRegex.FindStringSubmatch(l); m != nil {
+			fileGroups[m[1]] = true
+		}
+	}
+	return fileGroups, scanner.Err()
+}

cmd/octopus/octopus.go

@@ -11,33 +11,45 @@ import (
 )
 
 func main() {
+	command := flag.String("command", "", "(required) command to execute on remote hosts")
+	hostGroups := flag.String("host-groups", "",
+		"(required) named host groups on which to execute the command")
+	hostsFile := flag.String("hosts-file", defaultHostsFile, fmt.Sprintf(
+		"file which defines which remote hosts are available for execution (default: %s)", defaultHostsFile))
 	identityFile := flag.String("identity-file", "~/.ssh/id_rsa",
 		"identity file used to authenticate to remote hosts")
-	command := flag.String("command", "", "(required) command to execute on remote hosts")
 	flag.Parse()
 
 	if strings.Trim(*command, " \t") == "" {
 		fmt.Printf("ERROR! '-command' must be specified\n\n")
 		flag.PrintDefaults()
-		os.Exit(1)
+		os.Exit(-1)
+	}
+	if strings.Trim(*hostGroups, " \t") == "" {
+		fmt.Printf("ERROR! '-hosts' must be specified \n\n")
+		flag.PrintDefaults()
+		os.Exit(-1)
+	}
+
+	h := strings.Split(*hostGroups, ",")
+	hostAddrs, err := getAddrsFromHostsFile(h, *hostsFile)
+	if err != nil {
+		log.Fatalf("%v", err)
 	}
 
 	config, err := newCommandConfig(*identityFile)
 	if err != nil {
 		log.Fatalf("could not generate command config: %v", err)
 	}
 
-	hosts := []string{"10.86.1.87", "10.86.1.103"}
-	tentacles := make(chan tentacle, len(hosts))
-
-	for i := 0; i < len(hosts); i++ {
-		go runCommand(hosts[i], *command, config, tentacles)
+	tch := make(chan tentacle, len(hostAddrs))
+	for i := 0; i < len(hostAddrs); i++ {
+		go runCommand(hostAddrs[i], *command, config, tch)
 	}
 
 	numErrors := 0
-
-	for range hosts {
-		t := <-tentacles
+	for range hostAddrs {
+		t := <-tch
 		err := t.print()
 		if err != nil {
 			numErrors++