From e597db193cbfd271d867c9301026c1450afd3019 Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Tue, 17 Oct 2023 19:17:49 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20UPF=E6=8A=93=E5=8C=85=E8=84=9A=E6=9C=AC?=
 =?UTF-8?q?=E6=88=96telnet=E6=96=B9=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 features/trace/tcpdump.go | 79 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 73 insertions(+), 6 deletions(-)

diff --git a/features/trace/tcpdump.go b/features/trace/tcpdump.go
index ccab0b5b..3c1ffb0f 100644
--- a/features/trace/tcpdump.go
+++ b/features/trace/tcpdump.go
@@ -130,15 +130,15 @@ func TcpdumpNeUPFTask(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// 开始
-	if body.RunType == "start" {
+	// 开始telnet
+	if body.RunType == "start_telnet" {
 		// 创建TCP连接
 		conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", neInfo.Ip, 5002))
 		if err != nil {
-			conn.Close()
 			ctx.JSON(w, 200, result.ErrMsg(err.Error()))
 			return
 		}
+		defer conn.Close()
 
 		filePcapName := fmt.Sprintf("tmp_%s_%s.pcap", body.NeType, body.NeId)
 		cmdStr := fmt.Sprintf("pcap dispatch trace on max 100000 file %s", filePcapName)
@@ -169,15 +169,15 @@ func TcpdumpNeUPFTask(w http.ResponseWriter, r *http.Request) {
 		conn.Close()
 		return
 	}
-	// 停止
-	if body.RunType == "stop" {
+	// 停止telnet
+	if body.RunType == "stop_telnet" {
 		// 创建TCP连接
 		conn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", neInfo.Ip, 5002))
 		if err != nil {
-			conn.Close()
 			ctx.JSON(w, 200, result.ErrMsg(err.Error()))
 			return
 		}
+		defer conn.Close()
 
 		filePcapName := fmt.Sprintf("tmp_%s_%s.pcap", body.NeType, body.NeId)
 		cmdStr := "pcap dispatch trace off"
@@ -268,5 +268,72 @@ func TcpdumpNeUPFTask(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// 开始-脚本字符串
+	if body.RunType == "start_str" {
+		fileLogName := fmt.Sprintf("tmp_%s_%s.log", body.NeType, body.NeId)
+		filePcapName := fmt.Sprintf("tmp_%s_%s.pcap", body.NeType, body.NeId)
+		scriptStr := "#!/bin/expect\nset capcmd [lindex $argv 0]\nspawn telnet localhost 5002\nexpect \"upfd1# \"\nsend \"$capcmd\\n\"\nexpect \"upfd1# \"\nsend \"quit\\n\"\nexpect \"eof\""
+		writeLog := fmt.Sprintf(" > %s 2>&1 \ncat %s", fileLogName, fileLogName) // 执行信息写入日志文件输出，避免弹出code 127
+
+		capCmdStr := fmt.Sprintf("%s file %s", body.Cmd, filePcapName)
+
+		cmdStr := fmt.Sprintf("cd /tmp\n\necho '%s' > cap.sh\n\nchmod +x cap.sh\n\n./cap.sh '%s'%s", scriptStr, capCmdStr, writeLog)
+		usernameNe := conf.Get("ne.user").(string) // 网元统一用户
+		sshHost := fmt.Sprintf("%s@%s", usernameNe, neInfo.Ip)
+		msg, err := cmd.ExecWithCheck("ssh", sshHost, cmdStr)
+		if err != nil {
+			ctx.JSON(w, 200, result.ErrMsg(err.Error()))
+		} else {
+			s := strings.Index(msg, "pcap dispatch trace:")
+			if s != -1 {
+				e := strings.Index(msg, "\r\nupfd1#")
+				msg = msg[s:e]
+			} else {
+				msg = "Executed, please stop before proceeding"
+			}
+			ctx.JSON(w, 200, result.OkData(map[string]any{
+				"cmd":      capCmdStr,
+				"msg":      msg,
+				"fileName": filePcapName,
+			}))
+		}
+		return
+	}
+	// 停止-脚本字符串
+	if body.RunType == "stop_str" {
+		fileLogName := fmt.Sprintf("tmp_%s_%s.log", body.NeType, body.NeId)
+		filePcapName := fmt.Sprintf("tmp_%s_%s.pcap", body.NeType, body.NeId)
+		scriptStr := "#!/bin/expect\nset capcmd [lindex $argv 0]\nspawn telnet localhost 5002\nexpect \"upfd1# \"\nsend \"$capcmd\\n\"\nexpect \"upfd1# \"\nsend \"quit\\n\"\nexpect \"eof\""
+		writeLog := fmt.Sprintf(" > %s 2>&1 \ncat %s", fileLogName, fileLogName) // 执行信息写入日志文件输出，避免弹出code 127
+
+		capCmdStr := body.Cmd
+
+		cmdStr := fmt.Sprintf("cd /tmp\n\necho '%s' > cap.sh\n\nchmod +x cap.sh\n\n./cap.sh '%s'%s", scriptStr, capCmdStr, writeLog)
+
+		usernameNe := conf.Get("ne.user").(string) // 网元统一用户
+		sshHost := fmt.Sprintf("%s@%s", usernameNe, neInfo.Ip)
+		msg, err := cmd.ExecWithCheck("ssh", sshHost, cmdStr)
+		if err != nil {
+			ctx.JSON(w, 200, result.ErrMsg(err.Error()))
+		} else {
+			s := strings.Index(msg, "pcap dispatch trace:")
+			if s == -1 {
+				s = strings.Index(msg, "Write ")
+			}
+			if s != -1 {
+				e := strings.Index(msg, "\r\nupfd1#")
+				msg = msg[s:e]
+			} else {
+				msg = "No stoppable found"
+			}
+			ctx.JSON(w, 200, result.OkData(map[string]any{
+				"cmd":      capCmdStr,
+				"msg":      msg,
+				"fileName": filePcapName,
+			}))
+		}
+		return
+	}
+
 	ctx.JSON(w, 200, result.ErrMsg("runType is start or stop"))
 }