Add machine group

README.md

@@ -33,17 +33,20 @@ GLOBAL OPTIONS:
 See [releases](https://github.com/d3witt/viking/releases) for pre-built binaries.
 
 On Unix:
+
 ```
 env CGO_ENABLED=0 go install -ldflags="-s -w" github.com/d3witt/viking@latest
 ```
 
 On Windows cmd:
+
 ```
 set CGO_ENABLED=0
 go install -ldflags="-s -w" github.com/d3witt/viking@latest
 ```
 
 On Windows powershell:
+
 ```
 $env:CGO_ENABLED = '0'
 go install -ldflags="-s -w" github.com/d3witt/viking@latest
@@ -52,48 +55,74 @@ go install -ldflags="-s -w" github.com/d3witt/viking@latest
 ## 📄 Usage
 
 #### 🛰️ Add machine:
+
 ```
 $ viking machine add --name deathstar --key starkey 168.112.216.50
 Machine deathstar added.
 ```
-> [!NOTE]  
+
+> [!NOTE]
 > The key flag is not required. If a key is not specified, SSH Agent will be used to connect to the server.
 
 #### 📡 Exec command:
+
 ```
 $ viking exec deathstar echo 1234
 1234
 ```
 
 #### 📺 Connect to the machine:
+
 ```
 $ viking exec --tty deathstar /bin/bash
 root@deathstar:~$
 ```
 
+#### 🗂️ Machine group:
+
+```
+$ viking machine add -n dev -k starkey 168.112.216.50 117.51.181.37 24.89.193.43 77.79.125.157
+Machine dev added.
+
+$ viking exec dev echo 1234
+168.112.216.50: 1234
+117.51.181.37: 1234
+24.89.193.43: 1234
+77.79.125.157: 1234
+```
+
+> [!NOTE]
+> All machines in the group will run the same command at the same time. If there are errors, they will show up in the output. The execution will keep going despite the errors.
+
 #### 🔑 Add SSH key from a file
+
 ```
 $ viking key add --name starkey --passphrase dart ./id_rsa_star
 Key starkey added.
 ```
 
 #### 🆕 Generate SSH Key
+
 ```
 $ viking key generate --name starkey2
 Key starkey2 added.
 ```
 
 #### 📋 Copy public SSH Key
+
 ```
 $ viking key copy starkey2
 Public key copied to your clipboard.
 ```
 
 #### ⚙️ Custom config directory
+
 Viking saves data locally. Set `VIKING_CONFIG_DIR` env variable for a custom directory. Use `viking config` to check the current config folder.
 
-## 🤝  Missing a Feature?
+## 🤝 Missing a Feature?
+
 Feel free to open a new issue, or contact me.
 
-## 📘  License
+## 📘 License
+
 Viking is provided under the [MIT License](https://github.com/d3witt/viking/blob/main/LICENSE).

cli/command/cli.go

@@ -1,27 +1,12 @@
 package command
 
 import (
-	"errors"
-	"io"
-
 	"github.com/d3witt/viking/config"
-	"github.com/mattn/go-isatty"
-	"golang.org/x/term"
+	"github.com/d3witt/viking/streams"
 )
 
 type Cli struct {
 	Config   *config.Config
-	Out, Err io.Writer
-	In       io.ReadCloser
-	InFd     int
-	OutFd    int
-}
-
-// TerminalSize returns the width and height of the terminal.
-func (c *Cli) TerminalSize() (int, int, error) {
-	if !isatty.IsTerminal(uintptr(c.InFd)) {
-		return 0, 0, errors.New("not a terminal")
-	}
-
-	return term.GetSize(c.InFd)
+	Out, Err *streams.Out
+	In       *streams.In
 }

cli/command/machine/add.go

@@ -66,7 +66,7 @@ func runAdd(vikingCli *command.Cli, host, name, user, key string) error {
 
 	if err := vikingCli.Config.AddMachine(config.Machine{
 		Name:      name,
-		Host:      hostIp,
+		Host:      []net.IP{hostIp},
 		User:      user,
 		Key:       key,
 		CreatedAt: time.Now(),

cli/command/machine/execute.go

@@ -2,20 +2,20 @@ package machine
 
 import (
 	"fmt"
+	"io"
 	"strings"
+	"sync"
 
 	"github.com/d3witt/viking/cli/command"
 	"github.com/d3witt/viking/sshexec"
 	"github.com/urfave/cli/v2"
-	"golang.org/x/term"
 )
 
 func NewExecuteCmd(vikingCli *command.Cli) *cli.Command {
 	return &cli.Command{
 		Name:      "exec",
 		Usage:     "Execute shell command on machine",
-		Args:      true,
-		ArgsUsage: "MACHINE \"COMMAND\"",
+		ArgsUsage: "NAME \"COMMAND\"",
 		Flags: []cli.Flag{
 			&cli.BoolFlag{
 				Name:    "tty",
@@ -34,7 +34,7 @@ func NewExecuteCmd(vikingCli *command.Cli) *cli.Command {
 }
 
 func runExecute(vikingCli *command.Cli, machine string, cmd string, tty bool) error {
-	m, err := vikingCli.Config.GetMachine(machine)
+	m, err := vikingCli.Config.GetMachineByName(machine)
 	if err != nil {
 		return err
 	}
@@ -50,7 +50,41 @@ func runExecute(vikingCli *command.Cli, machine string, cmd string, tty bool) er
 		passphrase = key.Passphrase
 	}
 
-	client, err := sshexec.SshClient(m.Host.String(), m.User, private, passphrase)
+	if tty {
+		if len(m.Host) != 1 {
+			return fmt.Errorf("cannot allocate a pseudo-TTY to multiple hosts")
+		}
+
+		return executeTTY(vikingCli, m.Host[0].String(), m.User, private, passphrase, cmd)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(len(m.Host))
+
+	for _, host := range m.Host {
+		go func(ip string) {
+			defer wg.Done()
+
+			out := vikingCli.Out
+			errOut := vikingCli.Err
+			if len(m.Host) > 1 {
+				prefix := fmt.Sprintf("%s: ", ip)
+				out = out.WithPrefix(prefix)
+				errOut = errOut.WithPrefix(prefix + "error: ")
+			}
+
+			if err := execute(out, ip, m.User, private, passphrase, cmd); err != nil {
+				fmt.Fprintln(errOut, err.Error())
+			}
+		}(host.String())
+	}
+
+	wg.Wait()
+	return nil
+}
+
+func execute(out io.Writer, ip, user, private, passphrase, cmd string) error {
+	client, err := sshexec.SshClient(ip, user, private, passphrase)
 	if err != nil {
 		return err
 	}
@@ -59,32 +93,39 @@ func runExecute(vikingCli *command.Cli, machine string, cmd string, tty bool) er
 	sshCmd := sshexec.Command(sshexec.NewExecutor(client), cmd)
 	sshCmd.NoLogs = true
 
-	if tty {
-		w, h, err := vikingCli.TerminalSize()
-		if err != nil {
-			return err
-		}
+	output, err := sshCmd.CombinedOutput()
+	if handleSSHError(err) != nil {
+		return err
+	}
 
-		termState, err := term.GetState(vikingCli.OutFd)
-		if err != nil {
-			return fmt.Errorf("failed to get terminal state: %w", err)
-		}
-		defer term.Restore(vikingCli.OutFd, termState)
+	fmt.Fprint(out, string(output))
+	return nil
+}
 
-		term.MakeRaw(vikingCli.OutFd)
-		if err := sshCmd.StartInteractive(cmd, vikingCli.In, vikingCli.Out, vikingCli.Err, w, h); handleSSHError(err) != nil {
-			return err
-		}
+func executeTTY(vikingCli *command.Cli, ip, user, private, passphrase, cmd string) error {
+	client, err := sshexec.SshClient(ip, user, private, passphrase)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
 
-		return nil
+	sshCmd := sshexec.Command(sshexec.NewExecutor(client), cmd)
+	sshCmd.NoLogs = true
+
+	w, h, err := vikingCli.In.Size()
+	if err != nil {
+		return err
 	}
 
-	output, err := sshCmd.CombinedOutput()
-	if handleSSHError(err) != nil {
+	if err := vikingCli.Out.MakeRaw(); err != nil {
 		return err
 	}
+	defer vikingCli.Out.Restore()
 
-	fmt.Fprint(vikingCli.Out, string(output))
+	err = sshCmd.StartInteractive(cmd, vikingCli.In, vikingCli.Out, vikingCli.Err, w, h)
+	if handleSSHError(err) != nil {
+		return err
+	}
 
 	return nil
 }
@@ -93,6 +134,5 @@ func handleSSHError(err error) error {
 	if _, ok := err.(*sshexec.ExitError); ok {
 		return nil
 	}
-
 	return err
 }

cli/command/machine/list.go

@@ -1,7 +1,9 @@
 package machine
 
 import (
+	"net"
 	"sort"
+	"strings"
 
 	"github.com/d3witt/viking/cli/command"
 	"github.com/dustin/go-humanize"
@@ -28,16 +30,24 @@ func listMachines(vikingCli *command.Cli) error {
 	data := [][]string{
 		{
 			"NAME",
-			"HOST",
+			"HOSTS",
 			"KEY",
 			"CREATED",
 		},
 	}
 
 	for _, machine := range machines {
+		host := strings.Join(func(hosts []net.IP) []string {
+			strs := make([]string, len(hosts))
+			for i, h := range hosts {
+				strs[i] = h.String()
+			}
+			return strs
+		}(machine.Host), ", ")
+
 		data = append(data, []string{
 			machine.Name,
-			machine.Host.String(),
+			host,
 			machine.Key,
 			humanize.Time(machine.CreatedAt),
 		})

cli/command/machine/remove.go

@@ -12,7 +12,7 @@ func NewRmCmd(vikingCli *command.Cli) *cli.Command {
 		Name:      "rm",
 		Usage:     "Remove a machine",
 		Args:      true,
-		ArgsUsage: "MACHINE",
+		ArgsUsage: "NAME",
 		Action: func(ctx *cli.Context) error {
 			machine := ctx.Args().First()
 			return runRemove(vikingCli, machine)

config/machine.go

@@ -2,13 +2,12 @@ package config
 
 import (
 	"errors"
-	"net"
 	"time"
 )
 
 type Machine struct {
 	Name      string `toml:"-"`
-	Host      net.IP
+	Host      IPList
 	User      string
 	Key       string
 	CreatedAt time.Time
@@ -31,23 +30,6 @@ func (c *Config) ListMachines() []Machine {
 	return machines
 }
 
-// GetMachine returns a machine by name or host.
-func (c *Config) GetMachine(machine string) (Machine, error) {
-	if machine == "" {
-		return Machine{}, ErrMachineNameOrHostRequired
-	}
-
-	if m, err := c.GetMachineByName(machine); err == nil {
-		return m, nil
-	}
-
-	if m, err := c.GetMachineByHost(net.ParseIP(machine)); err == nil {
-		return m, nil
-	}
-
-	return Machine{}, ErrMachineNotFound
-}
-
 func (c *Config) GetMachineByName(name string) (Machine, error) {
 	if machine, ok := c.Machines[name]; ok {
 		machine.Name = name
@@ -57,19 +39,8 @@ func (c *Config) GetMachineByName(name string) (Machine, error) {
 	return Machine{}, ErrMachineNotFound
 }
 
-func (c *Config) GetMachineByHost(host net.IP) (Machine, error) {
-	for name, machine := range c.Machines {
-		if machine.Host.Equal(host) {
-			machine.Name = name
-			return machine, nil
-		}
-	}
-
-	return Machine{}, ErrMachineNotFound
-}
-
 func (c *Config) AddMachine(machine Machine) error {
-	_, err := c.GetMachine(machine.Name)
+	_, err := c.GetMachineByName(machine.Name)
 	if err == nil {
 		return ErrMachineAlreadyExists
 	}
@@ -81,7 +52,7 @@ func (c *Config) AddMachine(machine Machine) error {
 
 // RemoveMachine removes a machine from the config by name or host.
 func (c *Config) RemoveMachine(machine string) error {
-	m, err := c.GetMachine(machine)
+	m, err := c.GetMachineByName(machine)
 	if err != nil {
 		return err
 	}