From 35f43ca7c236ad55e22df3785640e6d0ec40ec22 Mon Sep 17 00:00:00 2001 From: Spencer Cai Date: Thu, 11 Nov 2021 11:33:59 +0800 Subject: [PATCH] make ssh actually work & better display for hosts (#6) Signed-off-by: spencercjh --- cmd/sshctx/fzf.go | 9 +- cmd/sshctx/list.go | 1 + cmd/sshctx/previous.go | 2 +- cmd/sshctx/switch.go | 114 +++++++++++++++++-------- go.mod | 7 +- internal/env/constants.go | 2 +- internal/sshconfig/sshconfig.go | 36 +++++--- internal/sshconfig/sshconfig_loader.go | 6 +- internal/sshconfig/sshconfig_test.go | 6 +- test/blank_config.yaml | 1 + test/config_example.yaml | 1 + 11 files changed, 121 insertions(+), 64 deletions(-) diff --git a/cmd/sshctx/fzf.go b/cmd/sshctx/fzf.go index ec1d12fc..8b74f092 100644 --- a/cmd/sshctx/fzf.go +++ b/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 } diff --git a/cmd/sshctx/list.go b/cmd/sshctx/list.go index 6714f2ac..1ebce1ae 100644 --- a/cmd/sshctx/list.go +++ b/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) } diff --git a/cmd/sshctx/previous.go b/cmd/sshctx/previous.go index ed4d8260..0139e071 100644 --- a/cmd/sshctx/previous.go +++ b/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 } diff --git a/cmd/sshctx/switch.go b/cmd/sshctx/switch.go index 6f4ed2fe..4921447c 100644 --- a/cmd/sshctx/switch.go +++ b/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,8 +26,8 @@ import ( "os" "os/exec" "strconv" - - "github.com/pkg/errors" + "strings" + "sync" ) // SwitchOp indicates intention to switch contexts. @@ -34,61 +35,104 @@ 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) { @@ -96,12 +140,12 @@ func connectPrevious(stderr io.Writer) (string, error) { }(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) } diff --git a/go.mod b/go.mod index 75090477..1bf6cb6d 100644 --- a/go.mod +++ b/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 ) diff --git a/internal/env/constants.go b/internal/env/constants.go index db410495..f2d08951 100644 --- a/internal/env/constants.go +++ b/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 diff --git a/internal/sshconfig/sshconfig.go b/internal/sshconfig/sshconfig.go index 9a566d8c..01f85387 100644 --- a/internal/sshconfig/sshconfig.go +++ b/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") diff --git a/internal/sshconfig/sshconfig_loader.go b/internal/sshconfig/sshconfig_loader.go index 178afaa8..6ac8b605 100644 --- a/internal/sshconfig/sshconfig_loader.go +++ b/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 diff --git a/internal/sshconfig/sshconfig_test.go b/internal/sshconfig/sshconfig_test.go index 88b62ff5..9eae38a9 100644 --- a/internal/sshconfig/sshconfig_test.go +++ b/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) { diff --git a/test/blank_config.yaml b/test/blank_config.yaml index 5d734063..ef953e71 100644 --- a/test/blank_config.yaml +++ b/test/blank_config.yaml @@ -2,3 +2,4 @@ previous: host: username: port: + displayname: diff --git a/test/config_example.yaml b/test/config_example.yaml index 24ac03d8..34803b90 100644 --- a/test/config_example.yaml +++ b/test/config_example.yaml @@ -1,4 +1,5 @@ previous: host: 10.115.40.97 username: root + displayname: test port: 22