refactor: 为SMBGhost插件添加包控制功能

- 在detectVulnerability和gatherSystemInfo方法中添加包控制检查
- 统一错误处理和包计数逻辑
- 确保所有TCP连接遵循发包限制
This commit is contained in:
ZacharyZcR 2025-09-02 12:03:27 +00:00
parent 5bf3f7a2d9
commit 2385a730eb

View File

@ -156,11 +156,18 @@ func (p *SmbGhostPlugin) detectVulnerability(ctx context.Context, info *common.H
addr := fmt.Sprintf("%s:445", info.Host) addr := fmt.Sprintf("%s:445", info.Host)
timeout := time.Duration(common.Timeout) * time.Second timeout := time.Duration(common.Timeout) * time.Second
// 检查发包限制
if canSend, reason := common.CanSendPacket(); !canSend {
return false, fmt.Errorf("发包受限: %s", reason)
}
// 建立TCP连接 // 建立TCP连接
conn, err := common.WrapperTcpWithTimeout("tcp", addr, timeout) conn, err := common.WrapperTcpWithTimeout("tcp", addr, timeout)
if err != nil { if err != nil {
common.IncrementTCPFailedPacketCount()
return false, fmt.Errorf("连接失败: %v", err) return false, fmt.Errorf("连接失败: %v", err)
} }
common.IncrementTCPSuccessPacketCount()
defer conn.Close() defer conn.Close()
// 设置连接超时 // 设置连接超时
@ -200,10 +207,17 @@ func (p *SmbGhostPlugin) gatherSystemInfo(ctx context.Context, info *common.Host
addr := fmt.Sprintf("%s:445", info.Host) addr := fmt.Sprintf("%s:445", info.Host)
timeout := time.Duration(common.Timeout) * time.Second timeout := time.Duration(common.Timeout) * time.Second
conn, err := common.WrapperTcpWithTimeout("tcp", addr, timeout) // 检查发包限制
if err != nil { if canSend, _ := common.CanSendPacket(); !canSend {
return "" return ""
} }
conn, err := common.WrapperTcpWithTimeout("tcp", addr, timeout)
if err != nil {
common.IncrementTCPFailedPacketCount()
return ""
}
common.IncrementTCPSuccessPacketCount()
defer conn.Close() defer conn.Close()
// 发送标准SMB协商请求获取系统信息 // 发送标准SMB协商请求获取系统信息