[Update] 修改结构

pkg/common/util.go

+package common
+
+import "fmt"

pkg/handler/banner.go

-package handler
-
-import (
-	"text/template"
-
-	"github.com/gliderlabs/ssh"
-
-	"cocogo/pkg/logger"
-)
-
-const bannerTemplate = `
-		{{.UserName}}	Welcome to use Jumpserver open source fortress system{{.EndLine}}
-{{.Tab}}1) Enter {{.ColorCode}}ID{{.ColorEnd}} directly login or enter {{.ColorCode}}part IP, Hostname, Comment{{.ColorEnd}} to search login(if unique). {{.EndLine}}
-{{.Tab}}2) Enter {{.ColorCode}}/{{.ColorEnd}} + {{.ColorCode}}IP, Hostname{{.ColorEnd}} or {{.ColorCode}}Comment{{.ColorEnd}} search, such as: /ip. {{.EndLine}}
-{{.Tab}}3) Enter {{.ColorCode}}p{{.ColorEnd}} to display the host you have permission.{{.EndLine}}
-{{.Tab}}4) Enter {{.ColorCode}}g{{.ColorEnd}} to display the node that you have permission.{{.EndLine}}
-{{.Tab}}5) Enter {{.ColorCode}}g{{.ColorEnd}} + {{.ColorCode}}NodeID{{.ColorEnd}} to display the host under the node, such as g1. {{.EndLine}}
-{{.Tab}}6) Enter {{.ColorCode}}s{{.ColorEnd}} Chinese-english proxy.{{.EndLine}}
-{{.Tab}}7) Enter {{.ColorCode}}h{{.ColorEnd}} help.{{.EndLine}}
-{{.Tab}}8) Enter {{.ColorCode}}r{{.ColorEnd}} to refresh your assets and nodes.{{.EndLine}}
-{{.Tab}}0) Enter {{.ColorCode}}q{{.ColorEnd}} exit.{{.EndLine}}
-`
-
-var displayTemplate = template.Must(template.New("display").Parse(bannerTemplate))
-
-type Banner struct {
-	UserName  string
-	ColorCode string
-	ColorEnd  string
-	Tab       string
-	EndLine   string
-}
-
-func (h *Banner) display(sess ssh.Session) {
-	e := displayTemplate.Execute(sess, h)
-	if e != nil {
-		logger.Warn("Display help info failed")
-	}
-}
-
-func NewBanner(userName string) *Banner {
-	return &Banner{
-		UserName:  userName,
-		ColorCode: GreenColorCode,
-		ColorEnd:  ColorEnd,
-		Tab:       Tab,
-		EndLine:   EndLine,
-	}
-}

pkg/model/assets.go

+package model
+
+import "cocogo/pkg/sdk"
+
+type AssetList []sdk.Asset
+
+func (a *AssetList) SortBy(tp string) AssetList {
+	switch tp {
+	case "ip":
+		return []sdk.Asset{}
+	default:
+		return []sdk.Asset{}
+	}
+}
+
+type NodeList []sdk.Node

pkg/sdk/client.go

@@ -6,7 +6,6 @@ import (
 
 	"cocogo/pkg/common"
 	"cocogo/pkg/config"
-	"cocogo/pkg/model"
 )
 
 type ClientAuth interface {
@@ -34,7 +33,7 @@ func (c *WrapperClient) LoadAuth() error {
 }
 
 func (c *WrapperClient) CheckAuth() error {
-	var user model.User
+	var user User
 	err := c.Http.Get("UserProfileUrl", &user)
 	if err != nil {
 		return err

pkg/service/perms.go

 package service
 
-import "cocogo/pkg/sdk"
+import (
+	"cocogo/pkg/model"
+)
 
-func GetUserAssets(userId string) (assets []sdk.Asset) {
-	return
+func GetUserAssets(userId string) (assets model.AssetList) {
+	return model.AssetList{{Id: "xxxxxxxxx", Hostname: "test", Ip: "192.168.244.185", Port: 22}}
 }
 
-func GetUserNodes(userId string) (nodes []sdk.Node) {
-	return
+func GetUserNodes(userId string) (nodes model.NodeList) {
+	return model.NodeList{{Id: "XXXXXXX", Name: "test"}}
+}
+
+func ValidateUserAssetPermission(userId, assetId, systemUserId string) bool {
+	return true
 }

pkg/sshd/handler/banner.go

+package handler
+
+import (
+	"cocogo/pkg/config"
+	"fmt"
+	"io"
+	"text/template"
+
+	"github.com/gliderlabs/ssh"
+
+	"cocogo/pkg/logger"
+)
+
+const defaultTitle = `Welcome to use Jumpserver open source fortress system`
+
+type MenuItem struct {
+	instruct string
+	helpText string
+}
+
+type Menu []MenuItem
+
+var menu = Menu{
+	{instruct: "ID", helpText: "directly login or enter."},
+	{instruct: "part IP, Hostname, Comment", helpText: "to search login if unique."},
+	{instruct: "/ + IP, Hostname, Comment", helpText: "to search, such as: /192.168"},
+	{instruct: "p", helpText: "display the host you have permission."},
+	{instruct: "g", helpText: "display the node that you have permission."},
+	{instruct: "r", helpText: "refresh your assets and nodes"},
+	{instruct: "s", helpText: "switch Chinese-english language."},
+	{instruct: "h", helpText: "print help"},
+	{instruct: "q", helpText: "exit"},
+}
+
+type Banner struct {
+	user string
+}
+
+type ColorMeta struct {
+	GreenBoldColor string
+	ColorEnd       string
+}
+
+func displayBanner(sess ssh.Session, user string) {
+	title := defaultTitle
+	if config.Conf.HeaderTitle != "" {
+		title = config.Conf.HeaderTitle
+	}
+	welcomeMsg := user + "  " + title
+	_, err := io.WriteString(sess, welcomeMsg)
+	if err != nil {
+		logger.Error("Send to client error, %s", err)
+	}
+	cm := ColorMeta{GreenBoldColor: "\033[1;32m", ColorEnd: "\033[0m"}
+	for i, v := range menu {
+		line := fmt.Sprintf("\t%d) Enter {{.GreenBoldColor}}%s{{.ColorEnd}} to %s.\r\n", i+1, v.instruct, v.helpText)
+		tmpl := template.Must(template.New("item").Parse(line))
+		err := tmpl.Execute(sess, cm)
+		if err != nil {
+			logger.Error("Send to client error, %s", err)
+		}
+	}
+}

pkg/handler/const.go → pkg/sshd/handler/const.go

@@ -2,9 +2,6 @@ package handler
 
 const (
 	GreenColorCode = "\033[32m"
-	ColorEnd       = "\033[0m"
-	Tab            = "\t"
-	EndLine        = "\r\n\r"
 )
 
 const (

pkg/handler/session.go → pkg/sshd/handler/session.go

@@ -4,7 +4,11 @@ import (
 	//"context"
 	//"strconv"
 
+	"cocogo/pkg/model"
+	"cocogo/pkg/proxy"
 	"cocogo/pkg/sdk"
+	"cocogo/pkg/service"
+	"cocogo/pkg/userhome"
 	"context"
 	"fmt"
 	"github.com/olekukonko/tablewriter"
@@ -37,56 +41,65 @@ import (
 )
 
 type InteractiveHandler struct {
-	sess                ssh.Session
-	term                *terminal.Terminal
-	user                sdk.User
-	banner              *Banner
-	currentSearchAssets []sdk.Asset
-	onceLoad            sync.Once
+	sess         ssh.Session
+	term         *terminal.Terminal
+	user         sdk.User
+	assets       model.AssetList
+	searchResult model.AssetList
+	nodes        model.NodeList
+	onceLoad     sync.Once
 	sync.RWMutex
 }
 
 func (i *InteractiveHandler) displayBanner() {
-	i.banner.display(i.sess)
+	displayBanner(i.sess, i.user.Name)
 }
 
-func (i *InteractiveHandler) Dispatch() {
+func (i *InteractiveHandler) preDispatch() {
 	i.displayBanner()
-	_, winCh, _ := i.sess.Pty()
+	i.onceLoad.Do(func() {
+		i.loadUserAssets()
+		i.loadUserAssetNodes()
+	})
+}
+
+func (i *InteractiveHandler) watchSizeChange(winCh <-chan ssh.Window) {
+	ctx, cancelFunc := context.WithCancel(i.sess.Context())
+	defer cancelFunc()
 	for {
-		ctx, cancelFunc := context.WithCancel(i.sess.Context())
-		go func() {
-			for {
-				select {
-				case <-ctx.Done():
-					logger.Info("ctx done")
-					return
-				case win, ok := <-winCh:
-					if !ok {
-						return
-					}
-					logger.Info("Term change:", win)
-					_ = i.term.SetSize(win.Width, win.Height)
-				}
+		select {
+		case <-ctx.Done():
+			logger.Info("ctx done")
+			return
+		case win, ok := <-winCh:
+			if !ok {
+				return
 			}
-		}()
+			logger.Info("Term change:", win)
+			_ = i.term.SetSize(win.Width, win.Height)
+		}
+	}
+}
+
+func (i *InteractiveHandler) Dispatch() {
+	i.preDispatch()
+	_, winCh, _ := i.sess.Pty()
+	for {
+		go i.watchSizeChange(winCh)
+
 		line, err := i.term.ReadLine()
-		cancelFunc()
 		if err != nil {
 			logger.Error("ReadLine done", err)
 			break
 		}
-		if line == "" {
-			continue
-		}
 
-		if len(line) == 1 {
+		switch len(line) {
+		case 0, 1:
 			switch line {
-			case "p", "P":
-				i.displayAssets([]sdk.Asset{})
-
+			case "", "p", "P":
+				i.displayAssets(i.assets)
 			case "g", "G":
-				i.displayNodes([]sdk.Node{})
+				i.displayNodes(i.nodes)
 			case "s", "S":
 				i.changeLanguage()
 			case "h", "H":
@@ -98,41 +111,28 @@ func (i *InteractiveHandler) Dispatch() {
 				return
 			default:
 				assets := i.searchAsset(line)
-				i.currentSearchAssets = assets
+				i.searchResult = assets
 				i.displayAssetsOrProxy(assets)
 			}
-			continue
-		}
-		if strings.Index(line, "/") == 0 {
-			searchWord := strings.TrimSpace(strings.TrimPrefix(line, "/"))
-			assets := i.searchAsset(searchWord)
-			i.currentSearchAssets = assets
-			i.displayAssets(assets)
-			continue
-		}
-
-		if strings.Index(line, "g") == 0 {
-			searchWord := strings.TrimSpace(strings.TrimPrefix(line, "g"))
-			if num, err := strconv.Atoi(searchWord); err == nil {
-				if num >= 0 {
-					assets := i.searchNodeAssets(num)
-					i.displayAssets(assets)
-					i.currentSearchAssets = assets
-					continue
+		default:
+			switch {
+			case strings.Index(line, "/") == 0:
+				searchWord := strings.TrimSpace(line[1:])
+				assets := i.searchAsset(searchWord)
+				i.searchResult = assets
+				i.displayAssets(assets)
+			case strings.Index(line, "g") == 0:
+				searchWord := strings.TrimSpace(strings.TrimPrefix(line, "g"))
+				if num, err := strconv.Atoi(searchWord); err == nil {
+					if num >= 0 {
+						assets := i.searchNodeAssets(num)
+						i.displayAssets(assets)
+						i.searchResult = assets
+						continue
+					}
 				}
 			}
 		}
-
-		if strings.Index(line, "join") == 0 {
-			roomID := strings.TrimSpace(strings.TrimPrefix(line, "join"))
-			i.JoinShareRoom(roomID)
-			continue
-		}
-
-		assets := i.searchAsset(line)
-		i.currentSearchAssets = assets
-		i.displayAssetsOrProxy(assets)
-
 	}
 }
 
@@ -180,7 +180,7 @@ func (i *InteractiveHandler) displayAssetsOrProxy(assets []sdk.Asset) {
 		if err != nil {
 			return
 		}
-		if ok := sdk.ValidateUserAssetPermission(i.user.Id, systemUser.Id, assets[0].Id); !ok {
+		if ok := service.ValidateUserAssetPermission(i.user.Id, systemUser.Id, assets[0].Id); !ok {
 			// 检查user 是否对该资产有权限
 			return
 		}
@@ -195,7 +195,7 @@ func (i *InteractiveHandler) displayAssetsOrProxy(assets []sdk.Asset) {
 	}
 }
 
-func (i *InteractiveHandler) displayAssets(assets []sdk.Asset) {
+func (i *InteractiveHandler) displayAssets(assets model.AssetList) {
 	if len(assets) == 0 {
 		_, _ = io.WriteString(i.sess, "\r\n No Assets\r\n\r")
 	} else {
@@ -235,99 +235,86 @@ func (i *InteractiveHandler) refreshAssetsAndNodesData() {
 	if err != nil {
 		logger.Error("refresh Assets  Nodes err:", err)
 	}
+}
 
+func (i *InteractiveHandler) loadUserAssets() {
+	i.assets = service.GetUserAssets(i.user.Id)
+}
+
+func (i *InteractiveHandler) loadUserAssetNodes() {
+	i.nodes = service.GetUserNodes(i.user.Id)
 }
 
-//
-//func (i *InteractiveHandler) loadUserAssets() {
-//	assets, err := appService.GetUserAssets(i.user.Id)
-//	if err != nil {
-//		logger.Error("load Assets failed")
-//		return
-//	}
-//	logger.Info("load Assets success")
-//	Cached.Store(i.user.Id, assets)
-//	i.assetData.Store(AssetsMapKey, assets)
-//}
-//
-//func (i *InteractiveHandler) loadUserAssetNodes() {
-//	assetNodes, err := appService.GetUserAssetNodes(i.user.Id)
-//	if err != nil {
-//		logger.Error("load Asset Nodes failed")
-//		return
-//	}
-//	logger.Info("load Asset Nodes success")
-//	i.assetData.Store(AssetNodesMapKey, assetNodes)
-//}
-//
 func (i *InteractiveHandler) changeLanguage() {
 
 }
 
-//
-//func (i *InteractiveHandler) JoinShareRoom(roomID string) {
-//	sshConn := userhome.NewSSHConn(i.sess)
-//	ctx, cancelFuc := context.WithCancel(i.sess.Context())
-//
-//	_, winCh, _ := i.sess.Pty()
-//	go func() {
-//		for {
-//			select {
-//			case <-ctx.Done():
-//				return
-//			case win, ok := <-winCh:
-//				if !ok {
-//					return
-//				}
-//				fmt.Println("join term change:", win)
-//			}
-//		}
-//	}()
-//	proxy.Manager.JoinShareRoom(roomID, sshConn)
-//	logger.Info("exit room id:", roomID)
-//	cancelFuc()
-//
-//}
-//
-//func (i *InteractiveHandler) searchAsset(key string) (assets []sdk.Asset) {
-//	if indexNum, err := strconv.Atoi(key); err == nil {
-//		if indexNum > 0 && indexNum <= len(i.currentSearchAssets) {
-//			return []sdk.Asset{i.currentSearchAssets[indexNum-1]}
-//		}
-//	}
-//
-//	if assetsData, ok := i.assetData.Load(AssetsMapKey); ok {
-//		for _, assetValue := range assetsData.([]sdk.Asset) {
-//			if isSubstring([]string{assetValue.Ip, assetValue.Hostname, assetValue.Comment}, key) {
-//				assets = append(assets, assetValue)
-//			}
-//		}
-//	} else {
-//		assetsData, _ := Cached.Load(i.user.Id)
-//		for _, assetValue := range assetsData.([]sdk.Asset) {
-//			if isSubstring([]string{assetValue.Ip, assetValue.Hostname, assetValue.Comment}, key) {
-//				assets = append(assets, assetValue)
-//			}
-//		}
-//	}
-//
-//	return assets
-//}
-//
-//func (i *InteractiveHandler) searchNodeAssets(num int) (assets []sdk.Asset) {
-//	var assetNodesData []sdk.Node
-//	if assetNodes, ok := i.assetData.Load(AssetNodesMapKey); ok {
-//		assetNodesData = assetNodes.([]sdk.Node)
-//		if num > len(assetNodesData) || num == 0 {
-//			return assets
-//		}
-//		return assetNodesData[num-1].AssetsGranted
-//	}
-//	return assets
-//
-//}
-//
-//func (i *InteractiveHandler) Proxy(asset sdk.Asset, systemUser sdk.SystemUserAuthInfo) error {
+func (i *InteractiveHandler) JoinShareRoom(roomID string) {
+	sshConn := userhome.NewSSHConn(i.sess)
+	ctx, cancelFuc := context.WithCancel(i.sess.Context())
+
+	_, winCh, _ := i.sess.Pty()
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case win, ok := <-winCh:
+				if !ok {
+					return
+				}
+				fmt.Println("join term change:", win)
+			}
+		}
+	}()
+	proxy.Manager.JoinShareRoom(roomID, sshConn)
+	logger.Info("exit room id:", roomID)
+	cancelFuc()
+
+}
+
+func (i *InteractiveHandler) searchAsset(key string) (assets []sdk.Asset) {
+	//if indexNum, err := strconv.Atoi(key); err == nil {
+	//	if indexNum > 0 && indexNum <= len(i.searchResult) {
+	//		return []sdk.Asset{i.searchResult[indexNum-1]}
+	//	}
+	//}
+	//
+	//if assetsData, ok := i.assetData.Load(AssetsMapKey); ok {
+	//	for _, assetValue := range assetsData.([]sdk.Asset) {
+	//		if isSubstring([]string{assetValue.Ip, assetValue.Hostname, assetValue.Comment}, key) {
+	//			assets = append(assets, assetValue)
+	//		}
+	//	}
+	//} else {
+	//	assetsData, _ := Cached.Load(i.user.Id)
+	//	for _, assetValue := range assetsData.([]sdk.Asset) {
+	//		if isSubstring([]string{assetValue.Ip, assetValue.Hostname, assetValue.Comment}, key) {
+	//			assets = append(assets, assetValue)
+	//		}
+	//	}
+	//}
+
+	return assets
+}
+
+func (i *InteractiveHandler) searchNodeAssets(num int) (assets []sdk.Asset) {
+	//var assetNodesData []sdk.Node
+	//if assetNodes, ok := i.assetData.Load(AssetNodesMapKey); ok {
+	//	assetNodesData = assetNodes.([]sdk.Node)
+	//	if num > len(assetNodesData) || num == 0 {
+	//		return assets
+	//	}
+	//	return assetNodesData[num-1].AssetsGranted
+	//}
+	return assets
+
+}
+
+func (i *InteractiveHandler) Proxy(asset sdk.Asset, systemUser sdk.SystemUserAuthInfo) error {
+	return nil
+}
+
 //	/*
 //		1. 创建SSHConn，符合core.Conn接口
 //		2. 创建一个session Home
@@ -433,15 +420,12 @@ func ConstructAssetNodeTree(assetNodes []sdk.Node) treeprint.Tree {
 func SessionHandler(sess ssh.Session) {
 	_, _, ptyOk := sess.Pty()
 	if ptyOk {
-		banner := NewBanner(sess.User())
 		handler := &InteractiveHandler{
-			sess:   sess,
-			term:   terminal.NewTerminal(sess, "Opt> "),
-			banner: banner,
+			sess: sess,
+			term: terminal.NewTerminal(sess, "Opt> "),
 		}
 
 		logger.Info("Accept one session")
-		handler.displayBanner()
 		handler.Dispatch()
 	} else {
 		_, err := io.WriteString(sess, "No PTY requested.\n")
@@ -449,5 +433,4 @@ func SessionHandler(sess ssh.Session) {
 			return
 		}
 	}
-
 }

pkg/sshd/handler/utils.go

+package handler
+
+import (
+	"fmt"
+	"strings"
+)
+
+const (
+	ColorStart = "\033["
+	Green      = "32m"
+	Red        = "41m"
+	White      = "47m"
+	ColorEnd   = "0m"
+	Tab        = "\t"
+	EndLine    = "\r\n"
+	Bold       = "1"
+)
+
+func WrapperString(text string, color string, bold bool) string {
+	wrapWith := make([]string, 1)
+	if bold {
+		wrapWith = append(wrapWith, Bold)
+	}
+	wrapWith = append(wrapWith, color)
+	return fmt.Sprintf("%s%s%s", strings.Join(wrapWith, ";"), text, ColorEnd)
+}
+
+func STitle(text string) string {
+	return WrapperString(text, Green, true)
+}
+
+func SWarn(text string) string {
+	return WrapperString(text, White, false)
+}

pkg/sshd/server.go

@@ -7,8 +7,9 @@ import (
 
 	"cocogo/pkg/auth"
 	"cocogo/pkg/config"
-	"cocogo/pkg/handler"
 	"cocogo/pkg/logger"
+
+	"./handler"
 )
 
 var (