diff --git a/core/alive_scanner.go b/core/alive_scanner.go index ce5532f..a8a66ff 100644 --- a/core/alive_scanner.go +++ b/core/alive_scanner.go @@ -54,7 +54,7 @@ func (s *AliveScanStrategy) Description() string { } // Execute 执行存活探测扫描策略 -func (s *AliveScanStrategy) Execute(_ context.Context, session *common.ScanSession, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (s *AliveScanStrategy) Execute(ctx context.Context, session *common.ScanSession, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 验证扫描目标(需要同时检查 -h 和 -hf 参数) if info.Host == "" && session.Params.HostsFile == "" { common.LogError(i18n.GetText("parse_error_target_empty")) @@ -62,14 +62,14 @@ func (s *AliveScanStrategy) Execute(_ context.Context, session *common.ScanSessi } // 执行存活探测 - s.performAliveScan(info, session) + s.performAliveScan(ctx, info, session) // 输出统计信息 s.outputStats() } // performAliveScan 执行存活探测 -func (s *AliveScanStrategy) performAliveScan(info common.HostInfo, session *common.ScanSession) { +func (s *AliveScanStrategy) performAliveScan(ctx context.Context, info common.HostInfo, session *common.ScanSession) { // 解析目标主机 hosts, err := parsers.ParseIP(info.Host, session.Params.HostsFile, session.Params.ExcludeHosts) if err != nil { @@ -89,7 +89,7 @@ func (s *AliveScanStrategy) performAliveScan(info common.HostInfo, session *comm // 执行存活检测 - aliveList := CheckLive(hosts, false, session) // 使用ICMP探测 + aliveList := CheckLive(ctx, hosts, false, session) // 使用ICMP探测 // 更新统计信息 s.stats.AliveHosts = len(aliveList) diff --git a/core/icmp.go b/core/icmp.go index 6aee615..8f4154a 100644 --- a/core/icmp.go +++ b/core/icmp.go @@ -40,7 +40,7 @@ var pingErrorKeywords = []string{ // CheckLive 检测主机存活状态 // 支持 ICMP/Ping 探测,并在响应率过低时自动启用 TCP 补充探测 -func CheckLive(hostslist []string, Ping bool, session *common.ScanSession) []string { +func CheckLive(ctx context.Context, hostslist []string, Ping bool, session *common.ScanSession) []string { config := session.Config state := session.State // 创建局部WaitGroup @@ -71,7 +71,7 @@ func CheckLive(hostslist []string, Ping bool, session *common.ScanSession) []str // TCP 补充探测:当 ICMP/Ping 响应率过低时自动启用 // 这对防火墙过滤 ICMP 的环境特别有用 - aliveHosts = tcpSupplementaryProbe(hostslist, aliveHosts, session) + aliveHosts = tcpSupplementaryProbe(ctx, hostslist, aliveHosts, session) // 输出存活统计信息 printAliveStats(aliveHosts, hostslist) @@ -81,7 +81,7 @@ func CheckLive(hostslist []string, Ping bool, session *common.ScanSession) []str // tcpSupplementaryProbe TCP 补充探测 // 当 ICMP 响应率过低时(<10%),对未响应主机进行 TCP 探测 -func tcpSupplementaryProbe(allHosts []string, aliveHosts []string, session *common.ScanSession) []string { +func tcpSupplementaryProbe(ctx context.Context, allHosts []string, aliveHosts []string, session *common.ScanSession) []string { totalHosts := len(allHosts) if totalHosts == 0 { return aliveHosts @@ -105,7 +105,7 @@ func tcpSupplementaryProbe(allHosts []string, aliveHosts []string, session *comm common.LogInfo(i18n.Tr("tcp_probe_low_icmp_rate", fmt.Sprintf("%.1f%%", responseRate*100), len(unrespondedHosts))) // 执行 TCP 补充探测 - tcpAliveHosts := runTcpProbeForHosts(unrespondedHosts, session) + tcpAliveHosts := runTcpProbeForHosts(ctx, unrespondedHosts, session) // 合并结果 if len(tcpAliveHosts) > 0 { @@ -685,10 +685,10 @@ const tcpProbeThreshold = 0.1 // 10% // tcpProbeAlive 使用 TCP 探测主机是否存活 // 尝试连接常用端口,任一端口响应即认为存活 -func tcpProbeAlive(session *common.ScanSession, host string) bool { +func tcpProbeAlive(ctx context.Context, session *common.ScanSession, host string) bool { for _, port := range tcpProbeCommonPorts { addr := fmt.Sprintf("%s:%d", host, port) - conn, err := session.DialTCP(context.Background(), "tcp", addr, tcpProbeTimeout) + conn, err := session.DialTCP(ctx, "tcp", addr, tcpProbeTimeout) if err == nil { _ = conn.Close() return true @@ -699,7 +699,7 @@ func tcpProbeAlive(session *common.ScanSession, host string) bool { // runTcpProbeForHosts 对指定主机列表进行 TCP 补充探测 // 返回存活的主机列表 -func runTcpProbeForHosts(hosts []string, session *common.ScanSession) []string { +func runTcpProbeForHosts(ctx context.Context, hosts []string, session *common.ScanSession) []string { config := session.Config if len(hosts) == 0 { return nil @@ -726,7 +726,7 @@ func runTcpProbeForHosts(hosts []string, session *common.ScanSession) []string { wg.Done() }() - if tcpProbeAlive(session, h) { + if tcpProbeAlive(ctx, session, h) { mu.Lock() aliveHosts = append(aliveHosts, h) mu.Unlock() diff --git a/core/port_scan.go b/core/port_scan.go index 64643b4..6459b7b 100644 --- a/core/port_scan.go +++ b/core/port_scan.go @@ -111,7 +111,7 @@ func (f *failedPortCollector) Count() int { // EnhancedPortScan 高性能端口扫描函数 // 使用滑动窗口调度 + 自适应线程池 + 流式迭代器 -func EnhancedPortScan(hosts []string, ports string, timeout int64, session *common.ScanSession) []string { +func EnhancedPortScan(ctx context.Context, hosts []string, ports string, timeout int64, session *common.ScanSession) []string { config := session.Config state := session.State common.LogDebug(fmt.Sprintf("[PortScan] 开始: %d个主机, 线程数=%d", len(hosts), config.ThreadNum)) @@ -182,7 +182,7 @@ func EnhancedPortScan(hosts []string, ports string, timeout int64, session *comm }() addr := fmt.Sprintf("%s:%d", taskInfo.host, taskInfo.port) - scanSinglePort(taskInfo.host, taskInfo.port, addr, to, &count, collector, failedCollector, session) + scanSinglePort(ctx, taskInfo.host, taskInfo.port, addr, to, &count, collector, failedCollector, session) common.UpdateProgressBar(1) }, state) if err != nil { @@ -347,10 +347,10 @@ func buildServiceLogMessage(addr string, serviceInfo *ServiceInfo, isWeb bool) s } // scanSinglePort 扫描单个端口并进行服务识别(重构后的简洁版本) -func scanSinglePort(host string, port int, addr string, timeout time.Duration, count *int64, collector *resultCollector, failedCollector *failedPortCollector, session *common.ScanSession) { +func scanSinglePort(ctx context.Context, host string, port int, addr string, timeout time.Duration, count *int64, collector *resultCollector, failedCollector *failedPortCollector, session *common.ScanSession) { config := session.Config // 步骤1:建立连接 - conn, err := connectWithRetry(context.Background(), session, addr, timeout, 3) + conn, err := connectWithRetry(ctx, session, addr, timeout, 3) if err != nil { handleConnectionFailure(err, host, port, addr, failedCollector) return @@ -369,7 +369,7 @@ func scanSinglePort(host string, port int, addr string, timeout time.Duration, c if common.IsProxyEnabled() && verifyMethod != "direct" { _ = conn.Close() // 重新建立干净的连接用于服务识别 - conn, err = connectWithRetry(context.Background(), session, addr, timeout, 3) + conn, err = connectWithRetry(ctx, session, addr, timeout, 3) if err != nil { handleConnectionFailure(err, host, port, addr, failedCollector) return @@ -382,12 +382,12 @@ func scanSinglePort(host string, port int, addr string, timeout time.Duration, c saveOpenPort(host, port) // 步骤3:服务识别(Scanner负责关闭连接,包括探测中可能创建的新连接) - scanner := NewSmartPortInfoScanner(host, port, conn, timeout, config, session) + scanner := NewSmartPortInfoScanner(ctx, host, port, conn, timeout, config, session) defer scanner.Close() serviceInfo, _ := scanner.SmartIdentify() // 步骤4:处理结果 - processServiceResult(host, port, addr, serviceInfo, config) + processServiceResult(host, port, addr, serviceInfo, config, session) } // handleConnectionFailure 处理连接失败 @@ -557,10 +557,10 @@ func saveOpenPort(host string, port int) { } // processServiceResult 处理服务识别结果 -func processServiceResult(host string, port int, addr string, serviceInfo *ServiceInfo, config *common.Config) { +func processServiceResult(host string, port int, addr string, serviceInfo *ServiceInfo, config *common.Config, session *common.ScanSession) { if serviceInfo == nil { // 服务识别失败,尝试 HTTP 回退探测 - if !tryHTTPFallbackDetection(host, port, addr, config) { + if !tryHTTPFallbackDetection(host, port, addr, config, session) { common.LogInfo(i18n.Tr("port_open", addr)) } return @@ -620,10 +620,10 @@ func buildServiceDetails(port int, info *ServiceInfo) map[string]interface{} { } // tryHTTPFallbackDetection 尝试HTTP回退探测,返回是否成功识别为HTTP服务 -func tryHTTPFallbackDetection(host string, port int, addr string, config *common.Config) bool { +func tryHTTPFallbackDetection(host string, port int, addr string, config *common.Config, session *common.ScanSession) bool { // 使用WebDetection进行HTTP协议探测 webDetector := GetWebPortDetector() - if !webDetector.DetectHTTPServiceOnly(host, port, config) { + if !webDetector.DetectHTTPServiceOnly(host, port, config, session) { return false } diff --git a/core/service_probe.go b/core/service_probe.go index d93f133..521d09c 100644 --- a/core/service_probe.go +++ b/core/service_probe.go @@ -76,6 +76,7 @@ type Info struct { Conn net.Conn // 网络连接 Result Result // 探测结果 Found bool // 是否成功识别服务 + ctx context.Context // 扫描级 context config *common.Config // 配置引用 session *common.ScanSession // 会话引用 readTimeoutMS int // 当前读取超时时间(毫秒) @@ -95,7 +96,7 @@ type SmartPortInfoScanner struct { // 预定义的基础探测器已在PortFinger.go中定义,这里不再重复定义 // NewSmartPortInfoScanner 创建智能服务识别器 -func NewSmartPortInfoScanner(addr string, port int, conn net.Conn, timeout time.Duration, config *common.Config, session *common.ScanSession) *SmartPortInfoScanner { +func NewSmartPortInfoScanner(ctx context.Context, addr string, port int, conn net.Conn, timeout time.Duration, config *common.Config, session *common.ScanSession) *SmartPortInfoScanner { return &SmartPortInfoScanner{ Address: addr, Port: port, @@ -107,6 +108,7 @@ func NewSmartPortInfoScanner(addr string, port int, conn net.Conn, timeout time. Address: addr, Port: port, Conn: conn, + ctx: ctx, config: config, session: session, Result: Result{ @@ -256,7 +258,7 @@ func (s *SmartPortInfoScanner) reconnectIfNeeded() { } // 重新建立连接 - newConn, err := s.session.DialTCP(context.Background(), "tcp", fmt.Sprintf("%s:%d", s.Address, s.Port), s.Timeout) + newConn, err := s.session.DialTCP(s.info.ctx, "tcp", fmt.Sprintf("%s:%d", s.Address, s.Port), s.Timeout) if err != nil { return } @@ -516,7 +518,7 @@ func (i *Info) Write(msg []byte) error { _ = oldConn.Close() // 尝试重新连接 - 支持SOCKS5代理 - newConn, retryErr := i.session.DialTCP(context.Background(), "tcp", fmt.Sprintf("%s:%d", i.Address, i.Port), time.Duration(6)*time.Second) + newConn, retryErr := i.session.DialTCP(i.ctx, "tcp", fmt.Sprintf("%s:%d", i.Address, i.Port), time.Duration(6)*time.Second) if retryErr != nil { return retryErr } diff --git a/core/service_probe_strategy_test.go b/core/service_probe_strategy_test.go index a014dac..5ab6b77 100644 --- a/core/service_probe_strategy_test.go +++ b/core/service_probe_strategy_test.go @@ -13,6 +13,7 @@ service_probe_strategy_test.go - SmartProbeStrategy 策略逻辑测试 */ import ( + "context" "testing" "time" @@ -155,7 +156,7 @@ func TestSmartPortInfoScanner_Creation(t *testing.T) { } // 使用 nil 连接(实际测试中会使用真实连接) - scanner := NewSmartPortInfoScanner("127.0.0.1", 80, nil, 3*time.Second, config, nil) + scanner := NewSmartPortInfoScanner(context.Background(), "127.0.0.1", 80, nil, 3*time.Second, config, nil) if scanner == nil { t.Fatal("Scanner 创建失败") diff --git a/core/service_scanner.go b/core/service_scanner.go index 8355fa5..cc13a91 100644 --- a/core/service_scanner.go +++ b/core/service_scanner.go @@ -141,7 +141,7 @@ func (s *ServiceScanStrategy) Execute(ctx context.Context, session *common.ScanS // performHostScan 执行主机扫描的完整流程 func (s *ServiceScanStrategy) performHostScan(ctx context.Context, session *common.ScanSession, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 发现目标主机和端口 - targetInfos, err := s.discoverTargets(info.Host, info, session) + targetInfos, err := s.discoverTargets(ctx, info.Host, info, session) if err != nil { common.LogError(err.Error()) return @@ -156,7 +156,7 @@ func (s *ServiceScanStrategy) performHostScan(ctx context.Context, session *comm // PrepareTargets 准备目标信息 func (s *ServiceScanStrategy) PrepareTargets(info common.HostInfo, session *common.ScanSession) []common.HostInfo { // 发现目标主机和端口 - targetInfos, err := s.discoverTargets(info.Host, info, session) + targetInfos, err := s.discoverTargets(context.Background(), info.Host, info, session) if err != nil { common.LogError(err.Error()) return nil @@ -215,7 +215,7 @@ func (s *ServiceScanStrategy) LogVulnerabilityPluginInfo(targets []common.HostIn // ============================================================================= // discoverTargets 发现目标主机和端口 -func (s *ServiceScanStrategy) discoverTargets(hostInput string, baseInfo common.HostInfo, session *common.ScanSession) ([]common.HostInfo, error) { +func (s *ServiceScanStrategy) discoverTargets(ctx context.Context, hostInput string, baseInfo common.HostInfo, session *common.ScanSession) ([]common.HostInfo, error) { config := session.Config state := session.State // 标准流程:解析目标主机 @@ -230,12 +230,12 @@ func (s *ServiceScanStrategy) discoverTargets(hostInput string, baseInfo common. if len(hosts) > 0 || len(state.GetHostPorts()) > 0 { // 主机存活检测 if s.shouldPerformLivenessCheck(hosts, config) { - hosts = CheckLive(hosts, false, session) + hosts = CheckLive(ctx, hosts, false, session) common.LogInfo(i18n.Tr("alive_hosts_count_info", len(hosts))) } // 端口扫描 - alivePorts := s.discoverAlivePorts(hosts, session) + alivePorts := s.discoverAlivePorts(ctx, hosts, session) if len(alivePorts) > 0 { targetInfos = s.convertToTargetInfos(alivePorts, baseInfo) } @@ -250,7 +250,7 @@ func (s *ServiceScanStrategy) shouldPerformLivenessCheck(hosts []string, config } // discoverAlivePorts 发现存活的端口 -func (s *ServiceScanStrategy) discoverAlivePorts(hosts []string, session *common.ScanSession) []string { +func (s *ServiceScanStrategy) discoverAlivePorts(ctx context.Context, hosts []string, session *common.ScanSession) []string { config := session.Config state := session.State var alivePorts []string @@ -266,7 +266,7 @@ func (s *ServiceScanStrategy) discoverAlivePorts(hosts []string, session *common // 根据扫描模式选择端口扫描方式 if len(hosts) > 0 { - alivePorts = EnhancedPortScan(hosts, config.Target.Ports, int64(config.Timeout.Seconds()), session) + alivePorts = EnhancedPortScan(ctx, hosts, config.Target.Ports, int64(config.Timeout.Seconds()), session) } return alivePorts diff --git a/core/web_scanner.go b/core/web_scanner.go index f921cad..935e377 100644 --- a/core/web_scanner.go +++ b/core/web_scanner.go @@ -30,9 +30,9 @@ func GetWebPortDetector() *WebPortDetector { // DetectHTTPScheme 智能检测HTTP/HTTPS协议 // 策略:TLS握手优先(快速且准确),失败后尝试HTTP // 返回: "https", "http", 或 "" (都不是Web服务) -func DetectHTTPScheme(host string, port int, config *common.Config) string { +func DetectHTTPScheme(host string, port int, config *common.Config, session *common.ScanSession) string { // 优化:先快速检测 TCP 连通性 - if !isPortReachable(host, port, config) { + if !isPortReachable(host, port, config, session) { return "" } @@ -125,10 +125,10 @@ func createHTTPClient(config *common.Config) *http.Client { } // DetectHTTPServiceOnly HTTP协议检测 - 保持API兼容,简化实现 -func (w *WebPortDetector) DetectHTTPServiceOnly(host string, port int, config *common.Config) bool { +func (w *WebPortDetector) DetectHTTPServiceOnly(host string, port int, config *common.Config, session *common.ScanSession) bool { // 优化:先快速检测 TCP 连通性,避免在不可达端口上浪费双倍超时时间 // 对于不存在的端口,这可以将检测时间从 2×timeout 减少到 1×timeout - if !isPortReachable(host, port, config) { + if !isPortReachable(host, port, config, session) { return false } @@ -149,11 +149,11 @@ func (w *WebPortDetector) DetectHTTPServiceOnly(host string, port int, config *c // isPortReachable 快速检测端口是否可达(TCP 连接测试) // 用于在 HTTP/HTTPS 检测前过滤不可达端口,避免双重超时 -func isPortReachable(host string, port int, config *common.Config) bool { +func isPortReachable(host string, port int, config *common.Config, session *common.ScanSession) bool { timeout := config.Network.WebTimeout addr := net.JoinHostPort(host, strconv.Itoa(port)) - conn, err := net.DialTimeout("tcp", addr, timeout) + conn, err := session.DialTCP(context.Background(), "tcp", addr, timeout) if err != nil { return false } diff --git a/core/web_scanner_test.go b/core/web_scanner_test.go index ee67d86..ae5ae5e 100644 --- a/core/web_scanner_test.go +++ b/core/web_scanner_test.go @@ -662,6 +662,8 @@ func TestDetectHTTPScheme(t *testing.T) { cfg.Network.WebTimeout = 2 * time.Second defer func() { cfg.Network.WebTimeout = oldTimeout }() + session := common.NewScanSession(cfg, common.NewState(), common.GetFlagVars()) + t.Run("HTTPS服务器检测", func(t *testing.T) { // 创建HTTPS测试服务器 server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -677,7 +679,7 @@ func TestDetectHTTPScheme(t *testing.T) { port, _ := strconv.Atoi(portStr) // 测试检测 - result := DetectHTTPScheme(host, port, cfg) + result := DetectHTTPScheme(host, port, cfg, session) if result != "https" { t.Errorf("DetectHTTPScheme() = %q, 期望 'https'", result) } @@ -698,7 +700,7 @@ func TestDetectHTTPScheme(t *testing.T) { port, _ := strconv.Atoi(portStr) // 测试检测 - result := DetectHTTPScheme(host, port, cfg) + result := DetectHTTPScheme(host, port, cfg, session) if result != "http" { t.Errorf("DetectHTTPScheme() = %q, 期望 'http'", result) } @@ -706,7 +708,7 @@ func TestDetectHTTPScheme(t *testing.T) { t.Run("不存在的服务", func(t *testing.T) { // 使用127.0.0.1的一个未使用端口 - result := DetectHTTPScheme("127.0.0.1", 65534, cfg) + result := DetectHTTPScheme("127.0.0.1", 65534, cfg, session) if result != "" { t.Errorf("不存在的服务应返回空字符串, 实际 %q", result) } @@ -736,7 +738,7 @@ func TestDetectHTTPScheme(t *testing.T) { port, _ := strconv.Atoi(portStr) // 测试检测 - result := DetectHTTPScheme("127.0.0.1", port, cfg) + result := DetectHTTPScheme("127.0.0.1", port, cfg, session) if result != "" { t.Logf("非Web服务检测返回: %q (预期空字符串,但立即关闭连接可能被误判)", result) } @@ -757,7 +759,7 @@ func TestDetectHTTPScheme(t *testing.T) { host, portStr, _ := net.SplitHostPort(server.Listener.Addr().String()) port, _ := strconv.Atoi(portStr) - result := DetectHTTPScheme(host, port, cfg) + result := DetectHTTPScheme(host, port, cfg, session) if result != "https" { t.Errorf("TLS 1.0服务器应被检测为https, 实际 %q", result) } diff --git a/plugins/services/smb.go b/plugins/services/smb.go index 9d69b75..c39d1aa 100644 --- a/plugins/services/smb.go +++ b/plugins/services/smb.go @@ -68,7 +68,7 @@ func (p *SmbPlugin) Scan(ctx context.Context, info *common.HostInfo, session *co auth := p.getAuthenticator(smbTarget.Protocol) // 4. 未授权访问检测 - if result := p.testUnauthorizedAccess(ctx, info, auth, config, state); result != nil && result.Success { + if result := p.testUnauthorizedAccess(ctx, info, auth, config, state, session); result != nil && result.Success { var successMsg string if config.Credentials.Domain != "" { successMsg = fmt.Sprintf("SMB %s 未授权访问 - %s\\%s:%s", target, config.Credentials.Domain, result.Username, result.Password) @@ -126,7 +126,7 @@ func (p *SmbPlugin) createAuthFunc(info *common.HostInfo, auth SMBAuthenticator, } // testUnauthorizedAccess 测试未授权访问 -func (p *SmbPlugin) testUnauthorizedAccess(ctx context.Context, info *common.HostInfo, auth SMBAuthenticator, config *common.Config, state *common.State) *ScanResult { +func (p *SmbPlugin) testUnauthorizedAccess(ctx context.Context, info *common.HostInfo, auth SMBAuthenticator, config *common.Config, state *common.State, session *common.ScanSession) *ScanResult { target := info.Target() unauthorizedCreds := []Credential{ @@ -136,7 +136,7 @@ func (p *SmbPlugin) testUnauthorizedAccess(ctx context.Context, info *common.Hos } for _, cred := range unauthorizedCreds { - shareInfo, err := auth.ListShares(ctx, info.Host, info.Port, cred, config.Credentials.Domain, config.Timeout) + shareInfo, err := auth.ListShares(ctx, info.Host, info.Port, cred, config.Credentials.Domain, config.Timeout, session) if err == nil && len(shareInfo) > 0 { var output strings.Builder displayUser := cred.Username diff --git a/plugins/services/smb_protocol.go b/plugins/services/smb_protocol.go index 1131054..d4aac21 100644 --- a/plugins/services/smb_protocol.go +++ b/plugins/services/smb_protocol.go @@ -391,7 +391,7 @@ func checkSMBGhost(ctx context.Context, host string, timeout time.Duration, sess // SMBAuthenticator 统一认证接口 type SMBAuthenticator interface { Authenticate(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration, session *common.ScanSession) (*AuthResult, error) - ListShares(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration) ([]string, error) + ListShares(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration, session *common.ScanSession) ([]string, error) } // SMB1Authenticator SMB1认证器 @@ -472,8 +472,8 @@ func (a *SMB1Authenticator) Authenticate(ctx context.Context, host string, port } // ListShares 列举SMB共享(SMB1使用SMB2库列举) -func (a *SMB1Authenticator) ListShares(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration) ([]string, error) { - return listSMBSharesInternal(host, port, cred, domain, timeout) +func (a *SMB1Authenticator) ListShares(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration, session *common.ScanSession) ([]string, error) { + return listSMBSharesInternal(ctx, host, port, cred, domain, timeout, session) } // SMB2Authenticator SMB2认证器 @@ -523,15 +523,15 @@ func (a *SMB2Authenticator) Authenticate(ctx context.Context, host string, port } // ListShares 列举SMB2共享 -func (a *SMB2Authenticator) ListShares(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration) ([]string, error) { - return listSMBSharesInternal(host, port, cred, domain, timeout) +func (a *SMB2Authenticator) ListShares(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration, session *common.ScanSession) ([]string, error) { + return listSMBSharesInternal(ctx, host, port, cred, domain, timeout, session) } // listSMBSharesInternal 内部共享列举实现 -func listSMBSharesInternal(host string, port int, cred Credential, domain string, timeout time.Duration) ([]string, error) { +func listSMBSharesInternal(ctx context.Context, host string, port int, cred Credential, domain string, timeout time.Duration, session *common.ScanSession) ([]string, error) { target := net.JoinHostPort(host, strconv.Itoa(port)) - conn, err := net.DialTimeout("tcp", target, timeout*2) + conn, err := session.DialTCP(ctx, "tcp", target, timeout*2) if err != nil { return nil, err } diff --git a/plugins/web/webtitle.go b/plugins/web/webtitle.go index 599aab5..410c815 100644 --- a/plugins/web/webtitle.go +++ b/plugins/web/webtitle.go @@ -41,7 +41,7 @@ func NewWebTitlePlugin() *WebTitlePlugin { // Scan 执行WebTitle扫描 func (p *WebTitlePlugin) Scan(ctx context.Context, info *common.HostInfo, session *common.ScanSession) *WebScanResult { config := session.Config - title, status, length, server, fingerprints, url, err := p.getWebTitle(ctx, info, config) + title, status, length, server, fingerprints, url, err := p.getWebTitle(ctx, info, config, session) if err != nil { return &WebScanResult{ Success: false, @@ -79,9 +79,9 @@ func (p *WebTitlePlugin) Scan(ctx context.Context, info *common.HostInfo, sessio } } -func (p *WebTitlePlugin) getWebTitle(ctx context.Context, info *common.HostInfo, config *common.Config) (string, int, int, string, []string, string, error) { +func (p *WebTitlePlugin) getWebTitle(ctx context.Context, info *common.HostInfo, config *common.Config, session *common.ScanSession) (string, int, int, string, []string, string, error) { // 智能协议检测 - protocol := p.detectProtocol(info, config) + protocol := p.detectProtocol(info, config, session) baseURL := fmt.Sprintf("%s://%s:%d", protocol, info.Host, info.Port) // 构建显示用URL(隐藏标准端口) @@ -235,7 +235,7 @@ func (p *WebTitlePlugin) formatHeaders(headers http.Header) string { } // detectProtocol 智能检测HTTP/HTTPS协议(基于服务识别和主动探测) -func (p *WebTitlePlugin) detectProtocol(info *common.HostInfo, config *common.Config) string { +func (p *WebTitlePlugin) detectProtocol(info *common.HostInfo, config *common.Config, session *common.ScanSession) string { host := info.Host port := info.Port @@ -262,7 +262,7 @@ func (p *WebTitlePlugin) detectProtocol(info *common.HostInfo, config *common.Co // 第三优先级:主动协议检测(TLS握手) // 对于-u模式或服务名为普通"http"的情况,进行主动检测确认 - detected := core.DetectHTTPScheme(host, port, config) + detected := core.DetectHTTPScheme(host, port, config, session) if detected != "" { // 缓存检测结果(避免重复检测) if exists { diff --git a/webscan/lib/Eval.go b/webscan/lib/Eval.go index 48cf091..8235738 100644 --- a/webscan/lib/Eval.go +++ b/webscan/lib/Eval.go @@ -12,6 +12,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/google/cel-go/cel" @@ -31,15 +32,13 @@ var ( baseProgramOpt []cel.ProgramOption ) -// 包级POC配置 -var ( - pocDNSLog bool // DNSLog配置缓存 -) +// 包级POC配置(atomic 保证并发安全) +var pocDNSLog atomic.Bool // InitPOCConfig 初始化POC配置(在扫描开始前调用) // 这样CEL回调函数可以使用包级变量而非GetGlobalConfig func InitPOCConfig(dnsLog bool) { - pocDNSLog = dnsLog + pocDNSLog.Store(dnsLog) } // NewEnv 创建一个新的 CEL 环境(使用缓存避免重复注册函数) @@ -349,7 +348,7 @@ func randomString(n int) string { // 使用包级pocDNSLog变量,由InitPOCConfig初始化 func reverseCheck(r *Reverse, timeout int64) bool { // 检查必要条件(使用包级配置变量) - if ceyeAPI == "" || r.Domain == "" || !pocDNSLog { + if ceyeAPI == "" || r.Domain == "" || !pocDNSLog.Load() { return false }