abra/pkg/ssh/ssh.go

package ssh

import (
	"os/user"
	"time"

	"github.com/AlecAivazis/survey/v2"
	"github.com/kevinburke/ssh_config"
	"github.com/sfreiberg/simplessh"
	"github.com/sirupsen/logrus"
)

// HostConfig is a SSH host config.
type HostConfig struct {
	Host         string
	IdentityFile string
	Port         string
	User         string
}

// GetHostConfig retrieves a ~/.ssh/config config for a host.
func GetHostConfig(hostname, username, port string) (HostConfig, error) {
	var hostConfig HostConfig

	var host, idf string

	if host = ssh_config.Get(hostname, "Hostname"); host == "" {
		logrus.Debugf("no hostname found in SSH config, assuming %s", hostname)
		host = hostname
	}

	if username == "" {
		if username = ssh_config.Get(hostname, "User"); username == "" {
			systemUser, err := user.Current()
			if err != nil {
				return hostConfig, err
			}
			logrus.Debugf("no username found in SSH config or passed on command-line, assuming %s", username)
			username = systemUser.Username
		}
	}

	if port == "" {
		if port = ssh_config.Get(hostname, "Port"); port == "" {
			logrus.Debugf("no port found in SSH config or passed on command-line, assuming 22")
			port = "22"
		}
	}

	idf = ssh_config.Get(hostname, "IdentityFile")

	hostConfig.Host = host
	if idf != "" {
		hostConfig.IdentityFile = idf
	}
	hostConfig.Port = port
	hostConfig.User = username

	logrus.Debugf("constructed SSH config %s for %s", hostConfig, hostname)

	return hostConfig, nil
}

// New creates a new SSH client connection.
func New(domainName, sshAuth, username, port string) (*simplessh.Client, error) {
	var client *simplessh.Client

	hostConfig, err := GetHostConfig(domainName, username, port)
	if err != nil {
		return client, err
	}

	if sshAuth == "identity-file" {
		var err error
		client, err = simplessh.ConnectWithAgentTimeout(hostConfig.Host, hostConfig.User, 5*time.Second)
		if err != nil {
			return client, err
		}
	} else {
		password := ""
		prompt := &survey.Password{
			Message: "SSH password?",
		}
		if err := survey.AskOne(prompt, &password); err != nil {
			return client, err
		}

		var err error
		client, err = simplessh.ConnectWithPasswordTimeout(hostConfig.Host, hostConfig.User, password, 5*time.Second)
		if err != nil {
			return client, err
		}
	}

	return client, nil
}