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