add local_shell playbook support

umputun · Feb 1, 2024 · 5b8c2de · 5b8c2de
1 parent c1d37f1
commit 5b8c2de
Show file tree

Hide file tree

Showing 6 changed files with 83 additions and 27 deletions.
diff --git a/README.md b/README.md
@@ -169,6 +169,7 @@ Spot supports the following command-line options:
 user: umputun                       # default ssh user. Can be overridden by -u flag or by inventory or host definition
 ssh_key: keys/id_rsa                # ssh key
 ssh_shell: /bin/bash                # shell to use for remote ssh execution, default is /bin/sh
+local_shell: /bin/bash              # shell to use for local execution, default is os shell
 inventory: /etc/spot/inventory.yml  # default inventory file. Can be overridden by --inventory flag
 
 # list of targets, i.e. hosts, inventory files or inventory URLs
@@ -239,6 +240,7 @@ In some cases, the rich syntax of the full playbook is not needed and can feel o
 user: umputun                       # default ssh user. Can be overridden by -u flag or by inventory or host definition
 ssh_key: keys/id_rsa                # ssh key
 ssh_shell: /bin/bash                # shell to use for remote ssh execution, default is /bin/sh
+local_shell: /bin/bash              # shell to use for local execution, default is os shell
 inventory: /etc/spot/inventory.yml  # default inventory file. Can be overridden by --inventory flag
 
 targets: ["devbox1", "devbox2", "h1.example.com:2222", "h2.example.com"] # list of host names from inventory and direct host ips

diff --git a/pkg/config/command.go b/pkg/config/command.go
@@ -32,8 +32,9 @@ type Cmd struct {
 	Register    []string          `yaml:"register" toml:"register"` // register variables from command
 	OnExit      string            `yaml:"on_exit" toml:"on_exit"`   // script to run on exit
 
-	Secrets  map[string]string `yaml:"-" toml:"-"` // loaded secrets, filled by playbook
-	SSHShell string            `yaml:"-" toml:"-"` // shell to use for ssh commands, filled by playbook
+	Secrets    map[string]string `yaml:"-" toml:"-"` // loaded secrets, filled by playbook
+	SSHShell   string            `yaml:"-" toml:"-"` // shell to use for ssh commands, filled by playbook
+	LocalShell string            `yaml:"-" toml:"-"` // shell to use for local commands, filled by playbooks
 }
 
 // CmdOptions defines options for a command
@@ -396,17 +397,24 @@ func (cmd *Cmd) validate() error {
 	return nil
 }
 
+// shell returns the shell to use for multi-line commands.
+// If Local is set, it returns LocalShell, otherwise SSHShell.
+// If LocalShell is not set, it returns OS default shell and if this one is not set, it returns /bin/sh.
+// For SSHShell, it returns /bin/sh if not sets.
 func (cmd *Cmd) shell() string {
-	if cmd.SSHShell == "" {
-		return "/bin/sh"
-	}
 	if cmd.Options.Local {
-		envShell := os.Getenv("SHELL")
-		if envShell == "" {
-			return "/bin/sh"
+		if cmd.LocalShell != "" {
+			return cmd.LocalShell
+		}
+		res := os.Getenv("SHELL")
+		if res != "" {
+			return res
 		}
-		return envShell
+		return "/bin/sh"
 	}
 
+	if cmd.SSHShell == "" {
+		return "/bin/sh"
+	}
 	return cmd.SSHShell
 }
diff --git a/pkg/config/command_test.go b/pkg/config/command_test.go
@@ -2,6 +2,7 @@ package config
 
 import (
 	"io"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -760,3 +761,25 @@ func TestHasShebang(t *testing.T) {
 		})
 	}
 }
+
+func TestCmd_shell(t *testing.T) {
+	t.Run("shell is not set", func(t *testing.T) {
+		c := Cmd{}
+		assert.Equal(t, "/bin/sh", c.shell())
+	})
+
+	t.Run("shell is set", func(t *testing.T) {
+		c := Cmd{SSHShell: "/bin/bash"}
+		assert.Equal(t, "/bin/bash", c.shell())
+	})
+
+	t.Run("shell is not set, local", func(t *testing.T) {
+		c := Cmd{Options: CmdOptions{Local: true}}
+		assert.Equal(t, os.Getenv("SHELL"), c.shell())
+	})
+
+	t.Run("shell is set, local", func(t *testing.T) {
+		c := Cmd{LocalShell: "/bin/bash", Options: CmdOptions{Local: true}}
+		assert.Equal(t, "/bin/bash", c.shell())
+	})
+}
diff --git a/pkg/config/playbook.go b/pkg/config/playbook.go
@@ -24,12 +24,13 @@ import (
 
 // PlayBook defines the top-level config object
 type PlayBook struct {
-	User      string            `yaml:"user" toml:"user"`           // ssh user
-	SSHKey    string            `yaml:"ssh_key" toml:"ssh_key"`     // ssh key
-	SSHShell  string            `yaml:"ssh_shell" toml:"ssh_shell"` // ssh shell to use
-	Inventory string            `yaml:"inventory" toml:"inventory"` // inventory file or url
-	Targets   map[string]Target `yaml:"targets" toml:"targets"`     // list of targets/environments
-	Tasks     []Task            `yaml:"tasks" toml:"tasks"`         // list of tasks
+	User       string            `yaml:"user" toml:"user"`               // ssh user
+	SSHKey     string            `yaml:"ssh_key" toml:"ssh_key"`         // ssh key
+	SSHShell   string            `yaml:"ssh_shell" toml:"ssh_shell"`     // ssh shell to use
+	LocalShell string            `yaml:"local_shell" toml:"local_shell"` // local shell to use
+	Inventory  string            `yaml:"inventory" toml:"inventory"`     // inventory file or url
+	Targets    map[string]Target `yaml:"targets" toml:"targets"`         // list of targets/environments
+	Tasks      []Task            `yaml:"tasks" toml:"tasks"`             // list of tasks
 
 	inventory       *InventoryData    // loaded inventory
 	overrides       *Overrides        // overrides passed from cli
@@ -45,12 +46,14 @@ type SecretsProvider interface {
 // SimplePlayBook defines simplified top-level config
 // It is used for unmarshalling only, and result used to make the usual PlayBook
 type SimplePlayBook struct {
-	User      string   `yaml:"user" toml:"user"`           // ssh user
-	SSHKey    string   `yaml:"ssh_key" toml:"ssh_key"`     // ssh key
-	Inventory string   `yaml:"inventory" toml:"inventory"` // inventory file or url
-	Targets   []string `yaml:"targets" toml:"targets"`     // list of names
-	Target    string   `yaml:"target" toml:"target"`       // a single target to run task on
-	Task      []Cmd    `yaml:"task" toml:"task"`           // single task is a list of commands
+	User       string   `yaml:"user" toml:"user"`               // ssh user
+	SSHKey     string   `yaml:"ssh_key" toml:"ssh_key"`         // ssh key
+	SSHShell   string   `yaml:"ssh_shell" toml:"ssh_shell"`     // ssh shell to uses
+	LocalShell string   `yaml:"local_shell" toml:"local_shell"` // local shell to use
+	Inventory  string   `yaml:"inventory" toml:"inventory"`     // inventory file or url
+	Targets    []string `yaml:"targets" toml:"targets"`         // list of names
+	Target     string   `yaml:"target" toml:"target"`           // a single target to run task on
+	Task       []Cmd    `yaml:"task" toml:"task"`               // single task is a list of commands
 }
 
 // Task defines multiple commands runs together
@@ -154,7 +157,9 @@ func New(fname string, overrides *Overrides, secProvider SecretsProvider) (res *
 	log.Printf("[INFO] playbook loaded with %d tasks", len(res.Tasks))
 	for i, tsk := range res.Tasks {
 		for j, c := range tsk.Commands {
-			res.Tasks[i].Commands[j].SSHShell = res.shell() // set secrets for all commands in task
+			// set shell (remote and local) for all commands in task
+			res.Tasks[i].Commands[j].SSHShell = res.remoteShell()
+			res.Tasks[i].Commands[j].LocalShell = res.localShell()
 			log.Printf("[DEBUG] load command %q (task: %s)", c.Name, tsk.Name)
 		}
 	}
@@ -300,7 +305,8 @@ func (p *PlayBook) Task(name string) (*Task, error) {
 		if name == "ad-hoc" && p.overrides.AdHocCommand != "" {
 			// special case for ad-hoc command, make a fake task with a single command from overrides.AdHocCommand
 			return &Task{Name: "ad-hoc", Commands: []Cmd{
-				{Name: "ad-hoc", Script: p.overrides.AdHocCommand, SSHShell: p.shell()}}}, nil
+				{Name: "ad-hoc", Script: p.overrides.AdHocCommand,
+					SSHShell: p.remoteShell(), LocalShell: p.localShell()}}}, nil
 		}
 		for _, t := range tsk {
 			if strings.EqualFold(t.Name, name) {
@@ -590,7 +596,8 @@ func (p *PlayBook) loadSecrets() error {
 	return nil
 }
 
-func (p *PlayBook) shell() string {
+// remoteShell returns the remoteShell to use for SSH commands, using overrides if set, or playbook's remoteShell if not.
+func (p *PlayBook) remoteShell() string {
 	if p.overrides != nil && p.overrides.SSHShell != "" {
 		return p.overrides.SSHShell
 	}
@@ -599,3 +606,16 @@ func (p *PlayBook) shell() string {
 	}
 	return "/bin/sh"
 }
+
+// localShell returns the local shell to use for local commands, using playbook's localShell if set,
+// or default to env:SHELL if not. If SHELL is not set in envs, it defaults to /bin/sh.
+func (p *PlayBook) localShell() string {
+	if p.LocalShell != "" {
+		return p.LocalShell
+	}
+	res := os.Getenv("SHELL")
+	if res != "" {
+		return res
+	}
+	return "/bin/sh"
+}
diff --git a/pkg/config/playbook_test.go b/pkg/config/playbook_test.go
@@ -23,6 +23,7 @@ func TestPlaybook_New(t *testing.T) {
 		assert.Equal(t, 1, len(c.Tasks), "single task")
 		assert.Equal(t, "umputun", c.User, "user")
 		assert.Equal(t, "/bin/sh", c.Tasks[0].Commands[0].SSHShell, "ssh shell")
+		assert.Equal(t, os.Getenv("SHELL"), c.Tasks[0].Commands[0].LocalShell, "local shell")
 
 		tsk := c.Tasks[0]
 		assert.Equal(t, 5, len(tsk.Commands), "5 commands")
@@ -197,19 +198,20 @@ func TestPlaybook_New(t *testing.T) {
 		t.Logf("%+v", c)
 		assert.Equal(t, 1, len(c.Tasks), "single task")
 		assert.Equal(t, "umputun", c.User, "user")
-		assert.Equal(t, "/bin/bash", c.Tasks[0].Commands[0].SSHShell, "ssh shell")
+		assert.Equal(t, "/bin/bash", c.Tasks[0].Commands[0].SSHShell, "remote ssh shell")
+		assert.Equal(t, "/bin/xxx", c.Tasks[0].Commands[0].LocalShell, "local local shell")
 
 		tsk := c.Tasks[0]
 		assert.Equal(t, 5, len(tsk.Commands), "5 commands")
 		assert.Equal(t, "deploy-remark42", tsk.Name, "task name")
 	})
-	t.Run("playbook with custom ssh shell override", func(t *testing.T) {
+	t.Run("playbook with custom shell overrides", func(t *testing.T) {
 		c, err := New("testdata/with-ssh-shell.yml", &Overrides{SSHShell: "/bin/zsh"}, nil)
 		require.NoError(t, err)
 		t.Logf("%+v", c)
 		assert.Equal(t, 1, len(c.Tasks), "single task")
 		assert.Equal(t, "umputun", c.User, "user")
-		assert.Equal(t, "/bin/zsh", c.Tasks[0].Commands[0].SSHShell, "ssh shell")
+		assert.Equal(t, "/bin/zsh", c.Tasks[0].Commands[0].SSHShell, "remote ssh shell")
 
 		tsk := c.Tasks[0]
 		assert.Equal(t, 5, len(tsk.Commands), "5 commands")

diff --git a/pkg/config/testdata/with-ssh-shell.yml b/pkg/config/testdata/with-ssh-shell.yml
@@ -1,5 +1,6 @@
 user: umputun
 ssh_shell: /bin/bash
+local_shell: /bin/xxx
 
 targets:
     remark42: