From 0c7c4f486f5bab53d9c77f91e3e7892e4bd1edb5 Mon Sep 17 00:00:00 2001
From: Eric <xplzv@126.com>
Date: Mon, 2 Sep 2019 12:32:57 +0800
Subject: [PATCH] [update]fix selecting asset bugs

---
 pkg/handler/pagination.go  | 68 ++++++++++++++++++++------------------
 pkg/handler/session.go     | 42 ++++++++++++++---------
 pkg/srvconn/connmanager.go |  2 +-
 3 files changed, 63 insertions(+), 49 deletions(-)

diff --git a/pkg/handler/pagination.go b/pkg/handler/pagination.go
index 109b66c..f7b3174 100644
--- a/pkg/handler/pagination.go
+++ b/pkg/handler/pagination.go
@@ -2,6 +2,7 @@ package handler
 
 import (
 	"fmt"
+	"io"
 	"strconv"
 	"strings"
 
@@ -176,36 +177,31 @@ func (p *AssetPagination) displayPageAssets() {
 }
 
 func (p *AssetPagination) displayTipsInfo() {
-	tips := []string{
-		i18n.T("\nTips: Enter the asset ID and log directly into the asset.\n"),
-		i18n.T("\nPage up: P/p	Page down: Enter|N/n	BACK: b.\n"),
-	}
-	for _, tip := range tips {
-		_, _ = p.term.Write([]byte(tip))
-	}
+	displayAssetPaginationTipsInfo(p.term)
 
 }
 
-func NewUserPagination(term *utils.Terminal, uid, search string, proxy bool) *UserAssetPagination {
+func NewUserPagination(term *utils.Terminal, uid, search string, policy bool) *UserAssetPagination {
 	return &UserAssetPagination{
-		UserID: uid,
-		offset: 0,
-		limit:  0,
-		search: search,
-		term:   term,
-		proxy:  proxy,
-		Data:   model.AssetsPaginationResponse{},
+		UserID:        uid,
+		offset:        0,
+		limit:         0,
+		search:        search,
+		term:          term,
+		displayPolicy: policy,
+		Data:          model.AssetsPaginationResponse{},
 	}
 }
 
 type UserAssetPagination struct {
-	UserID string
-	offset int
-	limit  int
-	search string
-	term   *utils.Terminal
-	proxy  bool
-	Data   model.AssetsPaginationResponse
+	UserID        string
+	offset        int
+	limit         int
+	search        string
+	term          *utils.Terminal
+	displayPolicy bool
+	Data          model.AssetsPaginationResponse
+	IsNeedProxy   bool
 }
 
 func (p *UserAssetPagination) Start() []model.Asset {
@@ -214,7 +210,8 @@ func (p *UserAssetPagination) Start() []model.Asset {
 	for {
 		p.retrieveData()
 
-		if p.proxy && p.Data.Total == 1 {
+		if p.displayPolicy && p.Data.Total == 1 {
+			p.IsNeedProxy = true
 			return p.Data.Data
 		}
 
@@ -251,6 +248,7 @@ func (p *UserAssetPagination) Start() []model.Asset {
 			default:
 				if indexID, err := strconv.Atoi(line); err == nil {
 					if indexID > 0 && indexID <= len(p.Data.Data) {
+						p.IsNeedProxy = true
 						return []model.Asset{p.Data.Data[indexID-1]}
 					}
 				}
@@ -259,6 +257,7 @@ func (p *UserAssetPagination) Start() []model.Asset {
 		default:
 			if indexID, err := strconv.Atoi(line); err == nil {
 				if indexID > 0 && indexID <= len(p.Data.Data) {
+					p.IsNeedProxy = true
 					return []model.Asset{p.Data.Data[indexID-1]}
 				}
 			}
@@ -348,25 +347,18 @@ func (p *UserAssetPagination) displayPageAssets() {
 }
 
 func (p *UserAssetPagination) displayTipsInfo() {
-	tips := []string{
-		i18n.T("\nTips: Enter the asset ID and log directly into the asset.\n"),
-		i18n.T("\nPage up: P/p	Page down: Enter|N/n	BACK: b.\n"),
-	}
-	for _, tip := range tips {
-		_, _ = p.term.Write([]byte(tip))
-	}
-
+	displayAssetPaginationTipsInfo(p.term)
 }
 
 func (p *UserAssetPagination) retrieveData() {
-	p.limit = GetPageSize(p.term)
+	p.limit = getPageSize(p.term)
 	if p.limit == 0 || p.offset < 0 || p.limit >= p.Data.Total {
 		p.offset = 0
 	}
 	p.Data = service.GetUserAssets(p.UserID, p.search, p.limit, p.offset)
 }
 
-func GetPageSize(term *utils.Terminal) int {
+func getPageSize(term *utils.Terminal) int {
 	var (
 		pageSize  int
 		minHeight = 8 // 分页显示的最小高度
@@ -391,3 +383,13 @@ func GetPageSize(term *utils.Terminal) int {
 	}
 	return pageSize
 }
+
+func displayAssetPaginationTipsInfo(w io.Writer) {
+	tips := []string{
+		i18n.T("\nTips: Enter the asset ID and log directly into the asset.\n"),
+		i18n.T("\nPage up: P/p	Page down: Enter|N/n	BACK: b.\n"),
+	}
+	for _, tip := range tips {
+		_, _ = w.Write([]byte(tip))
+	}
+}
diff --git a/pkg/handler/session.go b/pkg/handler/session.go
index 3e52377..e60189e 100644
--- a/pkg/handler/session.go
+++ b/pkg/handler/session.go
@@ -189,9 +189,15 @@ func (h *interactiveHandler) displayAllAssets() {
 		<-h.loadDataDone
 		h.displayAssets(h.allAssets)
 	default:
-		h.searchAsset("")
+		pag := NewUserPagination(h.term, h.user.ID, "", false)
+		result := pag.Start()
+		if pag.IsNeedProxy && len(result) == 1 {
+			h.searchResult = h.searchResult[:0]
+			h.ProxyAsset(result[0])
+		} else {
+			h.searchResult = result
+		}
 	}
-
 }
 
 func (h *interactiveHandler) chooseSystemUser(systemUsers []model.SystemUser) model.SystemUser {
@@ -304,7 +310,13 @@ func (h *interactiveHandler) searchAsset(key string) {
 		h.displayAssets(assets)
 	default:
 		pag := NewUserPagination(h.term, h.user.ID, key, false)
-		h.searchResult = pag.Start()
+		result := pag.Start()
+		if pag.IsNeedProxy && len(result) == 1 {
+			h.searchResult = h.searchResult[:0]
+			h.ProxyAsset(result[0])
+		} else {
+			h.searchResult = result
+		}
 	}
 }
 
@@ -312,11 +324,7 @@ func (h *interactiveHandler) searchAssetOrProxy(key string) {
 	if indexNum, err := strconv.Atoi(key); err == nil && len(h.searchResult) > 0 {
 		if indexNum > 0 && indexNum <= len(h.searchResult) {
 			assetSelect := h.searchResult[indexNum-1]
-			systemUsers := service.GetUserAssetSystemUsers(h.user.ID, assetSelect.ID)
-			systemUserSelect := h.chooseSystemUser(systemUsers)
-			h.systemUserSelect = &systemUserSelect
-			h.assetSelect = &assetSelect
-			h.Proxy(context.Background())
+			h.ProxyAsset(assetSelect)
 			return
 		}
 	}
@@ -342,14 +350,10 @@ func (h *interactiveHandler) searchAssetOrProxy(key string) {
 	}
 
 	if len(assets) == 1 {
-		systemUsers := service.GetUserAssetSystemUsers(h.user.ID, assets[0].ID)
-		systemUserSelect := h.chooseSystemUser(systemUsers)
-		h.systemUserSelect = &systemUserSelect
-		h.assetSelect = &assets[0]
-		h.Proxy(context.Background())
-		return
+		h.ProxyAsset(assets[0])
+	} else {
+		h.searchResult = assets
 	}
-	h.searchResult = assets
 }
 
 func (h *interactiveHandler) searchNodeAssets(num int) (assets model.AssetList) {
@@ -361,6 +365,14 @@ func (h *interactiveHandler) searchNodeAssets(num int) (assets model.AssetList)
 	return
 }
 
+func (h *interactiveHandler) ProxyAsset(assetSelect model.Asset) {
+	systemUsers := service.GetUserAssetSystemUsers(h.user.ID, assetSelect.ID)
+	systemUserSelect := h.chooseSystemUser(systemUsers)
+	h.systemUserSelect = &systemUserSelect
+	h.assetSelect = &assetSelect
+	h.Proxy(context.Background())
+}
+
 func (h *interactiveHandler) Proxy(ctx context.Context) {
 	p := proxy.ProxyServer{
 		UserConn:   h.sess,
diff --git a/pkg/srvconn/connmanager.go b/pkg/srvconn/connmanager.go
index 9dd3f25..fa10b5e 100644
--- a/pkg/srvconn/connmanager.go
+++ b/pkg/srvconn/connmanager.go
@@ -90,7 +90,7 @@ func KeepAlive(c *gossh.Client, closed <-chan struct{}, keepInterval time.Durati
 		case <-t.C:
 			_, _, err := c.SendRequest("keepalive@jumpserver.org", true, nil)
 			if err != nil {
-				logger.Error("SSH client %p keep alive err: ", c, err.Error())
+				logger.Errorf("SSH client %p keep alive err: ", c, err.Error())
 				return
 			}
 		}
-- 
2.18.0