restructuring inspector-driver integration

driver/driver.go

@@ -1,18 +1,28 @@
 package driver
 
-type fields struct {
-	// Supported inspector representations for specific driver
-	Supported []string
-	// Selected inspector representations
-	Selected []string
+// SystemInfo gives more insight into system details
+type SystemDetails struct {
+	IsWindows bool
+	IsLinux   bool
+	IsDarwin  bool
+	IsWeb     bool
+	Name      string
+	Extra     string
+}
+
+type driverBase struct {
 	// Polling interval between retrievals
 	PollInterval int64
+	Info         *SystemDetails
 }
 
+// Command represents the two commands ReadFile & RunCommand
+type Command func(string) (string, error)
+
 // Driver : specification of functions to be defined by every Driver
 type Driver interface {
 	ReadFile(path string) (string, error)
 	RunCommand(command string) (string, error)
 	// shows the driver details, not sure if we should be showing OS name
-	GetDetails() string
+	GetDetails() SystemDetails
 }

driver/local.go

@@ -1,20 +1,18 @@
 package driver
 
 import (
-	"fmt"
 	"io/ioutil"
 	"os"
 	"os/exec"
 	"runtime"
-	"strings"
 
 	log "github.com/sirupsen/logrus"
 )
 
 // Local : Driver for handling local executions
 type Local struct {
-	fields
-	Vars []string
+	driverBase
+	EnvVars []string
 }
 
 func (d *Local) ReadFile(path string) (string, error) {
@@ -26,26 +24,52 @@ func (d *Local) ReadFile(path string) (string, error) {
 	return string(content), nil
 }
 
+// RunCommand : For simple commands without shell variables, pipes, e.t.c
+// They can be passed directly. For complex commands e.g
+// `echo something | awk $var`, turn into a file to be saved
+// under ./shell/
 func (d *Local) RunCommand(command string) (string, error) {
 	// FIXME: If command contains a shell variable $ or glob
 	// type pattern, it would not be executed, see
 	// https://pkg.go.dev/os/exec for more information
-	cmdArgs := strings.Fields(command)
+	var cmd *exec.Cmd
 	log.Debugf("Running command `%s` ", command)
-	cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
+	if d.Info == nil {
+		d.GetDetails()
+	}
+	if d.Info.IsLinux || d.Info.IsDarwin {
+		cmd = exec.Command("bash", "-c", command)
+	} else {
+		cmd = exec.Command("cmd", "/C", command)
+	}
 	cmd.Env = os.Environ()
-	if len(d.Vars) != 0 {
-		for _, v := range d.Vars {
+	if len(d.EnvVars) != 0 {
+		for _, v := range d.EnvVars {
 			cmd.Env = append(cmd.Env, v)
 		}
 	}
+	_ = cmd.Wait()
 	out, err := cmd.Output()
 	if err != nil {
 		return ``, err
 	}
 	return string(out), nil
 }
 
-func (d *Local) GetDetails() string {
-	return fmt.Sprintf(`Local - %s`, runtime.GOOS)
+func (d *Local) GetDetails() SystemDetails {
+	if d.Info == nil {
+		details := &SystemDetails{}
+		details.Name = runtime.GOOS
+		switch details.Name {
+		case "windows":
+			details.IsWindows = true
+		case "linux":
+			details.IsLinux = true
+		case "darwin":
+			details.IsDarwin = true
+		}
+		details.Extra = runtime.GOARCH
+		d.Info = details
+	}
+	return *d.Info
 }

driver/local_unix_test.go

@@ -15,3 +15,11 @@ func TestUnixLocalRunCommand(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func TestUnixLocalSystemDetails(t *testing.T) {
+	d := Local{}
+	details := d.GetDetails()
+	if !(details.IsLinux || details.IsDarwin) {
+		t.Errorf("Expected Darwin or Linux on unix test, got %s", details.Name)
+	}
+}

driver/local_windows_test.go

@@ -13,10 +13,10 @@ func TestWindowsRunCommand(t *testing.T) {
 	}
 }
 
-func TestWindowsLocalGetDetails(t *testing.T) {
+func TestWindowsLocalSystemDetails(t *testing.T) {
 	d := Local{}
-	output := d.GetDetails()
-	if output != "Local - windows" {
-		t.Error(output)
+	details := d.GetDetails()
+	if !details.IsWindows {
+		t.Errorf("Expected windows got %s", details.Name)
 	}
 }

driver/ssh.go

@@ -13,7 +13,7 @@ var port = 22
 
 // SSH : Driver for handling ssh executions
 type SSH struct {
-	fields
+	driverBase
 	// User e.g root
 	User string
 	// Host name/ip e.g 171.23.122.1
@@ -27,7 +27,8 @@ type SSH struct {
 	// Check known hosts (only disable for tests
 	CheckKnownHosts bool
 	// set environmental vars for server e.g []string{"DEBUG=1", "FAKE=echo"}
-	Vars []string
+	EnvVars       []string
+	SessionClient *goph.Client
 }
 
 func (d *SSH) String() string {
@@ -36,33 +37,39 @@ func (d *SSH) String() string {
 
 // set the goph Client
 func (d *SSH) Client() (*goph.Client, error) {
-	var err error
-	var client *goph.Client
-	var callback ssh.HostKeyCallback
-	if d.Port != 0 {
-		port = d.Port
-	}
-	auth, err := goph.Key(d.KeyFile, d.KeyPass)
-	if err != nil {
-		return nil, err
-	}
-	if d.CheckKnownHosts {
-		callback, err = goph.DefaultKnownHosts()
+	if d.SessionClient == nil {
+		var err error
+		var client *goph.Client
+		var callback ssh.HostKeyCallback
+		if d.Port != 0 {
+			port = d.Port
+		}
+		auth, err := goph.Key(d.KeyFile, d.KeyPass)
 		if err != nil {
 			return nil, err
 		}
-	} else {
-		callback = ssh.InsecureIgnoreHostKey()
+		if d.CheckKnownHosts {
+			callback, err = goph.DefaultKnownHosts()
+			if err != nil {
+				return nil, err
+			}
+		} else {
+			callback = ssh.InsecureIgnoreHostKey()
+		}
+		client, err = goph.NewConn(&goph.Config{
+			User:     d.User,
+			Addr:     d.Host,
+			Port:     uint(port),
+			Auth:     auth,
+			Timeout:  goph.DefaultTimeout,
+			Callback: callback,
+		})
+		if err != nil {
+			d.SessionClient = client
+		}
+		return client, err
 	}
-	client, err = goph.NewConn(&goph.Config{
-		User:     d.User,
-		Addr:     d.Host,
-		Port:     uint(port),
-		Auth:     auth,
-		Timeout:  goph.DefaultTimeout,
-		Callback: callback,
-	})
-	return client, err
+	return d.SessionClient, nil
 }
 
 func (d *SSH) ReadFile(path string) (string, error) {
@@ -72,16 +79,16 @@ func (d *SSH) ReadFile(path string) (string, error) {
 }
 
 func (d *SSH) RunCommand(command string) (string, error) {
-	// FIXME: Do we retain client across all command runs?
+	// TODO: Ensure clients of all SSH drivers are closed on context end
+	// i.e d.SessionClient.Close()
 	log.Debugf("Running remote command %s", command)
 	client, err := d.Client()
 	if err != nil {
 		return ``, err
 	}
-	defer client.Close()
-	if len(d.Vars) != 0 {
+	if len(d.EnvVars) != 0 {
 		// add env variable to command
-		envline := strings.Join(d.Vars, ";")
+		envline := strings.Join(d.EnvVars, ";")
 		command = strings.Join([]string{envline, command}, ";")
 	}
 	out, err := client.Run(command)
@@ -91,20 +98,29 @@ func (d *SSH) RunCommand(command string) (string, error) {
 	return string(out), nil
 }
 
-func (d *SSH) GetDetails() string {
-	return fmt.Sprintf(`SSH - %s`, d.String())
-}
-
-func NewSSHForTest() *SSH {
-	return &SSH{
-		User:            "dev",
-		Host:            "127.0.0.1",
-		Port:            2222,
-		KeyFile:         "/home/deven/.ssh/id_rsa",
-		KeyPass:         "",
-		CheckKnownHosts: false,
-		fields: fields{
-			PollInterval: 5,
-		},
+func (d *SSH) GetDetails() SystemDetails {
+	if d.Info == nil {
+		uname, err := d.RunCommand(`uname`)
+		// try windows command
+		if err != nil {
+			windowsName, err := d.RunCommand(`systeminfo | findstr /B /C:"OS Name"`)
+			if err == nil {
+				if strings.Contains(strings.ToLower(windowsName), "windows") {
+					uname = "windows"
+				}
+			}
+		}
+		details := &SystemDetails{}
+		details.Name = uname
+		switch details.Name {
+		case "windows":
+			details.IsWindows = true
+		case "linux":
+			details.IsLinux = true
+		case "darwin":
+			details.IsDarwin = true
+		}
+		d.Info = details
 	}
+	return *d.Info
 }

driver/ssh_test.go

@@ -5,10 +5,32 @@ import (
 	"testing"
 )
 
+func NewSSHForTest() Driver {
+	return &SSH{
+		User:            "dev",
+		Host:            "127.0.0.1",
+		Port:            2222,
+		KeyFile:         "/home/diretnan/.ssh/id_rsa",
+		KeyPass:         "",
+		CheckKnownHosts: false,
+		driverBase: driverBase{
+			PollInterval: 5,
+		},
+	}
+}
+
 func TestSSHRunCommand(t *testing.T) {
 	d := NewSSHForTest()
 	output, err := d.RunCommand(`ps -A`)
 	if err != nil || !strings.Contains(output, "PID") {
 		t.Error(err)
 	}
 }
+
+func TestSSHSystemDetails(t *testing.T) {
+	d := NewSSHForTest()
+	details := d.GetDetails()
+	if !details.IsLinux {
+		t.Errorf("Expected linux server for ssh test got %s", details.Name)
+	}
+}

driver/web.go

@@ -21,9 +21,9 @@ const (
 	GET Request = "GET"
 )
 
-// Web : Driver for handling ssh executions
+// Web : Driver for handling web executions
 type Web struct {
-	fields
+	driverBase
 	// URL e.g https://google.com
 	URL string
 	// Method POST/GET
@@ -53,7 +53,7 @@ func (d *Web) RunCommand(command string) (string, error) {
 			res, err = http.Get(d.URL)
 		}
 		if err != nil || res.StatusCode < 200 || res.StatusCode > 299 {
-			message := fmt.Sprintf("Error %s running request: %s", err, string(res.StatusCode))
+			message := fmt.Sprintf("Error %s running request: %d", err, res.StatusCode)
 			return ``, errors.New(message)
 		}
 		elapsed := time.Since(start)
@@ -62,16 +62,14 @@ func (d *Web) RunCommand(command string) (string, error) {
 	return ``, errors.New("Cannot read file on web driver")
 }
 
-func (d *Web) GetDetails() string {
-	return fmt.Sprintf(`Web - %s`, d.String())
-}
-
-func NewWebForTest() *Web {
-	return &Web{
-		URL:    "https://duckduckgo.com",
-		Method: GET,
-		fields: fields{
-			PollInterval: 5,
-		},
+func (d *Web) GetDetails() SystemDetails {
+	if d.Info == nil {
+		details := &SystemDetails{
+			Name:  "web",
+			Extra: d.URL,
+			IsWeb: true,
+		}
+		d.Info = details
 	}
+	return *d.Info
 }