[Update] 修改命令记录

pkg/proxy/parser.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"cocogo/pkg/logger"
 	"cocogo/pkg/model"
+	"cocogo/pkg/utils"
+	"fmt"
 	"sync"
 )
 
@@ -28,6 +30,11 @@ type Parser struct {
 	cmdBuf    *bytes.Buffer
 	outputBuf *bytes.Buffer
 
+	userInputChan  chan []byte
+	userOutputChan chan []byte
+	srvInputChan   chan []byte
+	srvOutputChan  chan []byte
+
 	cmdCh chan *[2]string
 
 	inputInitial  bool
@@ -59,31 +66,59 @@ func (p *Parser) Initial() {
 	p.cmdOutputParser.Initial()
 }
 
+func (p *Parser) Parse() {
+	defer func() {
+		fmt.Println("Parse done")
+	}()
+	for {
+		select {
+		case ub, ok := <-p.userInputChan:
+			if !ok {
+				return
+			}
+			b := p.ParseUserInput(ub)
+			p.userOutputChan <- b
+
+		case sb, ok := <-p.srvInputChan:
+			if !ok {
+				return
+			}
+			b := p.ParseServerOutput(sb)
+			p.srvOutputChan <- b
+		}
+	}
+}
+
 // Todo: parseMultipleInput 依然存在问题
 
 // parseInputState 切换用户输入状态, 并结算命令和结果
-func (p *Parser) parseInputState(b []byte) {
+func (p *Parser) parseInputState(b []byte) []byte {
 	if p.inVimState || p.zmodemState != "" {
-		return
+		return b
 	}
 	p.inputPreState = p.inputState
 	if bytes.Contains(b, charEnter) {
 		p.inputState = false
 		// 用户输入了Enter，开始结算命令
 		p.parseCmdInput()
+		if p.IsCommandForbidden() {
+			p.srvOutputChan <- []byte("\r\nCommand ls is forbidden")
+			return []byte{utils.CharCleanLine, '\r'}
+		}
 	} else {
 		p.inputState = true
 		// 用户又开始输入，并上次不处于输入状态，开始结算上次命令的结果
 		if !p.inputPreState {
 			p.parseCmdOutput()
-			p.cmdCh <- &[2]string{p.command, p.output}
 		}
 	}
+	return b
 }
 
 func (p *Parser) parseCmdInput() {
 	data := p.cmdBuf.Bytes()
 	p.command = p.cmdInputParser.Parse(data)
+	fmt.Println("parse Command is ", p.command)
 	p.cmdBuf.Reset()
 	p.inputBuf.Reset()
 }
@@ -106,9 +141,8 @@ func (p *Parser) ParseUserInput(b []byte) []byte {
 		p.inputInitial = true
 	})
 	nb := p.replaceInputNewLine(b)
-	p.inputBuf.Write(nb)
-	p.parseInputState(nb)
-	return b
+	nb = p.parseInputState(nb)
+	return nb
 }
 
 func (p *Parser) parseZmodemState(b []byte) {
@@ -149,9 +183,10 @@ func (p *Parser) parseVimState(b []byte) {
 func (p *Parser) splitCmdStream(b []byte) {
 	p.parseVimState(b)
 	p.parseZmodemState(b)
-	if p.zmodemState != "" || p.inVimState || p.inputInitial {
+	if p.zmodemState != "" || p.inVimState || !p.inputInitial {
 		return
 	}
+	fmt.Println("Input state: ", p.inputState)
 	if p.inputState {
 		p.cmdBuf.Write(b)
 	} else {
@@ -167,3 +202,11 @@ func (p *Parser) ParseServerOutput(b []byte) []byte {
 func (p *Parser) SetCMDFilterRules(rules []model.SystemUserFilterRule) {
 	p.cmdFilterRules = rules
 }
+
+func (p *Parser) IsCommandForbidden() bool {
+	fmt.Println("Command is: ", p.command)
+	if p.command == "ls" {
+		return true
+	}
+	return false
+}

pkg/proxy/proxy.go

@@ -23,15 +23,18 @@ type ProxyServer struct {
 
 func (p *ProxyServer) getSystemUserAuthOrManualSet() {
 	info := service.GetSystemUserAssetAuthInfo(p.SystemUser.Id, p.Asset.Id)
+	p.SystemUser.Password = info.Password
+	p.SystemUser.PrivateKey = info.PrivateKey
+
 	if p.SystemUser.LoginMode == model.LoginModeManual ||
 		(p.SystemUser.Password == "" && p.SystemUser.PrivateKey == "") {
+		// Todo: terminal
 		logger.Info("Get password fom user input")
 	}
-	p.SystemUser.Password = info.Password
-	p.SystemUser.PrivateKey = info.PrivateKey
 }
 
 func (p *ProxyServer) getSystemUserUsernameIfNeed() {
+	// Todo: terminal
 
 }
 
@@ -74,6 +77,8 @@ func (p *ProxyServer) getTelnetConn() (srvConn *ServerSSHConnection, err error)
 }
 
 func (p *ProxyServer) getServerConn() (srvConn ServerConnection, err error) {
+	p.getSystemUserUsernameIfNeed()
+	p.getSystemUserAuthOrManualSet()
 	if p.Asset.Protocol == "telnet" {
 		return p.getTelnetConn()
 	} else {
@@ -138,8 +143,6 @@ func (p *ProxyServer) Proxy() {
 		logger.Error("Get system user filter rule error: ", err)
 	}
 	sw.parser.SetCMDFilterRules(cmdRules)
-	replayRecorder := NewReplyRecord(sw.Id)
-	sw.parser.SetReplayRecorder(replayRecorder)
 	_ = sw.Bridge()
 	_ = srvConn.Close()
 }

pkg/proxy/srvconn.go

@@ -121,7 +121,6 @@ func (sc *ServerSSHConnection) connect() (client *gossh.Client, err error) {
 	} else {
 		client, err = gossh.Dial("tcp", net.JoinHostPort(sc.host, sc.port), config)
 		if err != nil {
-			err = fmt.Errorf("connect host %sc error: %sc", sc.host, err)
 			return
 		}
 	}

pkg/proxy/switch.go

@@ -6,13 +6,18 @@ import (
 	"time"
 
 	"github.com/gliderlabs/ssh"
-	uuid "github.com/satori/go.uuid"
+	"github.com/satori/go.uuid"
 
 	"cocogo/pkg/logger"
 )
 
 func NewSwitchSession(userConn UserConnection, serverConn ServerConnection) (sw *SwitchSession) {
-	parser := new(Parser)
+	parser := &Parser{
+		userInputChan:  make(chan []byte, 1024),
+		userOutputChan: make(chan []byte, 1024),
+		srvInputChan:   make(chan []byte, 1024),
+		srvOutputChan:  make(chan []byte, 1024),
+	}
 	parser.Initial()
 	sw = &SwitchSession{userConn: userConn, serverConn: serverConn, parser: parser}
 	return sw
@@ -100,11 +105,12 @@ func (s *SwitchSession) readUserToServer(ctx context.Context) {
 			if !ok {
 				s.cancelFunc()
 			}
-			buf2 := s.parser.ParseUserInput(p)
-			_, err := s.serverTran.Write(buf2)
-			if err != nil {
-				return
+			s.parser.userInputChan <- p
+		case p, ok := <-s.parser.userOutputChan:
+			if !ok {
+				s.cancelFunc()
 			}
+			_, _ = s.serverTran.Write(p)
 		}
 	}
 }
@@ -122,11 +128,12 @@ func (s *SwitchSession) readServerToUser(ctx context.Context) {
 			if !ok {
 				s.cancelFunc()
 			}
-			buf2 := s.parser.ParseServerOutput(p)
-			_, err := s.userTran.Write(buf2)
-			if err != nil {
-				return
+			s.parser.srvInputChan <- p
+		case p, ok := <-s.parser.srvOutputChan:
+			if !ok {
+				s.cancelFunc()
 			}
+			_, _ = s.userConn.Write(p)
 		}
 	}
 }
@@ -138,6 +145,7 @@ func (s *SwitchSession) Bridge() (err error) {
 
 	s.userTran = NewDirectTransport("", s.userConn)
 	s.serverTran = NewDirectTransport("", s.serverConn)
+	go s.parser.Parse()
 	go s.watchWindowChange(ctx, winCh)
 	go s.readServerToUser(ctx)
 	s.readUserToServer(ctx)

pkg/service/init.go

@@ -32,7 +32,7 @@ func Initial() {
 func validateAccessAuth() {
 	maxTry := 30
 	count := 0
-	for count < maxTry {
+	for {
 		user, err := GetProfile()
 		if err == nil && user.Role == "App" {
 			break
@@ -40,12 +40,11 @@ func validateAccessAuth() {
 		if err != nil {
 			msg := "Connect server error or access key is invalid, remove %s run again"
 			logger.Errorf(msg, config.Conf.AccessKeyFile)
-		}
-		if user.Role != "App" {
+		} else if user.Role != "App" {
 			logger.Error("Access role is not App, is: ", user.Role)
 		}
-		time.Sleep(3 * time.Second)
 		count++
+		time.Sleep(3 * time.Second)
 		if count >= maxTry {
 			os.Exit(1)
 		}

pkg/utils/util.go

@@ -19,9 +19,10 @@ const (
 )
 
 const (
-	CharClear   = "\x1b[H\x1b[2J"
-	CharTab     = "\t"
-	CharNewLine = "\r\n"
+	CharClear     = "\x1b[H\x1b[2J"
+	CharTab       = "\t"
+	CharNewLine   = "\r\n"
+	CharCleanLine = '\x15'
 )
 
 func WrapperString(text string, color string, meta ...bool) string {