feat(ssh_executor): Sudo option (#488)

* feat(ssh_executor): Sudo option Signed-off-by: Louis-Quentin <[email protected]> * feat(ssh_executor): add sudo option Signed-off-by: Louis-Quentin <[email protected]> * feat(ssh_executor): typo in test Signed-off-by: Louis-Quentin <[email protected]>
ovh · Feb 7, 2022 · 3d44aa6 · 3d44aa6
1 parent bc593f6
commit 3d44aa6
Show file tree

Hide file tree

Showing 5 changed files with 176 additions and 21 deletions.
diff --git a/executors/ssh/README.md b/executors/ssh/README.md
@@ -13,6 +13,8 @@ In your yaml file, you can use:
   - user optional (default is OS username)
   - password optional (mandatory if no privatekey is found)
   - privatekey optional (default is $HOME/.ssh/id_rsa)
+  - sudo optional
+  - sudopassword optional (default to password)
 ```
 
 Example
@@ -30,7 +32,29 @@ testcases:
     - result.code ShouldEqual 0
     - result.timeseconds ShouldBeLessThan 1
 
+- name: Use specific privatekey
+  steps:
+  - type: ssh
+    host: 10.0.1.5:2222
+    command: echo 'foo'
+    user: bar
+    privatekey: /home/foo/.ssh/id_rsa
+    assertions:
+    - result.code ShouldEqual 0
+
+- name: Execute command as another user than bar
+  steps:
+  - type: ssh
+    host: 10.0.1.5:2222
+    command: echo 'foo'
+    user: bar
+    sudo: root
+    sudopassword: '{{.mypassword}}'
+    assertions:
+    - result.code ShouldEqual 0
+
 ```
+*NB: Sudo option uses a pseudotty*
 
 ## Output
 

diff --git a/executors/ssh/ssh.go b/executors/ssh/ssh.go
@@ -1,24 +1,24 @@
 package ssh
 
 import (
-	"bytes"
 	"context"
 	"fmt"
+	"github.com/mitchellh/mapstructure"
+	"github.com/ovh/venom"
+	"golang.org/x/crypto/ssh"
+	"io"
 	"os"
 	"os/user"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
-
-	"github.com/mitchellh/mapstructure"
-	"golang.org/x/crypto/ssh"
-
-	"github.com/ovh/venom"
+	"unicode/utf8"
 )
 
 // Name for test ssh
 const Name = "ssh"
+const sudoprompt = "sudo_venom"
 
 // New returns a new Test Exec
 func New() venom.Executor {
@@ -27,11 +27,13 @@ func New() venom.Executor {
 
 // Executor represents a Test Exec
 type Executor struct {
-	Host       string `json:"host,omitempty" yaml:"host,omitempty"`
-	Command    string `json:"command,omitempty" yaml:"command,omitempty"`
-	User       string `json:"user,omitempty" yaml:"user,omitempty"`
-	Password   string `json:"password,omitempty" yaml:"password,omitempty"`
-	PrivateKey string `json:"privatekey,omitempty" yaml:"privatekey,omitempty"`
+	Host         string `json:"host,omitempty" yaml:"host,omitempty"`
+	Command      string `json:"command,omitempty" yaml:"command,omitempty"`
+	User         string `json:"user,omitempty" yaml:"user,omitempty"`
+	Password     string `json:"password,omitempty" yaml:"password,omitempty"`
+	PrivateKey   string `json:"privatekey,omitempty" yaml:"privatekey,omitempty"`
+	Sudo         string `json:"sudo,omitempty" yaml:"sudo,omitempty"`
+	SudoPassword string `json:"sudopassword,omitempty" yaml:"sudopassword,omitempty"`
 }
 
 // Result represents a step result
@@ -67,17 +69,30 @@ func (Executor) Run(ctx context.Context, step venom.TestStep) (interface{}, erro
 	start := time.Now()
 	result := Result{}
 
-	client, session, err := connectToHost(e.User, e.Password, e.PrivateKey, e.Host)
+	client, session, err := connectToHost(e.User, e.Password, e.PrivateKey, e.Host, e.Sudo)
 	if err != nil {
 		result.Err = err.Error()
 	} else {
 		defer client.Close()
-		stdout := &bytes.Buffer{}
-		stderr := &bytes.Buffer{}
+		stdout := &Buffer{}
+		stderr := &Buffer{}
 
 		session.Stderr = stderr
 		session.Stdout = stdout
-		if err := session.Run(e.Command); err != nil {
+		stdin, _ := session.StdinPipe()
+
+		// Handle sudo password
+		command := e.Command
+		quit := make(chan bool)
+		if e.Sudo != "" {
+			command = "TERM=xterm-mono sudo -S -p " + sudoprompt + " -u " + e.Sudo + " " + command
+			if e.SudoPassword == "" {
+				e.SudoPassword = e.Password
+			}
+			go handleSudo(stdin, stdout, quit, e.SudoPassword)
+		}
+
+		if err := session.Run(command); err != nil {
 			if exiterr, ok := err.(*ssh.ExitError); ok {
 				status := exiterr.ExitStatus()
 				result.Code = strconv.Itoa(status)
@@ -92,8 +107,11 @@ func (Executor) Run(ctx context.Context, step venom.TestStep) (interface{}, erro
 			result.Code = "0"
 		}
 
-		result.Systemerr = stderr.String()
-		result.Systemout = stdout.String()
+		if e.Sudo != "" {
+			quit <- true
+		}
+		result.Systemerr = strings.TrimSpace(stderr.String())
+		result.Systemout = strings.TrimSpace(stdout.String())
 	}
 
 	elapsed := time.Since(start)
@@ -102,7 +120,26 @@ func (Executor) Run(ctx context.Context, step venom.TestStep) (interface{}, erro
 	return result, nil
 }
 
-func connectToHost(u, pass, key, host string) (*ssh.Client, *ssh.Session, error) {
+func handleSudo(in io.Writer, out *Buffer, quit chan bool, password string) {
+	sudopromptlen := len(sudoprompt)
+	for {
+		select {
+		case <-quit:
+			return
+		default:
+			content := out.String()
+			bufferLen := utf8.RuneCountInString(content)
+
+			// Check if we have to enter password
+			if bufferLen >= sudopromptlen && strings.Contains(content[bufferLen-sudopromptlen:], sudoprompt) {
+				in.Write([]byte(password + "\n"))
+				out.Truncate(0)
+			}
+		}
+	}
+}
+
+func connectToHost(u, pass, key, host, sudo string) (*ssh.Client, *ssh.Session, error) {
 	//Default user is current username
 	if u == "" {
 		osUser, err := user.Current()
@@ -112,9 +149,9 @@ func connectToHost(u, pass, key, host string) (*ssh.Client, *ssh.Session, error)
 		u = osUser.Username
 	}
 
-	//If password is set, use it
+	//If password is set, and we don't have key use it
 	var auth []ssh.AuthMethod
-	if pass != "" {
+	if pass != "" && key == "" {
 		auth = []ssh.AuthMethod{ssh.Password(pass)}
 	} else {
 		//Load the the private key
@@ -149,6 +186,18 @@ func connectToHost(u, pass, key, host string) (*ssh.Client, *ssh.Session, error)
 		return nil, nil, err
 	}
 
+	// Request PTY for sudo cmd
+	if sudo != "" {
+		modes := ssh.TerminalModes{
+			ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
+			ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
+		}
+
+		if err := session.RequestPty("xterm", 40, 80, modes); err != nil {
+			return nil, nil, err
+		}
+	}
+
 	return client, session, nil
 }
 

diff --git a/executors/ssh/sshutil.go b/executors/ssh/sshutil.go
@@ -0,0 +1,54 @@
+package ssh
+
+import (
+	"bytes"
+	"sync"
+)
+
+// Buffer thread safe
+type Buffer struct {
+	b bytes.Buffer
+	m sync.Mutex
+}
+
+// Read thread safe
+func (b *Buffer) Read(p []byte) (n int, err error) {
+	b.m.Lock()
+	defer b.m.Unlock()
+	return b.b.Read(p)
+}
+
+// Write thread safe
+func (b *Buffer) Write(p []byte) (n int, err error) {
+	b.m.Lock()
+	defer b.m.Unlock()
+	return b.b.Write(p)
+}
+
+// String thread safe
+func (b *Buffer) String() string {
+	b.m.Lock()
+	defer b.m.Unlock()
+	return b.b.String()
+}
+
+// Bytes thread safe
+func (b *Buffer) Bytes() []byte {
+	b.m.Lock()
+	defer b.m.Unlock()
+	return b.b.Bytes()
+}
+
+// Len thread safe
+func (b *Buffer) Len() int {
+	b.m.Lock()
+	defer b.m.Unlock()
+	return b.b.Len()
+}
+
+// Truncate thread safe
+func (b *Buffer) Truncate(n int) {
+	b.m.Lock()
+	defer b.m.Unlock()
+	b.b.Truncate(n)
+}
diff --git a/tests/Makefile b/tests/Makefile
@@ -42,7 +42,7 @@ venom-kafka.cid:
 venom-rabbit.cid:
 	$(call docker_run,rabbitmq,rabbitmq,-p 5672:5672 -p 15672:15672)
 venom-sshd.cid:
-	$(call docker_run,ghcr.io/linuxserver/openssh-server,sshd,-p 2222:2222 -e PUID=1000 -e PGID=1000 -e TZ=Europe/London -e PUBLIC_KEY="$(shell cat ~/.ssh/id_rsa.pub)" -e USER_NAME=venom )
+	$(call docker_run,ghcr.io/linuxserver/openssh-server,sshd,-p 2222:2222 -e PUID=1000 -e PGID=1000 -e TZ=Europe/London -e PUBLIC_KEY="$(shell cat ~/.ssh/id_rsa.pub)" -e USER_NAME=venom -e USER_PASSWORD=testvenom -e SUDO_ACCESS=true)
 venom-mqtt.cid:
 	$(call docker_run,eclipse-mosquitto,mqtt-broker,-p 1883:1883 -p 9001:9001 -v $(shell realpath mqtt/mosquitto.conf):/mosquitto/config/mosquitto.conf:ro)
 venom-qpid.cid:

diff --git a/tests/ssh.yml b/tests/ssh.yml
@@ -10,3 +10,31 @@ testcases:
     assertions:
     - result.code ShouldEqual 0
     - result.timeseconds ShouldBeLessThan 10
+
+- name: ssh sudo as root
+  steps:
+  - type: ssh
+    user: venom
+    host: localhost:2222
+    privatekey: "$HOME/.ssh/id_rsa"
+    command: whoami
+    sudo: root
+    sudopassword: testvenom
+    assertions:
+    - result.code ShouldEqual 0
+    - result.systemout ShouldEqual root
+    - result.timeseconds ShouldBeLessThan 10
+
+- name: ssh sudo as self
+  steps:
+  - type: ssh
+    user: venom
+    host: localhost:2222
+    privatekey: "$HOME/.ssh/id_rsa"
+    command: whoami
+    sudo: venom
+    sudopassword: testvenom
+    assertions:
+    - result.code ShouldEqual 0
+    - result.systemout ShouldEqual venom
+    - result.timeseconds ShouldBeLessThan 10