merge 4.0.0-RC

1 files changed, 48 insertions, 25 deletions

diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go
index bbfb7be..6128018 100644
--- a/internal/ssh/client/authmethods.go
+++ b/internal/ssh/client/authmethods.go
@@ -1,93 +1,116 @@
 package client
 
 import (
+	"fmt"
 	"os"
 
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/ssh"
 
 	gossh "golang.org/x/crypto/ssh"
 )
 
+const addedPathStr string = "Added path to list of auth methods, not adding further methods"
+
 // InitSSHAuthMethods initialises all known SSH auth methods on the client side.
-func InitSSHAuthMethods(sshAuthMethods []gossh.AuthMethod, hostKeyCallback gossh.HostKeyCallback, trustAllHosts bool, throttleCh chan struct{}, privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
+func InitSSHAuthMethods(sshAuthMethods []gossh.AuthMethod,
+	hostKeyCallback gossh.HostKeyCallback, trustAllHosts bool, throttleCh chan struct{},
+	privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
+
 	if len(sshAuthMethods) > 0 {
 		simpleCallback, err := NewSimpleCallback()
 		if err != nil {
-			logger.FatalExit(err)
+			dlog.Client.FatalPanic(err)
 		}
 		return sshAuthMethods, simpleCallback
 	}
-
 	return initKnownHostsAuthMethods(trustAllHosts, throttleCh, privateKeyPath)
 }
 
-func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{}, privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
+func initIntegrationTestKnownHostsAuthMethods() []gossh.AuthMethod {
 	var sshAuthMethods []gossh.AuthMethod
+	privateKeyPath := "./id_rsa"
 
-	knownHostsPath := os.Getenv("HOME") + "/.ssh/known_hosts"
-	knownHostsCallback, err := NewKnownHostsCallback(knownHostsPath, trustAllHosts, throttleCh)
+	GeneratePrivatePublicKeyPairIfNotExists(privateKeyPath, 4096)
+	authMethod, err := ssh.PrivateKey(privateKeyPath)
 	if err != nil {
-		logger.FatalExit(knownHostsPath, err)
+		dlog.Client.FatalPanic("Unable to use private SSH key", privateKeyPath, err)
 	}
-	logger.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsPath)
 
-	if config.Common.ExperimentalFeaturesEnable {
-		sshAuthMethods = append(sshAuthMethods, gossh.Password("experimental feature test"))
-		logger.Debug("initKnownHostsAuthMethods", "Added experimental method to list of auth methods")
+	sshAuthMethods = append(sshAuthMethods, authMethod)
+	dlog.Client.Debug("initKnownHostsAuthMethods", addedPathStr, privateKeyPath)
+	return sshAuthMethods
+}
+
+func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
+	privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
+
+	var sshAuthMethods []gossh.AuthMethod
+	knownHostsFile := fmt.Sprintf("%s/.ssh/known_hosts", os.Getenv("HOME"))
+	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
+		// In case of integration test, override known hosts file path.
+		knownHostsFile = "./known_hosts"
 	}
 
-	// First try to read custom private key path.
+	knownHostsCallback, err := NewKnownHostsCallback(knownHostsFile, trustAllHosts, throttleCh)
+	if err != nil {
+		dlog.Client.FatalPanic(knownHostsFile, err)
+	}
+	dlog.Client.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsFile)
+
+	if config.Env("DTAIL_INTEGRATION_TEST_RUN_MODE") {
+		return initIntegrationTestKnownHostsAuthMethods(), knownHostsCallback
+	}
+
+	// Try to read custom private key path.
 	if privateKeyPath != "" {
 		authMethod, err := ssh.PrivateKey(privateKeyPath)
 		if err == nil {
 			sshAuthMethods = append(sshAuthMethods, authMethod)
-			logger.Debug("initKnownHostsAuthMethods", "Added path to list of auth methods, not adding further methods", privateKeyPath)
+			dlog.Client.Debug("initKnownHostsAuthMethods", addedPathStr, privateKeyPath)
 			return sshAuthMethods, knownHostsCallback
 		}
-		logger.FatalExit("Unable to use private SSH key", privateKeyPath, err)
+		dlog.Client.FatalPanic("Unable to use private SSH key", privateKeyPath, err)
 	}
 
 	// Second, try SSH Agent
 	authMethod, err := ssh.Agent()
 	if err == nil {
 		sshAuthMethods = append(sshAuthMethods, authMethod)
-		logger.Debug("initKnownHostsAuthMethods", "Added SSH Agent (SSH_AUTH_SOCK) to list of auth methods, not adding further methods")
+		dlog.Client.Debug("initKnownHostsAuthMethods", "Added SSH Agent (SSH_AUTH_SOCK)"+
+			"to list of auth methods, not adding further methods")
 		return sshAuthMethods, knownHostsCallback
 	}
-	logger.Debug("initKnownHostsAuthMethods", "Unable to init SSH Agent auth method", err)
+	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to init SSH Agent auth method", err)
 
 	// Third, try Linux/UNIX default key paths
 	privateKeyPath = os.Getenv("HOME") + "/.ssh/id_rsa"
 	authMethod, err = ssh.PrivateKey(privateKeyPath)
 	if err == nil {
 		sshAuthMethods = append(sshAuthMethods, authMethod)
-		logger.Debug("initKnownHostsAuthmethods", "Added path to list of auth methods, not adding further methods", privateKeyPath)
+		dlog.Client.Debug("initKnownHostsAuthmethods", addedPathStr, privateKeyPath)
 		return sshAuthMethods, knownHostsCallback
 	}
-	logger.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
+	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
 
 	privateKeyPath = os.Getenv("HOME") + "/.ssh/id_dsa"
 	authMethod, err = ssh.PrivateKey(privateKeyPath)
 	if err == nil {
 		sshAuthMethods = append(sshAuthMethods, authMethod)
-		logger.Debug("initKnownHostsAuthmethods", "Added path to list of auth methods, not adding further methods", privateKeyPath)
+		dlog.Client.Debug("initKnownHostsAuthmethods", addedPathStr, privateKeyPath)
 		return sshAuthMethods, knownHostsCallback
 	}
-	logger.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
 
 	privateKeyPath = os.Getenv("HOME") + "/.ssh/id_ecdsa"
 	authMethod, err = ssh.PrivateKey(privateKeyPath)
 	if err == nil {
 		sshAuthMethods = append(sshAuthMethods, authMethod)
-		logger.Debug("initKnownHostsAuthmethods", "Added path to list of auth methods, not adding further methods", privateKeyPath)
+		dlog.Client.Debug("initKnownHostsAuthmethods", addedPathStr, privateKeyPath)
 		return sshAuthMethods, knownHostsCallback
 	}
-	logger.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
-
-	logger.FatalExit("Unable to find private SSH key information")
 
+	dlog.Client.FatalPanic("Unable to find private SSH key information", privateKeyPath, err)
 	// Never reach this point.
 	return sshAuthMethods, knownHostsCallback
 }