refactor: 重构Memcached和RabbitMQ插件使用统一发包控制

- 修改Memcached插件,在TCP连接和服务识别中添加发包控制
- 修改RabbitMQ插件,在AMQP连接、HTTP连接和管理接口中添加发包控制
- 统一包计数逻辑,确保TCP连接成功和失败都正确计数
- 保持现有缓存服务和消息队列检测功能
This commit is contained in:
ZacharyZcR 2025-09-02 11:52:59 +00:00
parent 98a9a4e1c2
commit 36f0e5076d
2 changed files with 52 additions and 0 deletions

View File

@ -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 { func (p *MemcachedPlugin) connectToMemcached(ctx context.Context, info *common.HostInfo) net.Conn {
target := fmt.Sprintf("%s:%s", info.Host, info.Ports) 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 timeout := time.Duration(common.Timeout) * time.Second
conn, err := net.DialTimeout("tcp", target, timeout) conn, err := net.DialTimeout("tcp", target, timeout)
if err != nil { if err != nil {
common.IncrementTCPFailedPacketCount()
return nil return nil
} }
common.IncrementTCPSuccessPacketCount()
conn.SetDeadline(time.Now().Add(timeout)) 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 { func (p *MemcachedPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
target := fmt.Sprintf("%s:%s", info.Host, info.Ports) 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) conn := p.connectToMemcached(ctx, info)
if conn == nil { if conn == nil {
return &ScanResult{ return &ScanResult{

View File

@ -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 { func (p *RabbitMQPlugin) testAMQPProtocol(ctx context.Context, info *common.HostInfo) *plugins.Result {
target := fmt.Sprintf("%s:%s", info.Host, info.Ports) 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端口 // 连接到AMQP端口
conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second) conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
if err != nil { if err != nil {
common.IncrementTCPFailedPacketCount()
return &plugins.Result{ return &plugins.Result{
Success: false, Success: false,
Service: "rabbitmq", Service: "rabbitmq",
Error: err, Error: err,
} }
} }
common.IncrementTCPSuccessPacketCount()
defer conn.Close() 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 { 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) baseURL := fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
client := &http.Client{ client := &http.Client{
@ -190,8 +208,10 @@ func (p *RabbitMQPlugin) testCredential(ctx context.Context, info *common.HostIn
resp, err := client.Do(req) resp, err := client.Do(req)
if err != nil { if err != nil {
common.IncrementTCPFailedPacketCount()
return false return false
} }
common.IncrementTCPSuccessPacketCount()
defer resp.Body.Close() defer resp.Body.Close()
return resp.StatusCode == 200 return resp.StatusCode == 200
@ -217,6 +237,17 @@ func (p *RabbitMQPlugin) identifyService(ctx context.Context, info *common.HostI
// testManagementInterface 检测RabbitMQ管理界面 // testManagementInterface 检测RabbitMQ管理界面
func (p *RabbitMQPlugin) testManagementInterface(ctx context.Context, info *common.HostInfo) *plugins.Result { func (p *RabbitMQPlugin) testManagementInterface(ctx context.Context, info *common.HostInfo) *plugins.Result {
target := fmt.Sprintf("%s:%s", info.Host, info.Ports) 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) baseURL := fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
client := &http.Client{ client := &http.Client{
@ -234,12 +265,14 @@ func (p *RabbitMQPlugin) testManagementInterface(ctx context.Context, info *comm
resp, err := client.Do(req) resp, err := client.Do(req)
if err != nil { if err != nil {
common.IncrementTCPFailedPacketCount()
return &plugins.Result{ return &plugins.Result{
Success: false, Success: false,
Service: "rabbitmq", Service: "rabbitmq",
Error: err, Error: err,
} }
} }
common.IncrementTCPSuccessPacketCount()
defer resp.Body.Close() defer resp.Body.Close()
var banner string var banner string