diff --git a/core/alive_scanner.go b/core/alive_scanner.go index a5faa24..47999be 100644 --- a/core/alive_scanner.go +++ b/core/alive_scanner.go @@ -1,6 +1,7 @@ package core import ( + "context" "fmt" "sync" "time" @@ -53,7 +54,7 @@ func (s *AliveScanStrategy) Description() string { } // Execute 执行存活探测扫描策略 -func (s *AliveScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (s *AliveScanStrategy) Execute(_ context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 验证扫描目标(需要同时检查 -h 和 -hf 参数) fv := common.GetFlagVars() if info.Host == "" && fv.HostsFile == "" { diff --git a/core/local_scanner.go b/core/local_scanner.go index 08ec528..ee67a2b 100644 --- a/core/local_scanner.go +++ b/core/local_scanner.go @@ -1,6 +1,7 @@ package core import ( + "context" "sync" "github.com/shadow1ng/fscan/common" @@ -41,7 +42,7 @@ func (s *LocalScanStrategy) Description() string { } // Execute 执行本地扫描策略 -func (s *LocalScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (s *LocalScanStrategy) Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 输出扫描开始信息 s.LogScanStart() @@ -66,7 +67,7 @@ func (s *LocalScanStrategy) Execute(config *common.Config, state *common.State, targets := s.PrepareTargets(info) // 执行扫描任务 - ExecuteScanTasks(config, state, targets, s, ch, wg) + ExecuteScanTasks(ctx, config, state, targets, s, ch, wg) } // PrepareTargets 准备本地扫描目标 diff --git a/core/scanner.go b/core/scanner.go index 09e8bc4..e19c17d 100644 --- a/core/scanner.go +++ b/core/scanner.go @@ -18,7 +18,7 @@ import ( // ScanStrategy 定义扫描策略接口 type ScanStrategy interface { - Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) + Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) GetPlugins(config *common.Config) ([]string, bool) IsPluginApplicableByName(pluginName string, targetHost string, targetPort int, isCustomMode bool, config *common.Config) bool } @@ -73,7 +73,10 @@ func selectStrategy(config *common.Config, state *common.State, info common.Host } // RunScan 执行整体扫描流程 -func RunScan(info common.HostInfo, config *common.Config, state *common.State) { +func RunScan(ctx context.Context, info common.HostInfo, config *common.Config, state *common.State) { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + // 初始化HTTP客户端(静默,无需日志) if err := lib.Inithttp(config); err != nil { common.LogError(i18n.Tr("http_client_init_failed", err)) @@ -88,7 +91,7 @@ func RunScan(info common.HostInfo, config *common.Config, state *common.State) { wg := sync.WaitGroup{} // 执行策略 - strategy.Execute(config, state, info, ch, &wg) + strategy.Execute(ctx, config, state, info, ch, &wg) // 等待所有扫描完成 wg.Wait() @@ -111,6 +114,8 @@ func RunScan(info common.HostInfo, config *common.Config, state *common.State) { signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) <-sigChan common.LogInfo(i18n.GetText("received_exit_signal")) + cancel() + time.Sleep(500 * time.Millisecond) } // 完成扫描 @@ -134,7 +139,7 @@ func finishScan(config *common.Config, state *common.State) { } // ExecuteScanTasks 任务执行通用框架 -func ExecuteScanTasks(config *common.Config, state *common.State, targets []common.HostInfo, strategy ScanStrategy, ch chan struct{}, wg *sync.WaitGroup) { +func ExecuteScanTasks(ctx context.Context, config *common.Config, state *common.State, targets []common.HostInfo, strategy ScanStrategy, ch chan struct{}, wg *sync.WaitGroup) { // 获取要执行的插件 pluginsToRun, isCustomMode := strategy.GetPlugins(config) @@ -149,6 +154,13 @@ func ExecuteScanTasks(config *common.Config, state *common.State, targets []comm // 流式执行任务,避免预构建大量任务对象 for _, target := range targets { + // 检查取消 + select { + case <-ctx.Done(): + return + default: + } + targetPort := target.Port for _, pluginName := range pluginsToRun { @@ -159,7 +171,7 @@ func ExecuteScanTasks(config *common.Config, state *common.State, targets []comm // 检查插件是否适用于当前目标 if strategy.IsPluginApplicableByName(pluginName, target.Host, targetPort, isCustomMode, config) { - executeScanTask(config, state, pluginName, target, ch, wg) + executeScanTask(ctx, config, state, pluginName, target, ch, wg) } } } @@ -182,10 +194,42 @@ func countApplicableTasks(targets []common.HostInfo, pluginsToRun []string, isCu return count } +// longRunningPlugins 长驻插件,不加入 scan WaitGroup,通过 ctx 取消退出 +var longRunningPlugins = map[string]bool{ + "forwardshell": true, + "socks5proxy": true, + "reverseshell": true, +} + // executeScanTask 执行单个扫描任务 -func executeScanTask(config *common.Config, state *common.State, pluginName string, target common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func executeScanTask(ctx context.Context, config *common.Config, state *common.State, pluginName string, target common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { + // 检查取消 + select { + case <-ctx.Done(): + return + default: + } + + // 长驻插件不进 WaitGroup,通过 ctx 管理生命周期 + if longRunningPlugins[pluginName] { + go func() { + plugin := plugins.Get(pluginName) + if plugin != nil { + plugin.Scan(ctx, &target, config, state) + } + }() + return + } + wg.Add(1) - ch <- struct{}{} // 获取并发槽位 + + // 获取并发槽位,支持取消 + select { + case ch <- struct{}{}: + case <-ctx.Done(): + wg.Done() + return + } go func() { // 开始监控插件任务 @@ -210,7 +254,7 @@ func executeScanTask(config *common.Config, state *common.State, pluginName stri plugin := plugins.Get(pluginName) if plugin != nil { - result := plugin.Scan(context.Background(), &target, config, state) + result := plugin.Scan(ctx, &target, config, state) if result != nil { if result.Success { // 保存成功的扫描结果到文件 diff --git a/core/scanner_test.go b/core/scanner_test.go index 163bcba..a3fbf78 100644 --- a/core/scanner_test.go +++ b/core/scanner_test.go @@ -1,6 +1,7 @@ package core import ( + "context" "fmt" "sync" "testing" @@ -214,7 +215,7 @@ type mockStrategy struct { applicablePlugins map[string]bool // pluginName -> isApplicable } -func (m *mockStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (m *mockStrategy) Execute(_ context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { } func (m *mockStrategy) GetPlugins() ([]string, bool) { diff --git a/core/service_scanner.go b/core/service_scanner.go index 67d2eac..113b4c5 100644 --- a/core/service_scanner.go +++ b/core/service_scanner.go @@ -1,6 +1,7 @@ package core import ( + "context" "fmt" "strconv" "strings" @@ -112,7 +113,7 @@ func (s *ServiceScanStrategy) Description() string { } // Execute 执行服务扫描策略 -func (s *ServiceScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (s *ServiceScanStrategy) Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 验证扫描目标(需要同时检查 -h 和 -hf 参数) fv := common.GetFlagVars() if info.Host == "" && fv.HostsFile == "" { @@ -133,11 +134,11 @@ func (s *ServiceScanStrategy) Execute(config *common.Config, state *common.State s.LogPluginInfo(config) // 执行主机扫描流程 - s.performHostScan(config, state, info, ch, wg) + s.performHostScan(ctx, config, state, info, ch, wg) } // performHostScan 执行主机扫描的完整流程 -func (s *ServiceScanStrategy) performHostScan(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (s *ServiceScanStrategy) performHostScan(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 发现目标主机和端口 targetInfos, err := s.discoverTargets(info.Host, info, config, state) if err != nil { @@ -147,7 +148,7 @@ func (s *ServiceScanStrategy) performHostScan(config *common.Config, state *comm // 执行漏洞扫描 if len(targetInfos) > 0 { - ExecuteScanTasks(config, state, targetInfos, s, ch, wg) + ExecuteScanTasks(ctx, config, state, targetInfos, s, ch, wg) } } diff --git a/core/web_scanner.go b/core/web_scanner.go index 3e4f44d..bd63869 100644 --- a/core/web_scanner.go +++ b/core/web_scanner.go @@ -1,6 +1,7 @@ package core import ( + "context" "crypto/tls" "fmt" "net" @@ -303,7 +304,7 @@ func (s *WebScanStrategy) Description() string { } // Execute 执行Web扫描策略 -func (s *WebScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { +func (s *WebScanStrategy) Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { // 输出扫描开始信息 s.LogScanStart() @@ -320,7 +321,7 @@ func (s *WebScanStrategy) Execute(config *common.Config, state *common.State, in s.LogPluginInfo(config) // 执行扫描任务 - ExecuteScanTasks(config, state, targets, s, ch, wg) + ExecuteScanTasks(ctx, config, state, targets, s, ch, wg) } // PrepareTargets 准备URL目标列表 diff --git a/main.go b/main.go index b145efc..8f391d1 100644 --- a/main.go +++ b/main.go @@ -1,6 +1,7 @@ package main import ( + "context" "os" "os/signal" "syscall" @@ -66,5 +67,5 @@ func main() { defer func() { _ = common.Cleanup() }() // 执行扫描 - core.RunScan(*result.Info, result.Config, result.State) + core.RunScan(context.Background(), *result.Info, result.Config, result.State) } diff --git a/web/api/scan.go b/web/api/scan.go index 1b50f73..8634cc7 100644 --- a/web/api/scan.go +++ b/web/api/scan.go @@ -3,6 +3,7 @@ package api import ( + "context" "encoding/json" "net/http" "sync" @@ -73,7 +74,7 @@ type ScanHandler struct { hub *ws.Hub state int32 startTime time.Time - stopChan chan struct{} + cancelFn context.CancelFunc mu sync.RWMutex results *ResultStore } @@ -122,7 +123,6 @@ func (h *ScanHandler) Start(w http.ResponseWriter, r *http.Request) { h.mu.Lock() h.startTime = time.Now() - h.stopChan = make(chan struct{}) h.mu.Unlock() // 清空旧结果 @@ -145,8 +145,18 @@ func (h *ScanHandler) Start(w http.ResponseWriter, r *http.Request) { // runScan 执行扫描 func (h *ScanHandler) runScan(req ScanRequest) { + ctx, cancel := context.WithCancel(context.Background()) + + h.mu.Lock() + h.cancelFn = cancel + h.mu.Unlock() + defer func() { - common.ClearResultCallback() // 清除回调 + cancel() + h.mu.Lock() + h.cancelFn = nil + h.mu.Unlock() + common.ClearResultCallback() atomic.StoreInt32(&h.state, int32(ScanStateIdle)) h.hub.Broadcast(ws.MsgScanCompleted, map[string]interface{}{ "duration": time.Since(h.startTime).Seconds(), @@ -209,7 +219,7 @@ func (h *ScanHandler) runScan(req ScanRequest) { }) // 执行扫描 - core.RunScan(info, config, state) + core.RunScan(ctx, info, config, state) } // Stop 停止扫描 @@ -229,8 +239,8 @@ func (h *ScanHandler) Stop(w http.ResponseWriter, r *http.Request) { atomic.StoreInt32(&h.state, int32(ScanStateStopping)) h.mu.Lock() - if h.stopChan != nil { - close(h.stopChan) + if h.cancelFn != nil { + h.cancelFn() } h.mu.Unlock()