From ec9a30d78c557ac740976c5a629b4651b0c04b81 Mon Sep 17 00:00:00 2001
From: TsMask <340112800@qq.com>
Date: Tue, 23 Jan 2024 18:43:26 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20ws=E8=8E=B7=E5=8F=96=E7=BD=91=E7=BB=9C?=
 =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E8=BF=9B=E7=A8=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/modules/ws/model/net_connect.go       | 20 +++++++++
 src/modules/ws/processor/net_connect.go   | 53 +++++++++++++++++++++++
 src/modules/ws/service/ws_receive.impl.go |  6 +++
 3 files changed, 79 insertions(+)
 create mode 100644 src/modules/ws/model/net_connect.go
 create mode 100644 src/modules/ws/processor/net_connect.go

diff --git a/src/modules/ws/model/net_connect.go b/src/modules/ws/model/net_connect.go
new file mode 100644
index 00000000..8116c285
--- /dev/null
+++ b/src/modules/ws/model/net_connect.go
@@ -0,0 +1,20 @@
+package model
+
+import "github.com/shirou/gopsutil/v3/net"
+
+// NetConnectData 网络连接进程数据
+type NetConnectData struct {
+	Type   string   `json:"type"`
+	Status string   `json:"status"`
+	Laddr  net.Addr `json:"localaddr"`
+	Raddr  net.Addr `json:"remoteaddr"`
+	PID    int32    `json:"PID"`
+	Name   string   `json:"name"`
+}
+
+// NetConnectQuery 网络连接进程查询
+type NetConnectQuery struct {
+	Port        int32  `json:"port"`
+	ProcessName string `json:"processName"`
+	ProcessID   int32  `json:"processID"`
+}
diff --git a/src/modules/ws/processor/net_connect.go b/src/modules/ws/processor/net_connect.go
new file mode 100644
index 00000000..16b0c6a5
--- /dev/null
+++ b/src/modules/ws/processor/net_connect.go
@@ -0,0 +1,53 @@
+package processor
+
+import (
+	"encoding/json"
+	"strings"
+
+	"ems.agt/src/modules/ws/model"
+	"github.com/shirou/gopsutil/v3/net"
+	"github.com/shirou/gopsutil/v3/process"
+)
+
+// GetNetConnections 获取网络连接进程
+func GetNetConnections(data any) ([]byte, error) {
+	msgByte, _ := json.Marshal(data)
+	var query model.NetConnectQuery
+	err := json.Unmarshal(msgByte, &query)
+	if err != nil {
+		return nil, err
+	}
+
+	result := []model.NetConnectData{}
+	for _, netType := range [...]string{"tcp", "udp"} {
+		connections, _ := net.Connections(netType)
+		if err == nil {
+			for _, conn := range connections {
+				if query.ProcessID > 0 && query.ProcessID != conn.Pid {
+					continue
+				}
+				proc, err := process.NewProcess(conn.Pid)
+				if err == nil {
+					name, _ := proc.Name()
+					if name != "" && query.ProcessName != "" && !strings.Contains(name, query.ProcessName) {
+						continue
+					}
+					if query.Port > 0 && query.Port != int32(conn.Laddr.Port) && query.Port != int32(conn.Raddr.Port) {
+						continue
+					}
+					result = append(result, model.NetConnectData{
+						Type:   netType,
+						Status: conn.Status,
+						Laddr:  conn.Laddr,
+						Raddr:  conn.Raddr,
+						PID:    conn.Pid,
+						Name:   name,
+					})
+				}
+
+			}
+		}
+	}
+	resultByte, err := json.Marshal(result)
+	return resultByte, err
+}
diff --git a/src/modules/ws/service/ws_receive.impl.go b/src/modules/ws/service/ws_receive.impl.go
index bc548c2b..c6dda62c 100644
--- a/src/modules/ws/service/ws_receive.impl.go
+++ b/src/modules/ws/service/ws_receive.impl.go
@@ -23,6 +23,12 @@ func (s *WSReceiveImpl) Receive(client *model.WSClient, reqMsg model.WSRequest)
 			return err
 		}
 		client.MsgChan <- res
+	case "net":
+		res, err := processor.GetNetConnections(reqMsg.Data)
+		if err != nil {
+			return err
+		}
+		client.MsgChan <- res
 	default:
 		return fmt.Errorf("message type not supported")
 	}