From 767cb2dd21397bbe0662f7b07fb8512481c19f0f Mon Sep 17 00:00:00 2001
From: ZacharyZcR <PayasoNorahC@protonmail.com>
Date: Tue, 2 Sep 2025 11:58:52 +0000
Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84NetBIOS=E5=92=8CF?=
 =?UTF-8?q?indNet=E6=8F=92=E4=BB=B6=E4=BD=BF=E7=94=A8=E7=BB=9F=E4=B8=80?=
 =?UTF-8?q?=E5=8F=91=E5=8C=85=E6=8E=A7=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 修改NetBIOS插件，在UDP和TCP连接中添加发包控制
- 修改FindNet插件，在RPC TCP连接中添加发包控制
- 统一包计数逻辑，支持UDP和TCP协议的包计数
- 保持现有Windows网络发现和主机信息收集功能
---
 plugins/services/findnet.go | 12 ++++++++++++
 plugins/services/netbios.go | 18 ++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/plugins/services/findnet.go b/plugins/services/findnet.go
index b34f751..dbfd344 100644
--- a/plugins/services/findnet.go
+++ b/plugins/services/findnet.go
@@ -44,15 +44,27 @@ func (p *FindNetPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanRe
 		}
 	}
 
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("FindNet连接 %s 受限: %s", target, reason))
+		return &ScanResult{
+			Success: false,
+			Service: "findnet",
+			Error:   fmt.Errorf("发包受限: %s", reason),
+		}
+	}
+	
 	// 建立连接
 	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return &ScanResult{
 			Success: false,
 			Service: "findnet",
 			Error:   fmt.Errorf("连接RPC端口失败: %v", err),
 		}
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	// 设置超时
diff --git a/plugins/services/netbios.go b/plugins/services/netbios.go
index 521ace1..31247be 100644
--- a/plugins/services/netbios.go
+++ b/plugins/services/netbios.go
@@ -158,10 +158,19 @@ func (p *NetBIOSPlugin) queryNetBIOSNames(host string) (*NetBIOSInfo, error) {
 	}
 
 	target := fmt.Sprintf("%s:137", host)
+	
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("NetBIOS UDP连接 %s 受限: %s", target, reason))
+		return nil, fmt.Errorf("发包受限: %s", reason)
+	}
+	
 	conn, err := net.DialTimeout("udp", target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
+		common.IncrementUDPPacketCount()
 		return nil, fmt.Errorf("连接NetBIOS名称服务失败: %v", err)
 	}
+	common.IncrementUDPPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))
@@ -183,10 +192,19 @@ func (p *NetBIOSPlugin) queryNetBIOSNames(host string) (*NetBIOSInfo, error) {
 // queryNetBIOSSession 查询NetBIOS会话服务(TCP 139)
 func (p *NetBIOSPlugin) queryNetBIOSSession(host string) (*NetBIOSInfo, error) {
 	target := fmt.Sprintf("%s:139", host)
+	
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("NetBIOS TCP连接 %s 受限: %s", target, reason))
+		return nil, fmt.Errorf("发包受限: %s", reason)
+	}
+	
 	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return nil, fmt.Errorf("连接NetBIOS会话服务失败: %v", err)
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))