make ssh actually work & better display for hosts (#6)

Signed-off-by: spencercjh <shouspencercjh@foxmail.com>

cmd/sshctx/fzf.go

@@ -18,7 +18,6 @@ import (
 	"bytes"
 	"fmt"
 	"github.com/spencercjh/sshctx/internal/env"
-	"github.com/spencercjh/sshctx/internal/printer"
 	"github.com/spencercjh/sshctx/internal/sshconfig"
 	"io"
 	"os"
@@ -32,7 +31,7 @@ type InteractiveSwitchOp struct {
 	SelfCmd string
 }
 
-func (op InteractiveSwitchOp) Run(_, stderr io.Writer) error {
+func (op InteractiveSwitchOp) Run(stdout, stderr io.Writer) error {
 	// parse sshconfig just to see if it can be loaded
 	sc := new(sshconfig.SSHConfig).WithLoader(sshconfig.DefaultLoader)
 	defer func(sshConfig *sshconfig.SSHConfig) {
@@ -62,10 +61,12 @@ func (op InteractiveSwitchOp) Run(_, stderr io.Writer) error {
 	if choice == "" {
 		return errors.New("you did not choose any of the options")
 	}
-	host, err := connectTarget(choice, stderr)
+	displayName, sshPara, err := connectTarget(choice, stderr)
 	if err != nil {
 		return errors.Wrap(err, "failed to switch host")
 	}
-	_ = printer.Success(stderr, "Switched to host \"%s\".", printer.SuccessColor.Sprint(host))
+	if err := savePreviousHost(stdout, displayName, sshPara); err != nil {
+		return errors.Wrap(err, "failed to save previous host")
+	}
 	return nil
 }

cmd/sshctx/list.go

@@ -51,6 +51,7 @@ func (op ListOp) Run(stdout, _ io.Writer) error {
 			_ = printer.Warning(stdout, "%s is an illegal ssh parameter", str)
 			continue
 		}
+		str = "💻: " + h.DisplayName + "#" + str
 		if h == sc.PreviousHost {
 			str = printer.ActiveItemColor.Sprint(str)
 		}

cmd/sshctx/previous.go

@@ -37,6 +37,6 @@ func (op PreviousOp) Run(stdout, _ io.Writer) error {
 	if sc.PreviousHost == sshconfig.EmptyHost {
 		return errors.New("No previous host in sshctx")
 	}
-	_ = printer.Success(stdout, "Previous host: %s", sc.PreviousHost.ToSSHParameter())
+	_ = printer.Success(stdout, "Previous host: %s#%s", sc.PreviousHost.DisplayName, sc.PreviousHost.ToSSHParameter())
 	return nil
 }

cmd/sshctx/switch.go

@@ -15,7 +15,8 @@
 package main
 
 import (
-	"bytes"
+	"fmt"
+	"github.com/pkg/errors"
 	"github.com/spencercjh/sshctx/internal/env"
 	"github.com/spencercjh/sshctx/internal/printer"
 	"github.com/spencercjh/sshctx/internal/sshconfig"
@@ -25,83 +26,126 @@ import (
 	"os"
 	"os/exec"
 	"strconv"
-
-	"github.com/pkg/errors"
+	"strings"
+	"sync"
 )
 
 // SwitchOp indicates intention to switch contexts.
 type SwitchOp struct {
 	Target string // '-' for back and forth, or NAME
 }
 
-func (op SwitchOp) Run(_, stderr io.Writer) error {
-	var target string
+func (op SwitchOp) Run(stdout, stderr io.Writer) error {
+	var displayName string
+	var sshPara string
 	var err error
 	if op.Target == "-" {
-		target, err = connectPrevious(stderr)
+		displayName, sshPara, err = connectPrevious(stderr)
 	} else {
-		target, err = connectTarget(op.Target, stderr)
+		displayName, sshPara, err = connectTarget(op.Target, stderr)
 	}
 	if err != nil {
 		return errors.Wrap(err, "failed to connect host")
 	}
-	// save previous host
-	e := savePreviousHost(target)
-	if e != nil {
-		return errors.Wrap(e, "failed to save previous host")
+	if err := savePreviousHost(stdout, displayName, sshPara); err != nil {
+		return errors.Wrap(err, "failed to save previous host")
 	}
 	return nil
 }
 
-func savePreviousHost(target string) error {
-	matches := env.SSHParameterRegexp.FindStringSubmatch(target)
-	port, _ := strconv.Atoi(matches[2])
-	var host = map[string]sshconfig.Host{"previous": {matches[1], matches[0], port}}
+func deleteEmpty(s []string) []string {
+	var r []string
+	for _, str := range s {
+		if str != "" {
+			r = append(r, str)
+		}
+	}
+	return r
+}
+
+func savePreviousHost(stdin io.Writer, displayName, sshPara string) error {
+	matches := env.SSHParameterRegexp.FindStringSubmatch(sshPara)
+	matches = deleteEmpty(matches)
+	var host map[string]sshconfig.Host
+	switch {
+	case len(matches) == 5:
+		port, _ := strconv.Atoi(matches[4])
+		host = map[string]sshconfig.Host{"previous": {Host: matches[2], DisplayName: displayName, Username: matches[1], Port: port}}
+	case len(matches) == 3:
+		host = map[string]sshconfig.Host{"previous": {Host: matches[2], DisplayName: displayName, Username: matches[1]}}
+	default:
+		return fmt.Errorf("illegal SSH parameter: %s", sshPara)
+	}
+
 	data, err := yaml.Marshal(&host)
 
 	if err != nil {
-		return errors.Wrap(err, "failed to marshal host")
+		return errors.Wrap(err, fmt.Sprintf("failed to marshal host: %v", host))
 	}
 
 	sshCtxDataPath, _ := sshconfig.GetSSHCtxDataPath()
-	if err := ioutil.WriteFile(sshCtxDataPath, data, 0); err != nil {
+	if err := ioutil.WriteFile(sshCtxDataPath, data, 0666); err != nil {
 		return errors.Wrap(err, "failed to write host to sshctxData file")
 	}
+	_ = printer.Success(stdin, "Saved previous host successfully: %v", host)
 	return nil
 }
 
-// connectTarget switches to specified context name.
-func connectTarget(target string, stderr io.Writer) (string, error) {
-	_ = printer.Success(stderr, "Switched to target \"%s\".", printer.SuccessColor.Sprint(target))
-	cmd := exec.Command("ssh", target)
-	var out bytes.Buffer
-	cmd.Stdin = os.Stdin
-	cmd.Stderr = stderr
-	cmd.Stdout = &out
-	if err := cmd.Run(); err != nil {
-		var exitError *exec.ExitError
-		if ok := errors.Is(err, exitError); !ok {
-			return target, err
-		}
+// extract
+func extract(target string) (string, string, error) {
+	sshParaBeginIndex := strings.IndexAny(target, "#")
+	if sshParaBeginIndex == -1 {
+		return "", "", errors.New("invalid target")
+	}
+	displayName := target[:sshParaBeginIndex]
+	sshPara := target[sshParaBeginIndex+1:]
+	return displayName, sshPara, nil
+}
+
+// connectTarget
+func connectTarget(target string, stderr io.Writer) (string, string, error) {
+	displayName, sshPara, err := extract(target)
+	if err != nil {
+		return "", "", err
 	}
-	return target, nil
+	return connectTargetWithDisplayName(displayName, sshPara, stderr)
+}
+
+// connectTargetWithDisplayName
+func connectTargetWithDisplayName(displayName string, sshPara string, stderr io.Writer) (string, string, error) {
+	_ = printer.Success(stderr, "Switched to target %s.", printer.SuccessColor.Sprint(displayName))
+
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(1)
+	go func() {
+		cmd := exec.Command("ssh", "-t", "-t", sshPara)
+		cmd.Stdin = os.Stdin
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = stderr
+		if err := cmd.Run(); err != nil {
+			_ = printer.Error(stderr, "Failed to connect to target %s because: %v.", printer.ErrorColor.Sprint(displayName), err)
+		}
+		waitGroup.Done()
+	}()
+	waitGroup.Wait()
+	return displayName, sshPara, nil
 }
 
 // connectPrevious switches to previously switch context.
-func connectPrevious(stderr io.Writer) (string, error) {
+func connectPrevious(stderr io.Writer) (string, string, error) {
 	sc := new(sshconfig.SSHConfig).WithLoader(sshconfig.DefaultLoader)
 
 	defer func(sshConfig *sshconfig.SSHConfig) {
 		_ = sshConfig.Close()
 	}(sc)
 
 	if err := sc.Parse(); err != nil {
-		return "", errors.Wrap(err, "sshconfig error")
+		return "", "", errors.Wrap(err, "sshconfig error")
 	}
 
 	if sc.PreviousHost == sshconfig.EmptyHost {
-		return "", errors.New("No previous host")
+		return "", "", errors.New("No previous host")
 	}
 
-	return connectTarget(sc.PreviousHost.ToSSHParameter(), stderr)
+	return connectTargetWithDisplayName(sc.PreviousHost.DisplayName, sc.PreviousHost.ToSSHParameter(), stderr)
 }

go.mod

@@ -6,12 +6,9 @@ require (
 	facette.io/natsort v0.0.0-20181210072756-2cd4dd1e2dcb
 	github.com/fatih/color v1.9.0
 	github.com/google/go-cmp v0.5.6
+	github.com/mattn/go-colorable v0.1.4 // indirect
 	github.com/mattn/go-isatty v0.0.12
 	github.com/pkg/errors v0.9.1
-	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c
-)
-
-require (
-	github.com/mattn/go-colorable v0.1.4 // indirect
 	golang.org/x/sys v0.0.0-20200116001909-b77594299b42 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c
 )

internal/env/constants.go

@@ -16,7 +16,7 @@ package env
 
 import "regexp"
 
-var SSHParameterRegexp = regexp.MustCompile(`(?m)^(\w+)@((?:1[0-9][0-9]\.|2[0-4][0-9]\.|25[0-5]\.|[1-9][0-9]\.|[0-9]\.){3}(?:1[0-9][0-9]|2[0-4][0-9]|25[0-5]|[1-9][0-9]|[0-9])+|\w+[^\s]+\.[^\s]+)+:(\d+)$`)
+var SSHParameterRegexp = regexp.MustCompile(`(?m)^(\w+)@((?:1[0-9][0-9]\.|2[0-4][0-9]\.|25[0-5]\.|[1-9][0-9]\.|[0-9]\.){3}(?:1[0-9][0-9]|2[0-4][0-9]|25[0-5]|[1-9][0-9]|[0-9])+|\w+[^\s]+\.[^\s]+)+( -p (\d+))?$`)
 
 const (
 	// FZFIgnore describes the environment variable to set to disable

internal/sshconfig/sshconfig.go

@@ -41,13 +41,17 @@ type SSHConfig struct {
 }
 
 type Host struct {
-	Host     string
-	Username string
-	Port     int
+	Host        string
+	DisplayName string
+	Username    string
+	Port        int
 }
 
 func (h *Host) ToSSHParameter() string {
-	return h.Username + "@" + h.Host + ":" + strconv.Itoa(h.Port)
+	if h.Port > 0 {
+		return h.Username + "@" + h.Host + " -p " + strconv.Itoa(h.Port)
+	}
+	return h.Username + "@" + h.Host
 }
 
 var EmptyHost = Host{}
@@ -155,11 +159,11 @@ func extractConfigItem(itemBeginIndex int, itemEndIndex int, rows []string) Host
 	var hostname string
 	for k := itemBeginIndex; k < itemEndIndex; k++ {
 		itemRow := rows[k]
-		if strings.HasPrefix(itemRow, "Host") {
-			host = strings.TrimSpace(itemRow[4:])
+		if strings.HasPrefix(itemRow, "Host ") {
+			host = strings.TrimSpace(itemRow[5:])
 		}
-		if strings.HasPrefix(itemRow, "Hostname") {
-			hostname = strings.TrimSpace(itemRow[8:])
+		if strings.HasPrefix(itemRow, "Hostname ") {
+			hostname = strings.TrimSpace(itemRow[9:])
 		}
 		if strings.HasPrefix(itemRow, "User") {
 			configItem.Username = strings.TrimSpace(itemRow[4:])
@@ -168,10 +172,18 @@ func extractConfigItem(itemBeginIndex int, itemEndIndex int, rows []string) Host
 			configItem.Port, _ = strconv.Atoi(strings.TrimSpace(itemRow[4:]))
 		}
 	}
-	if hostname != "" {
+	switch {
+	case host == "" && hostname != "":
 		configItem.Host = hostname
-	} else if host != "" {
+		configItem.DisplayName = hostname
+	case host != "" && hostname == "":
 		configItem.Host = host
+		configItem.DisplayName = host
+	case host != "" && host != "name" && host != "*" && hostname != "":
+		configItem.Host = hostname
+		configItem.DisplayName = host
+	default:
+		return EmptyHost
 	}
 	// no host and hostname
 	if configItem.Host == "name" ||
@@ -184,9 +196,6 @@ func extractConfigItem(itemBeginIndex int, itemEndIndex int, rows []string) Host
 	if configItem.Username == "" {
 		configItem.Username = os.Getenv("USER")
 	}
-	if configItem.Port == 0 {
-		configItem.Port = 22
-	}
 	return configItem
 }
 
@@ -223,6 +232,7 @@ func previousConfig(rootNode *yaml.Node) (Host, error) {
 	host := Host{}
 	host.Host = valueOf(previous, "host").Value
 	host.Username = valueOf(previous, "username").Value
+	host.DisplayName = valueOf(previous, "displayname").Value
 	port, err := strconv.Atoi(valueOf(previous, "port").Value)
 	if err != nil {
 		return EmptyHost, errors.Wrap(err, "Can't parse port in the previous node")

internal/sshconfig/sshconfig_loader.go

@@ -70,11 +70,13 @@ func (*StandardLoader) LoadSSHCTXData() (io.ReadWriteCloser, error) {
 			}
 			// try to create sshctxData
 			file, err = os.Create(defaultPath)
-			_ = os.Chmod(defaultPath, 0777)
-			// TODO: consider to ignore the error
 			if err != nil {
 				return nil, errors.Wrap(err, fmt.Sprintf("Can't create sshctxData file: %s", defaultPath))
 			}
+			err = os.Chmod(defaultPath, 0777)
+			if err != nil {
+				return nil, errors.Wrap(err, fmt.Sprintf("Can't chmod sshctxData file : %s", defaultPath))
+			}
 		}
 	}
 	return io.ReadWriteCloser(file), nil

internal/sshconfig/sshconfig_test.go

@@ -15,9 +15,9 @@ func TestHost_ToSSHParameter(t *testing.T) {
 		host Host
 		want string
 	}{
-		{name: "ipv4", host: Host{Host: "192.168.1.1", Username: "test", Port: 22}, want: "test@192.168.1.1:22"},
-		{name: "domain", host: Host{Host: "test.com", Username: "test", Port: 22}, want: "test@test.com:22"},
-		{name: "localhost", host: Host{Host: "localhost", Username: "test", Port: 22}, want: "test@localhost:22"},
+		{name: "ipv4", host: Host{Host: "192.168.1.1", Username: "test", Port: 22}, want: "test@192.168.1.1 -p 22"},
+		{name: "domain", host: Host{Host: "test.com", Username: "test", Port: 22}, want: "test@test.com -p 22"},
+		{name: "localhost", host: Host{Host: "localhost", Username: "test", Port: 22}, want: "test@localhost -p 22"},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

test/blank_config.yaml

@@ -2,3 +2,4 @@ previous:
   host:
   username:
   port:
+  displayname:

test/config_example.yaml

@@ -1,4 +1,5 @@
 previous:
   host: 10.115.40.97
   username: root
+  displayname: test
   port: 22