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

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

View File

@ -105,6 +105,11 @@ func (p *Smb2Plugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResul
// testCredential 测试单个凭据 - 简化的SMB2协议检测
func (p *Smb2Plugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
// 检查发包限制
if canSend, _ := common.CanSendPacket(); !canSend {
return false
}
// 基于TCP连接的简单SMB2检测
timeoutCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
defer cancel()
@ -112,8 +117,10 @@ func (p *Smb2Plugin) testCredential(ctx context.Context, info *common.HostInfo,
dialer := &net.Dialer{}
conn, err := dialer.DialContext(timeoutCtx, "tcp", fmt.Sprintf("%s:445", info.Host))
if err != nil {
common.IncrementTCPFailedPacketCount()
return false
}
common.IncrementTCPSuccessPacketCount()
defer conn.Close()
// 发送SMB2协商请求
@ -153,10 +160,17 @@ func (p *Smb2Plugin) testCredential(ctx context.Context, info *common.HostInfo,
// detectProtocol 检测SMB2协议信息
func (p *Smb2Plugin) detectProtocol(ctx context.Context, info *common.HostInfo) string {
conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:445", info.Host), 5*time.Second)
if err != nil {
// 检查发包限制
if canSend, _ := common.CanSendPacket(); !canSend {
return ""
}
conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:445", info.Host), 5*time.Second)
if err != nil {
common.IncrementTCPFailedPacketCount()
return ""
}
common.IncrementTCPSuccessPacketCount()
defer conn.Close()
// 发送SMB2协商请求获取版本信息