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

- 修改SMTP插件，在多个连接点添加发包控制
- 修改Telnet插件，在identifyService中使用SafeTCPDial包装器
- 保持现有功能不变，统一发包控制逻辑
---
 plugins/services/smtp.go   | 32 ++++++++++++++++++++++++++------
 plugins/services/telnet.go | 14 +++++++++++++-
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/plugins/services/smtp.go b/plugins/services/smtp.go
index c52d895..50a1eb0 100644
--- a/plugins/services/smtp.go
+++ b/plugins/services/smtp.go
@@ -90,12 +90,23 @@ func (p *SMTPPlugin) Scan(ctx context.Context, info *common.HostInfo) *plugins.R
 func (p *SMTPPlugin) testOpenRelay(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("SMTP连接 %s 受限: %s", target, reason))
+		return nil
+	}
+	
 	// 设置超时的Dialer
 	dialer := &net.Dialer{
 		Timeout: time.Duration(common.Timeout) * time.Second,
 	}
 	
 	conn, err := dialer.DialContext(ctx, "tcp", target)
+	if err == nil {
+		common.IncrementTCPSuccessPacketCount()
+	} else {
+		common.IncrementTCPFailedPacketCount()
+	}
 	if err != nil {
 		return nil
 	}
@@ -132,12 +143,23 @@ func (p *SMTPPlugin) testCredential(ctx context.Context, info *common.HostInfo,
 	
 	common.LogDebug(fmt.Sprintf("SMTP测试凭据: %s:%s", cred.Username, cred.Password))
 	
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("SMTP凭据测试 %s 受限: %s", target, reason))
+		return false
+	}
+
 	// 设置连接超时
 	dialer := &net.Dialer{
 		Timeout: timeout,
 	}
 	
 	conn, err := dialer.DialContext(ctx, "tcp", target)
+	if err == nil {
+		common.IncrementTCPSuccessPacketCount()
+	} else {
+		common.IncrementTCPFailedPacketCount()
+	}
 	if err != nil {
 		common.LogDebug(fmt.Sprintf("SMTP连接失败: %v", err))
 		return false
@@ -177,7 +199,8 @@ func (p *SMTPPlugin) testCredential(ctx context.Context, info *common.HostInfo,
 func (p *SMTPPlugin) getServerInfo(ctx context.Context, info *common.HostInfo) string {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 	
-	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
+	// 使用统一TCP包装器
+	conn, err := common.SafeTCPDial(target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
 		return ""
 	}
@@ -207,11 +230,8 @@ func (p *SMTPPlugin) identifyService(ctx context.Context, info *common.HostInfo)
 	if serverInfo != "" {
 		banner = fmt.Sprintf("SMTP邮件服务 (%s)", serverInfo)
 	} else {
-		// 简单连接测试
-		dialer := &net.Dialer{
-			Timeout: time.Duration(common.Timeout) * time.Second,
-		}
-		conn, err := dialer.DialContext(ctx, "tcp", target)
+		// 使用统一TCP包装器进行简单连接测试
+		conn, err := common.SafeTCPDial(target, time.Duration(common.Timeout)*time.Second)
 		if err != nil {
 			return &plugins.Result{
 				Success: false,
diff --git a/plugins/services/telnet.go b/plugins/services/telnet.go
index b22abbd..d294d33 100644
--- a/plugins/services/telnet.go
+++ b/plugins/services/telnet.go
@@ -65,8 +65,19 @@ func (p *TelnetPlugin) Scan(ctx context.Context, info *common.HostInfo) *plugins
 func (p *TelnetPlugin) testTelnetCredential(ctx context.Context, info *common.HostInfo, cred plugins.Credential) bool {
 	address := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 	
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("Telnet连接 %s 受限: %s", address, reason))
+		return false
+	}
+	
 	// 创建带超时的连接
 	conn, err := net.DialTimeout("tcp", address, time.Duration(common.Timeout)*time.Second)
+	if err == nil {
+		common.IncrementTCPSuccessPacketCount()
+	} else {
+		common.IncrementTCPFailedPacketCount()
+	}
 	if err != nil {
 		return false
 	}
@@ -304,7 +315,8 @@ func (p *TelnetPlugin) isLoginFailed(data string) bool {
 func (p *TelnetPlugin) identifyService(ctx context.Context, info *common.HostInfo) *plugins.Result {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
 
-	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
+	// 使用统一TCP包装器
+	conn, err := common.SafeTCPDial(target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
 		return &plugins.Result{
 			Success: false,