From 1febb54fe6dad6e12ed31ce28d7b457147669280 Mon Sep 17 00:00:00 2001
From: ZacharyZcR <PayasoNorahC@protonmail.com>
Date: Tue, 2 Sep 2025 11:48:52 +0000
Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84SMB=E5=92=8CLDAP?=
 =?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

- 修改SMB插件，在testCredential和identifyService中添加发包控制
- 修改LDAP插件，在connectLDAP中添加发包控制和包计数
- 统一包计数逻辑，确保TCP连接成功和失败都正确计数
- 保持现有功能不变，提升网络操作一致性
---
 plugins/services/ldap.go | 33 +++++++++++++++++++++++++++++++--
 plugins/services/smb.go  | 19 +++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/plugins/services/ldap.go b/plugins/services/ldap.go
index daee1af..d9cf6da 100644
--- a/plugins/services/ldap.go
+++ b/plugins/services/ldap.go
@@ -84,17 +84,46 @@ func (p *LDAPPlugin) testCredential(ctx context.Context, info *common.HostInfo,
 }
 
 func (p *LDAPPlugin) connectLDAP(ctx context.Context, info *common.HostInfo, creds plugins.Credential) (*ldaplib.Conn, error) {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("LDAP连接 %s:%s 受限: %s", info.Host, info.Ports, reason))
+		return nil, fmt.Errorf("发包受限: %s", reason)
+	}
+	
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 	
+	var conn *ldaplib.Conn
+	var err error
+	
 	if info.Ports == "636" {
-		return ldaplib.DialTLS("tcp", target, nil)
+		conn, err = ldaplib.DialTLS("tcp", target, nil)
+	} else {
+		conn, err = ldaplib.Dial("tcp", target)
 	}
-	return ldaplib.Dial("tcp", target)
+	
+	// 统计包数量
+	if err != nil {
+		common.IncrementTCPFailedPacketCount()
+	} else {
+		common.IncrementTCPSuccessPacketCount()
+	}
+	
+	return conn, err
 }
 
 func (p *LDAPPlugin) identifyService(ctx context.Context, info *common.HostInfo) *plugins.Result {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("LDAP识别 %s 受限: %s", target, reason))
+		return &plugins.Result{
+			Success: false,
+			Service: "ldap",
+			Error:   fmt.Errorf("发包受限: %s", reason),
+		}
+	}
+
 	conn, err := p.connectLDAP(ctx, info, plugins.Credential{})
 	if err != nil {
 		return &plugins.Result{
diff --git a/plugins/services/smb.go b/plugins/services/smb.go
index 22702f0..c9751e9 100644
--- a/plugins/services/smb.go
+++ b/plugins/services/smb.go
@@ -105,6 +105,12 @@ func (p *SmbPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult
 
 // testCredential 测试单个凭据
 func (p *SmbPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("SMB连接 %s:%s 受限: %s", info.Host, info.Ports, reason))
+		return false
+	}
+	
 	options := smb.Options{
 		Host:        info.Host,
 		Port:        445,
@@ -124,8 +130,10 @@ func (p *SmbPlugin) testCredential(ctx context.Context, info *common.HostInfo, c
 		session, err := smb.NewSession(options, false)
 		if err == nil {
 			defer session.Close()
+			common.IncrementTCPSuccessPacketCount()
 			resultChan <- session.IsAuthenticated
 		} else {
+			common.IncrementTCPFailedPacketCount()
 			resultChan <- false
 		}
 	}()
@@ -199,6 +207,17 @@ func (p *SmbPlugin) testShareAccess(ctx context.Context, info *common.HostInfo,
 
 // identifyService 服务识别
 func (p *SmbPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
+		common.LogError(fmt.Sprintf("SMB识别 %s 受限: %s", target, reason))
+		return &ScanResult{
+			Success: false,
+			Service: "smb",
+			Error:   fmt.Errorf("发包受限: %s", reason),
+		}
+	}
+	
 	if p.testCredential(ctx, info, Credential{Username: "", Password: ""}) {
 		target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 		banner := "SMB文件共享服务"