can specify a private ssh key via command line, no password support yet though

9 files changed, 101 insertions, 59 deletions

diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go
index 42b29bf..1ab67c2 100644
--- a/cmd/dcat/main.go
+++ b/cmd/dcat/main.go
@@ -26,6 +26,7 @@ func main() {
 	var quietEnable bool
 	var sshPort int
 	var trustAllHosts bool
+	var privateKeyPathFile string
 
 	userName := user.Name()
 
@@ -41,6 +42,7 @@ func main() {
 	flag.StringVar(&files, "files", "", "File(s) to read")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
+	flag.StringVar(&privateKeyPathFile, "key", "", "Path to private key")
 
 	flag.Parse()
 
@@ -55,12 +57,13 @@ func main() {
 	logger.Start(ctx, logger.Modes{Debug: debugEnable || config.Common.DebugEnable, Quiet: quietEnable})
 
 	args := clients.Args{
-		ConnectionsPerCPU: connectionsPerCPU,
-		ServersStr:        serversStr,
-		Discovery:         discovery,
-		UserName:          userName,
-		What:              files,
-		TrustAllHosts:     trustAllHosts,
+		ConnectionsPerCPU:  connectionsPerCPU,
+		ServersStr:         serversStr,
+		Discovery:          discovery,
+		UserName:           userName,
+		What:               files,
+		TrustAllHosts:      trustAllHosts,
+		PrivateKeyPathFile: privateKeyPathFile,
 	}
 
 	client, err := clients.NewCatClient(args)
diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index c136b43..faffe76 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -27,6 +27,7 @@ func main() {
 	var quietEnable bool
 	var sshPort int
 	var trustAllHosts bool
+	var privateKeyPathFile string
 
 	userName := user.Name()
 
@@ -43,6 +44,7 @@ func main() {
 	flag.StringVar(&regex, "regex", ".", "Regular expression")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
+	flag.StringVar(&privateKeyPathFile, "key", "", "Path to private key")
 
 	flag.Parse()
 
@@ -57,13 +59,14 @@ func main() {
 	logger.Start(ctx, logger.Modes{Debug: debugEnable || config.Common.DebugEnable, Quiet: quietEnable})
 
 	args := clients.Args{
-		ConnectionsPerCPU: connectionsPerCPU,
-		ServersStr:        serversStr,
-		Discovery:         discovery,
-		UserName:          userName,
-		What:              files,
-		TrustAllHosts:     trustAllHosts,
-		Regex:             regex,
+		ConnectionsPerCPU:  connectionsPerCPU,
+		ServersStr:         serversStr,
+		Discovery:          discovery,
+		UserName:           userName,
+		What:               files,
+		TrustAllHosts:      trustAllHosts,
+		Regex:              regex,
+		PrivateKeyPathFile: privateKeyPathFile,
 	}
 
 	client, err := clients.NewGrepClient(args)
diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index cae1552..d7b5fae 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -29,6 +29,7 @@ func main() {
 	var sshPort int
 	var timeout int
 	var trustAllHosts bool
+	var privateKeyPathFile string
 
 	userName := user.Name()
 
@@ -46,6 +47,7 @@ func main() {
 	flag.StringVar(&queryStr, "query", "", "Map reduce query")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
+	flag.StringVar(&privateKeyPathFile, "key", "", "Path to private key")
 
 	flag.Parse()
 
@@ -60,14 +62,15 @@ func main() {
 	logger.Start(ctx, logger.Modes{Debug: debugEnable || config.Common.DebugEnable, Quiet: quietEnable})
 
 	args := clients.Args{
-		ConnectionsPerCPU: connectionsPerCPU,
-		ServersStr:        serversStr,
-		Discovery:         discovery,
-		UserName:          userName,
-		What:              files,
-		TrustAllHosts:     trustAllHosts,
-		Mode:              omode.MapClient,
-		Timeout:           timeout,
+		ConnectionsPerCPU:  connectionsPerCPU,
+		ServersStr:         serversStr,
+		Discovery:          discovery,
+		UserName:           userName,
+		What:               files,
+		TrustAllHosts:      trustAllHosts,
+		Mode:               omode.MapClient,
+		Timeout:            timeout,
+		PrivateKeyPathFile: privateKeyPathFile,
 	}
 
 	client, err := clients.NewMaprClient(args, queryStr)
diff --git a/cmd/drun/main.go b/cmd/drun/main.go
index 039c104..8c78e8f 100644
--- a/cmd/drun/main.go
+++ b/cmd/drun/main.go
@@ -31,6 +31,7 @@ func main() {
 	var sshPort int
 	var timeout int
 	var trustAllHosts bool
+	var privateKeyPathFile string
 
 	userName := user.Name()
 
@@ -49,6 +50,7 @@ func main() {
 	flag.StringVar(&jobName, "name", "", "The job name (if run in background)")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
+	flag.StringVar(&privateKeyPathFile, "key", "", "Path to private key")
 
 	flag.Parse()
 
@@ -64,14 +66,15 @@ func main() {
 
 	command, commandArgs := readCommand(command)
 	args := clients.Args{
-		ConnectionsPerCPU: connectionsPerCPU,
-		ServersStr:        serversStr,
-		Discovery:         discovery,
-		UserName:          userName,
-		What:              command,
-		Arguments:         commandArgs,
-		TrustAllHosts:     trustAllHosts,
-		Timeout:           timeout,
+		ConnectionsPerCPU:  connectionsPerCPU,
+		ServersStr:         serversStr,
+		Discovery:          discovery,
+		UserName:           userName,
+		What:               command,
+		Arguments:          commandArgs,
+		TrustAllHosts:      trustAllHosts,
+		Timeout:            timeout,
+		PrivateKeyPathFile: privateKeyPathFile,
 	}
 
 	client, err := clients.NewRunClient(args, background, jobName)
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index e69e85a..699784b 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -36,6 +36,7 @@ func main() {
 	var sshPort int
 	var timeout int
 	var trustAllHosts bool
+	var privateKeyPathFile string
 
 	userName := user.Name()
 
@@ -56,6 +57,7 @@ func main() {
 	flag.StringVar(&regex, "regex", ".", "Regular expression")
 	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&userName, "user", userName, "Your system user name")
+	flag.StringVar(&privateKeyPathFile, "key", "", "Path to private key")
 
 	flag.Parse()
 
@@ -83,15 +85,16 @@ func main() {
 	}
 
 	args := clients.Args{
-		ConnectionsPerCPU: connectionsPerCPU,
-		ServersStr:        serversStr,
-		Discovery:         discovery,
-		UserName:          userName,
-		What:              files,
-		TrustAllHosts:     trustAllHosts,
-		Regex:             regex,
-		Mode:              omode.TailClient,
-		Timeout:           timeout,
+		ConnectionsPerCPU:  connectionsPerCPU,
+		ServersStr:         serversStr,
+		Discovery:          discovery,
+		UserName:           userName,
+		What:               files,
+		TrustAllHosts:      trustAllHosts,
+		Regex:              regex,
+		Mode:               omode.TailClient,
+		Timeout:            timeout,
+		PrivateKeyPathFile: privateKeyPathFile,
 	}
 
 	var client clients.Client
diff --git a/internal/clients/args.go b/internal/clients/args.go
index b4852d4..1145a4b 100644
--- a/internal/clients/args.go
+++ b/internal/clients/args.go
@@ -20,4 +20,5 @@ type Args struct {
 	Timeout            int
 	SSHAuthMethods     []gossh.AuthMethod
 	SSHHostKeyCallback gossh.HostKeyCallback
+	PrivateKeyPathFile string
 }
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index 10a5559..1c75205 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -40,7 +40,7 @@ func (c *baseClient) init(maker maker) {
 	logger.Info("Initiating base client")
 
 	c.maker = maker
-	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts, c.throttleCh)
+	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts, c.throttleCh, c.Args.PrivateKeyPathFile)
 
 	discoveryService := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle)
 
diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go
index a310159..2ff80b2 100644
--- a/internal/ssh/client/authmethods.go
+++ b/internal/ssh/client/authmethods.go
@@ -11,7 +11,7 @@ import (
 )
 
 // InitSSHAuthMethods initialises all known SSH auth methods on the client side.
-func InitSSHAuthMethods(sshAuthMethods []gossh.AuthMethod, hostKeyCallback gossh.HostKeyCallback, trustAllHosts bool, throttleCh chan struct{}) ([]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 {
@@ -20,39 +20,65 @@ func InitSSHAuthMethods(sshAuthMethods []gossh.AuthMethod, hostKeyCallback gossh
 		return sshAuthMethods, simpleCallback
 	}
 
-	return initKnownHostsAuthMethods(trustAllHosts, throttleCh)
+	return initKnownHostsAuthMethods(trustAllHosts, throttleCh, privateKeyPath)
 }
 
-func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{}) ([]gossh.AuthMethod, HostKeyCallback) {
+func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{}, privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
 	var sshAuthMethods []gossh.AuthMethod
+
+	knownHostsPath := os.Getenv("HOME") + "/.ssh/known_hosts"
+	knownHostsCallback, err := NewKnownHostsCallback(knownHostsPath, trustAllHosts, throttleCh)
+	if err != nil {
+		logger.FatalExit(knownHostsPath, err)
+	}
+	logger.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsPath)
+
 	if config.Common.ExperimentalFeaturesEnable {
 		sshAuthMethods = append(sshAuthMethods, gossh.Password("experimental feature test"))
-		logger.Debug("Added experimental method to list of auth methods")
+		logger.Debug("initKnownHostsAuthMethods", "Added experimental method to list of auth methods")
 	}
 
-	keyPath := os.Getenv("HOME") + "/.ssh/id_rsa"
-	if authMethod, err := ssh.PrivateKey(keyPath); err == nil {
-		sshAuthMethods = append(sshAuthMethods, authMethod)
-		logger.Debug("Added path to list of auth methods", keyPath)
+	// First 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)
+			return sshAuthMethods, knownHostsCallback
+		}
+		logger.FatalExit("Unable to use private SSH key", privateKeyPath, err)
 	}
 
-	keyPath = os.Getenv("HOME") + "/.ssh/id_dsa"
-	if authMethod, err := ssh.PrivateKey(keyPath); err == nil {
+	// Second, try SSH Agent
+	authMethod, err := ssh.Agent()
+	if err == nil {
 		sshAuthMethods = append(sshAuthMethods, authMethod)
-		logger.Debug("Added path to list of auth methods", keyPath)
+		logger.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)
 
-	if authMethod, err := ssh.Agent(); err == nil {
+	// 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("Added SSH Agent to list of auth methods")
+		logger.Debug("initKnownHostsAuthmethods", "Added path to list of auth methods, not adding further methods", privateKeyPath)
+		return sshAuthMethods, knownHostsCallback
 	}
+	logger.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
 
-	knownHostsPath := os.Getenv("HOME") + "/.ssh/known_hosts"
-	knownHostsCallback, err := NewKnownHostsCallback(knownHostsPath, trustAllHosts, throttleCh)
-	if err != nil {
-		logger.FatalExit(knownHostsPath, 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)
+		return sshAuthMethods, knownHostsCallback
 	}
-	logger.Debug("Added known hosts file path", knownHostsPath)
+	logger.Debug("initKnownHostsAuthMethods", "Unable to use private key", privateKeyPath, err)
+
+	logger.FatalExit("Unable to find private SSH key information")
 
+	// Never reach this point.
 	return sshAuthMethods, knownHostsCallback
 }
diff --git a/internal/version/version.go b/internal/version/version.go
index 6e1940c..07f75ec 100644
--- a/internal/version/version.go
+++ b/internal/version/version.go
@@ -11,9 +11,9 @@ const (
 	// Name of DTail.
 	Name string = "DTail"
 	// Version of DTail.
-	Version string = "2.2.1"
+	Version string = "2.2.2"
 	// Additional information for DTail
-	Additional string = ""
+	Additional string = "develop"
 	// ProtocolCompat -ibility version.
 	ProtocolCompat string = "2"
 )