From 36f0e5076d982e89c4931f61927173673424bfc0 Mon Sep 17 00:00:00 2001
From: ZacharyZcR <PayasoNorahC@protonmail.com>
Date: Tue, 2 Sep 2025 11:52:59 +0000
Subject: [PATCH] =?UTF-8?q?refactor:=20=E9=87=8D=E6=9E=84Memcached?=
 =?UTF-8?q?=E5=92=8CRabbitMQ=E6=8F=92=E4=BB=B6=E4=BD=BF=E7=94=A8=E7=BB=9F?=
 =?UTF-8?q?=E4=B8=80=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

- 修改Memcached插件，在TCP连接和服务识别中添加发包控制
- 修改RabbitMQ插件，在AMQP连接、HTTP连接和管理接口中添加发包控制
- 统一包计数逻辑，确保TCP连接成功和失败都正确计数
- 保持现有缓存服务和消息队列检测功能
---
 plugins/services/memcached.go | 19 +++++++++++++++++++
 plugins/services/rabbitmq.go  | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/plugins/services/memcached.go b/plugins/services/memcached.go
index 44d3840..c2a8df0 100644
--- a/plugins/services/memcached.go
+++ b/plugins/services/memcached.go
@@ -59,12 +59,21 @@ func (p *MemcachedPlugin) Scan(ctx context.Context, info *common.HostInfo) *Scan
 
 func (p *MemcachedPlugin) connectToMemcached(ctx context.Context, info *common.HostInfo) net.Conn {
 	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
+	
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("Memcached连接 %s 受限: %s", target, reason))
+		return nil
+	}
+	
 	timeout := time.Duration(common.Timeout) * time.Second
 
 	conn, err := net.DialTimeout("tcp", target, timeout)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return nil
 	}
+	common.IncrementTCPSuccessPacketCount()
 
 	conn.SetDeadline(time.Now().Add(timeout))
 
@@ -100,6 +109,16 @@ func (p *MemcachedPlugin) testBasicCommand(conn net.Conn) bool {
 func (p *MemcachedPlugin) 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("Memcached识别 %s 受限: %s", target, reason))
+		return &ScanResult{
+			Success: false,
+			Service: "memcached",
+			Error:   fmt.Errorf("发包受限: %s", reason),
+		}
+	}
+
 	conn := p.connectToMemcached(ctx, info)
 	if conn == nil {
 		return &ScanResult{
diff --git a/plugins/services/rabbitmq.go b/plugins/services/rabbitmq.go
index 907d887..80a546b 100644
--- a/plugins/services/rabbitmq.go
+++ b/plugins/services/rabbitmq.go
@@ -90,15 +90,27 @@ func (p *RabbitMQPlugin) Scan(ctx context.Context, info *common.HostInfo) *plugi
 func (p *RabbitMQPlugin) testAMQPProtocol(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("RabbitMQ AMQP连接 %s 受限: %s", target, reason))
+		return &plugins.Result{
+			Success: false,
+			Service: "rabbitmq",
+			Error:   fmt.Errorf("发包受限: %s", reason),
+		}
+	}
+	
 	// 连接到AMQP端口
 	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return &plugins.Result{
 			Success: false,
 			Service: "rabbitmq",
 			Error:   err,
 		}
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer conn.Close()
 
 	// 设置超时
@@ -174,6 +186,12 @@ func min(a, b int) int {
 }
 
 func (p *RabbitMQPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred plugins.Credential) bool {
+	// 检查发包限制
+	if canSend, reason := common.CanSendPacket(); !canSend {
+		common.LogError(fmt.Sprintf("RabbitMQ HTTP连接 %s:%s 受限: %s", info.Host, info.Ports, reason))
+		return false
+	}
+	
 	baseURL := fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
 	
 	client := &http.Client{
@@ -190,8 +208,10 @@ func (p *RabbitMQPlugin) testCredential(ctx context.Context, info *common.HostIn
 
 	resp, err := client.Do(req)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return false
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer resp.Body.Close()
 
 	return resp.StatusCode == 200
@@ -217,6 +237,17 @@ func (p *RabbitMQPlugin) identifyService(ctx context.Context, info *common.HostI
 // testManagementInterface 检测RabbitMQ管理界面
 func (p *RabbitMQPlugin) testManagementInterface(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("RabbitMQ管理接口 %s 受限: %s", target, reason))
+		return &plugins.Result{
+			Success: false,
+			Service: "rabbitmq",
+			Error:   fmt.Errorf("发包受限: %s", reason),
+		}
+	}
+	
 	baseURL := fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
 	
 	client := &http.Client{
@@ -234,12 +265,14 @@ func (p *RabbitMQPlugin) testManagementInterface(ctx context.Context, info *comm
 
 	resp, err := client.Do(req)
 	if err != nil {
+		common.IncrementTCPFailedPacketCount()
 		return &plugins.Result{
 			Success: false,
 			Service: "rabbitmq",
 			Error:   err,
 		}
 	}
+	common.IncrementTCPSuccessPacketCount()
 	defer resp.Body.Close()
 
 	var banner string