From 3207fb8cc7ba5fbf2e5cc09a4feb02826ff42e3d Mon Sep 17 00:00:00 2001
From: Jason-ZW <zhenyang@rancher.com>
Date: Mon, 28 Sep 2020 13:30:05 +0800
Subject: [PATCH] fix(autok3s): fix ssh terminal bug   1. add tab completion  
 2. fix terminal select logic bug

Signed-off-by: Jason-ZW <zhenyang@rancher.com>
---
 pkg/hosts/tunnel.go              | 45 +++++++++++++++++---------------
 pkg/providers/alibaba/alibaba.go | 33 ++++++++++++++++++++---
 pkg/utils/util.go                | 16 +++---------
 3 files changed, 58 insertions(+), 36 deletions(-)

diff --git a/pkg/hosts/tunnel.go b/pkg/hosts/tunnel.go
index b680ccc3..805fa4a0 100644
--- a/pkg/hosts/tunnel.go
+++ b/pkg/hosts/tunnel.go
@@ -6,6 +6,7 @@ import (
 	"os"
 
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/terminal"
 )
 
 type Tunnel struct {
@@ -41,40 +42,42 @@ func (t *Tunnel) Cmd(cmd string) *Tunnel {
 
 func (t *Tunnel) Terminal() error {
 	session, err := t.conn.NewSession()
-	if err != nil {
-		return err
-	}
-
 	defer func() {
 		_ = session.Close()
 	}()
-
-	if t.Stdin == nil {
-		session.Stdin = os.Stdin
-	} else {
-		session.Stdin = t.Stdin
-	}
-	if t.Stdout == nil {
-		session.Stdout = os.Stdout
-	} else {
-		session.Stdout = t.Stdout
-	}
-	if t.Stderr == nil {
-		session.Stderr = os.Stderr
-	} else {
-		session.Stderr = t.Stderr
+	if err != nil {
+		return err
 	}
 
 	term := os.Getenv("TERM")
 	if term == "" {
 		t.Term = "xterm-256color"
 	}
-	t.Height = 40
-	t.Weight = 80
 	t.Modes = ssh.TerminalModes{
+		ssh.ECHO:          1,
+		ssh.VSTATUS:       1,
 		ssh.TTY_OP_ISPEED: 14400,
 		ssh.TTY_OP_OSPEED: 14400,
 	}
+
+	fd := int(os.Stdin.Fd())
+	oldState, err := terminal.MakeRaw(fd)
+	defer func() {
+		_ = terminal.Restore(fd, oldState)
+	}()
+	if err != nil {
+		return err
+	}
+
+	t.Weight, t.Height, err = terminal.GetSize(fd)
+	if err != nil {
+		return err
+	}
+
+	session.Stdin = os.Stdin
+	session.Stdout = os.Stdout
+	session.Stderr = os.Stderr
+
 	if err := session.RequestPty(t.Term, t.Height, t.Weight, t.Modes); err != nil {
 		return err
 	}
diff --git a/pkg/providers/alibaba/alibaba.go b/pkg/providers/alibaba/alibaba.go
index 2566c22d..f8436d58 100644
--- a/pkg/providers/alibaba/alibaba.go
+++ b/pkg/providers/alibaba/alibaba.go
@@ -350,13 +350,40 @@ func (p *Alibaba) SSHK3sNode() error {
 	ids := make(map[string]string, len(response.Instances.Instance))
 	for _, instance := range response.Instances.Instance {
 		if instance.EipAddress.IpAddress != "" {
-			ids[instance.InstanceId] = instance.EipAddress.IpAddress
+			for _, t := range instance.Tags.Tag {
+				switch t.TagKey {
+				case "master":
+					if t.TagValue == "true" {
+						ids[instance.InstanceId] = instance.EipAddress.IpAddress + " (master)"
+					}
+				case "worker":
+					if t.TagValue == "true" {
+						ids[instance.InstanceId] = instance.EipAddress.IpAddress + " (worker)"
+					}
+				default:
+					continue
+				}
+			}
 		} else if instance.EipAddress.IpAddress == "" && len(instance.PublicIpAddress.IpAddress) > 0 {
-			ids[instance.InstanceId] = instance.PublicIpAddress.IpAddress[0]
+			for _, t := range instance.Tags.Tag {
+				switch t.TagKey {
+				case "master":
+					if t.TagValue == "true" {
+						ids[instance.InstanceId] = instance.PublicIpAddress.IpAddress[0] + " (master)"
+					}
+				case "worker":
+					if t.TagValue == "true" {
+						ids[instance.InstanceId] = instance.PublicIpAddress.IpAddress[0] + " (worker)"
+					}
+				default:
+					continue
+				}
+			}
+
 		}
 	}
 
-	ip := utils.AskForSelectItem(fmt.Sprintf("[%s] choose ssh node to connect", p.GetProviderName()), ids)
+	ip := strings.Split(utils.AskForSelectItem(fmt.Sprintf("[%s] choose ssh node to connect", p.GetProviderName()), ids), " (")[0]
 
 	if ip == "" {
 		return fmt.Errorf("[%s] choose incorrect ssh node", p.GetProviderName())
diff --git a/pkg/utils/util.go b/pkg/utils/util.go
index 0e3d53a2..d11e6a49 100644
--- a/pkg/utils/util.go
+++ b/pkg/utils/util.go
@@ -7,7 +7,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"os"
-	"strconv"
 	"strings"
 	"text/template"
 
@@ -16,8 +15,8 @@ import (
 
 const (
 	tmpl = `
-{{- range $key, $element := .}}
-  {{$key}}: {{$element}}
+{{- range $key, $value := .}}
+  - {{$key}}: {{$value}}
 {{- end}}
 `
 )
@@ -76,17 +75,10 @@ func AskForSelectItem(s string, ss map[string]string) string {
 	if err := t.Execute(buffer, ss); err != nil {
 		return ""
 	}
-	fmt.Printf("%s: \n \t%s\n[choose 0-%d]:", s, buffer.String(), len(ss))
+	fmt.Printf("%s: \n \t%s\n[choose one id]: ", s, buffer.String())
 	response, err := reader.ReadString('\n')
 	if err != nil {
 		logrus.Fatal(err)
 	}
-	index := 0
-	for _, v := range ss {
-		if strconv.Itoa(index) == strings.ToLower(strings.TrimSpace(response)) {
-			return v
-		}
-		index++
-	}
-	return ""
+	return ss[strings.ToLower(strings.TrimSpace(response))]
 }