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

- 修改VNC插件，在所有网络连接点添加发包控制和计数
- 修改RDP插件，在testRDPConnection和checkNLAStatus中添加发包控制
- 统一包计数逻辑，确保TCP连接成功和失败都正确计数
- 保持现有远程桌面检测功能完整性
---
 plugins/services/rdp.go | 20 +++++++++++++++++---
 plugins/services/vnc.go | 36 ++++++++++++++++++++++++++++++++----
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/plugins/services/rdp.go b/plugins/services/rdp.go
index 78c36a2..214fac6 100644
--- a/plugins/services/rdp.go
+++ b/plugins/services/rdp.go
@@ -100,10 +100,18 @@ func (p *RDPPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult
 func (p *RDPPlugin) testRDPConnection(ctx context.Context, info *common.HostInfo) bool {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 	
-	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
-	if err != nil {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("RDP连接 %s 受限: %s", target, reason))
 		return false
 	}
+	
+	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
+	if err != nil {
+		common.IncrementTCPFailedPacketCount()
+		return false
+	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))
@@ -136,13 +144,19 @@ func (p *RDPPlugin) testRDPConnection(ctx context.Context, info *common.HostInfo
 
 // checkNLAStatus 检查网络级别身份验证状态
 func (p *RDPPlugin) checkNLAStatus(ctx context.Context, info *common.HostInfo) string {
-	// 简化实现，实际需要解析RDP协商响应
+	// 检查发包限制
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("RDP NLA检测 %s 受限: %s", target, reason))
+		return "检测失败"
+	}
 	
 	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return "检测失败"
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))
diff --git a/plugins/services/vnc.go b/plugins/services/vnc.go
index aea239f..f637dfb 100644
--- a/plugins/services/vnc.go
+++ b/plugins/services/vnc.go
@@ -70,10 +70,18 @@ func (p *VNCPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult
 func (p *VNCPlugin) testUnauthAccess(ctx context.Context, info *common.HostInfo) *ScanResult {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 	
-	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
-	if err != nil {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("VNC未授权检测 %s 受限: %s", target, reason))
 		return nil
 	}
+	
+	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
+	if err != nil {
+		common.IncrementTCPFailedPacketCount()
+		return nil
+	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))
@@ -105,10 +113,18 @@ func (p *VNCPlugin) testUnauthAccess(ctx context.Context, info *common.HostInfo)
 func (p *VNCPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 	
-	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
-	if err != nil {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("VNC认证测试 %s 受限: %s", target, reason))
 		return false
 	}
+	
+	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
+	if err != nil {
+		common.IncrementTCPFailedPacketCount()
+		return false
+	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))
@@ -229,14 +245,26 @@ func (p *VNCPlugin) reverseBits(b byte) byte {
 func (p *VNCPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("VNC识别 %s 受限: %s", target, reason))
+		return &ScanResult{
+			Success: false,
+			Service: "vnc",
+			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: "vnc",
 			Error:   err,
 		}
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))