perf: ws模块文件实时查看业务处理

2024-08-07 19:35:42 +08:00
parent a5363b1ce1
commit 0d312e61fa
2 changed files with 170 additions and 61 deletions
--- a/src/modules/ws/processor/shell_command.go
+++ b/src/modules/ws/processor/shell_command.go
@@ -0,0 +1,75 @@
+package processor
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"be.ems/src/framework/logger"
+)
+
+// ParseCat 解析拼装cat命令
+func ParseCat(reqData any) (string, error) {
+	msgByte, _ := json.Marshal(reqData)
+	var data struct {
+		FilePath   string `json:"filePath"`   // 文件地址
+		ShowNumber bool   `json:"showNumber"` // 显示文件的行号，从 1 开始
+		ShowTabs   bool   `json:"showTabs"`   // 在每行的结尾显示 $ 符号
+		ShowAll    bool   `json:"showAll"`    // 结合 -vET 参数，显示所有特殊字符，包括行尾符、制表符等
+	}
+	if err := json.Unmarshal(msgByte, &data); err != nil {
+		logger.Warnf("ws processor ParseCat err: %s", err.Error())
+		return "", fmt.Errorf("query data structure error")
+	}
+	if data.FilePath == "" {
+		return "", fmt.Errorf("query data filePath empty")
+	}
+
+	command := []string{"cat"}
+	if data.ShowNumber {
+		command = append(command, "-n")
+	}
+	if data.ShowTabs {
+		command = append(command, "-T")
+	}
+	if data.ShowAll {
+		command = append(command, "-A")
+	}
+
+	command = append(command, data.FilePath)
+	command = append(command, "\n")
+	return strings.Join(command, " "), nil
+}
+
+// ParseTail 解析拼装tail命令
+func ParseTail(reqData any) (string, error) {
+	msgByte, _ := json.Marshal(reqData)
+	var data struct {
+		FilePath string `json:"filePath"` // 文件地址
+		Lines    int    `json:"lines"`    // 显示文件末尾的指定行数
+		Char     int    `json:"char"`     // 显示文件末尾的指定字数
+		Follow   bool   `json:"follow"`   // 输出文件末尾的内容，并继续监视文件的新增内容
+	}
+	if err := json.Unmarshal(msgByte, &data); err != nil {
+		logger.Warnf("ws processor ParseTail err: %s", err.Error())
+		return "", fmt.Errorf("query data structure error")
+	}
+	if data.FilePath == "" {
+		return "", fmt.Errorf("query data filePath empty")
+	}
+
+	command := []string{"tail"}
+	if data.Follow {
+		command = append(command, "-f")
+	}
+	if data.Lines > 0 {
+		command = append(command, fmt.Sprintf("-n %d", data.Lines))
+	}
+	if data.Char > 0 {
+		command = append(command, fmt.Sprintf("-c %d", data.Char))
+	}
+
+	command = append(command, data.FilePath)
+	command = append(command, "\n")
+	return strings.Join(command, " "), nil
+}
--- a/src/modules/ws/service/ws_receive.impl.go
+++ b/src/modules/ws/service/ws_receive.impl.go
@@ -65,7 +65,7 @@ func (s *WSReceiveImpl) Commont(client *model.WSClient, reqMsg model.WSRequest)
 	case "ne_state":
 		resByte, err = processor.GetNeState(reqMsg.RequestID, reqMsg.Data)
 	default:
-		err = fmt.Errorf("message type not supported")
+		err = fmt.Errorf("message type %s not supported", reqMsg.Type)
 	}

 	if err != nil {
@@ -99,40 +99,23 @@ func (s *WSReceiveImpl) Shell(client *model.WSClient, reqMsg model.WSRequest) {
 		return
 	case "ssh":
 		// SSH会话消息接收写入会话
-		if client.ChildConn == nil {
-			err = fmt.Errorf("message type ssh not connected")
-		} else {
-			command := reqMsg.Data.(string)
-			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-			_, err = sshClientSession.Write(command)
-		}
+		command := reqMsg.Data.(string)
+		sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+		_, err = sshClientSession.Write(command)
 	case "ssh_resize":
 		// SSH会话窗口重置
-		if client.ChildConn == nil {
-			err = fmt.Errorf("message type ssh not connected")
-		} else {
-			msgByte, _ := json.Marshal(reqMsg.Data)
-			var data struct {
-				Cols int `json:"cols"`
-				Rows int `json:"rows"`
-			}
-			err = json.Unmarshal(msgByte, &data)
-			if err == nil {
-				sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-				err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
-			}
+		msgByte, _ := json.Marshal(reqMsg.Data)
+		var data struct {
+			Cols int `json:"cols"`
+			Rows int `json:"rows"`
 		}
-	case "telnet":
-		// Telnet会话消息接收写入会话
-		if client.ChildConn == nil {
-			err = fmt.Errorf("message type telnet not connected")
-		} else {
-			command := reqMsg.Data.(string)
-			telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
-			_, err = telnetClientSession.Write(command)
+		err = json.Unmarshal(msgByte, &data)
+		if err == nil {
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
 		}
 	default:
-		err = fmt.Errorf("message type not supported")
+		err = fmt.Errorf("message type %s not supported", reqMsg.Type)
 	}

 	if err != nil {
@@ -153,6 +136,73 @@ func (s *WSReceiveImpl) Shell(client *model.WSClient, reqMsg model.WSRequest) {

 // ShellView 接收查看文件终端交互业务处理
 func (s *WSReceiveImpl) ShellView(client *model.WSClient, reqMsg model.WSRequest) {
+	// 必传requestId确认消息
+	if reqMsg.RequestID == "" {
+		msg := "message requestId is required"
+		logger.Infof("ws ShellView UID %s err: %s", client.BindUid, msg)
+		msgByte, _ := json.Marshal(result.ErrMsg(msg))
+		client.MsgChan <- msgByte
+		return
+	}
+
+	var resByte []byte
+	var err error
+
+	switch reqMsg.Type {
+	case "close":
+		s.close(client)
+		return
+	case "cat":
+	case "tail":
+		var command string
+		if reqMsg.Type == "cat" {
+			command, err = processor.ParseCat(reqMsg.Data)
+		}
+		if reqMsg.Type == "tail" {
+			command, err = processor.ParseTail(reqMsg.Data)
+		}
+		if command != "" && err == nil {
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			_, err = sshClientSession.Write(command)
+		}
+	case "ctrl-c":
+		// 模拟按下 Ctrl+C
+		sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+		_, err = sshClientSession.Write("\x03")
+	case "resize":
+		// 会话窗口重置
+		msgByte, _ := json.Marshal(reqMsg.Data)
+		var data struct {
+			Cols int `json:"cols"`
+			Rows int `json:"rows"`
+		}
+		err = json.Unmarshal(msgByte, &data)
+		if err == nil {
+			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
+			err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
+		}
+	default:
+		err = fmt.Errorf("message type %s not supported", reqMsg.Type)
+	}
+
+	if err != nil {
+		logger.Warnf("ws ShellView UID %s err: %s", client.BindUid, err.Error())
+		msgByte, _ := json.Marshal(result.ErrMsg(err.Error()))
+		client.MsgChan <- msgByte
+		if err == io.EOF {
+			// 等待1s后关闭连接
+			time.Sleep(1 * time.Second)
+			client.StopChan <- struct{}{}
+		}
+		return
+	}
+	if len(resByte) > 0 {
+		client.MsgChan <- resByte
+	}
+}
+
+// Telnet 接收终端交互业务处理
+func (s *WSReceiveImpl) Telnet(client *model.WSClient, reqMsg model.WSRequest) {
 	// 必传requestId确认消息
 	if reqMsg.RequestID == "" {
 		msg := "message requestId is required"
@@ -169,42 +219,26 @@ func (s *WSReceiveImpl) ShellView(client *model.WSClient, reqMsg model.WSRequest
 	case "close":
 		s.close(client)
 		return
-	case "ssh":
-		// SSH会话消息接收写入会话
-		if client.ChildConn == nil {
-			err = fmt.Errorf("message type ssh not connected")
-		} else {
-			command := reqMsg.Data.(string)
-			sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-			_, err = sshClientSession.Write(command)
-		}
-	case "ssh_resize":
-		// SSH会话窗口重置
-		if client.ChildConn == nil {
-			err = fmt.Errorf("message type ssh not connected")
-		} else {
-			msgByte, _ := json.Marshal(reqMsg.Data)
-			var data struct {
-				Cols int `json:"cols"`
-				Rows int `json:"rows"`
-			}
-			err = json.Unmarshal(msgByte, &data)
-			if err == nil {
-				sshClientSession := client.ChildConn.(*ssh.SSHClientSession)
-				err = sshClientSession.Session.WindowChange(data.Rows, data.Cols)
-			}
-		}
 	case "telnet":
 		// Telnet会话消息接收写入会话
-		if client.ChildConn == nil {
-			err = fmt.Errorf("message type telnet not connected")
-		} else {
-			command := reqMsg.Data.(string)
+		command := reqMsg.Data.(string)
+		telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
+		_, err = telnetClientSession.Write(command)
+	case "telnet_resize":
+		// Telnet会话窗口重置
+		msgByte, _ := json.Marshal(reqMsg.Data)
+		var data struct {
+			Cols int `json:"cols"`
+			Rows int `json:"rows"`
+		}
+		err = json.Unmarshal(msgByte, &data)
+		if err == nil {
 			telnetClientSession := client.ChildConn.(*telnet.TelnetClientSession)
-			_, err = telnetClientSession.Write(command)
+			err = telnetClientSession.WindowChange(data.Rows, data.Cols)
+			_ = telnetClientSession.Read()
 		}
 	default:
-		err = fmt.Errorf("message type not supported")
+		err = fmt.Errorf("message type %s not supported", reqMsg.Type)
 	}

 	if err != nil {