From 65ff74cc5efc39eea88e686dc9c1d8fb55cc57d3 Mon Sep 17 00:00:00 2001 From: zhizhuo Date: Tue, 11 Aug 2026 20:38:17 +0800 Subject: [PATCH 1/5] =?UTF-8?q?=E9=87=8D=E6=9E=84=20SDK=20=E8=87=B3=20pkg/?= =?UTF-8?q?sdk=EF=BC=8C=E8=A1=A5=E9=BD=90=E8=83=BD=E5=8A=9B=E5=B9=B6?= =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=BC=95=E6=93=8E=E6=B3=84=E6=BC=8F=E4=B8=8E?= =?UTF-8?q?=E7=AB=9E=E6=80=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 把根目录的 afrog.go 迁移为独立的 pkg/sdk 包，改用函数式选项， 让 SDK 可以安全地嵌入长驻进程。 SDK - New(ctx, opts...) 函数式选项，context 贯穿生命周期 - 哨兵错误、Close() 幂等释放、channel 流式订阅 - 默认零控制台输出，构造函数不再创建目录或改写用户配置 - 结果视图可直接 JSON 序列化，携带完整请求/响应报文 - 新增 WithCheckpoint(断点续扫)、WithCyberspace(空间测绘取目标)、 WithTargetPreProbe、WithExecutionMonitor、WithTaskTimeout， 以及 OOB 轮询间隔与命中保留时长 PoC 输入 - PocFile/PocPaths/AppendPoc 合并解析，修复同时指定 -P 与 -ap 时 -ap 被静默丢弃的问题 - 支持 glob 通配，加载失败原因通过 PocDiagnostics 返回而非只打印 引擎 - OOB 轮询协程不再在扫描结束后存活，新增 Runner.Release 释放资源 - ticker 与 oobMgr 改为原子指针，消除 Stop 与调度之间的竞态 - 极高速率下 ticker 间隔为 0 导致 panic 的问题 - 新增 OnFailure 回调，PoC 执行失败不再被静默吞掉 - 响应体被 MaxRespBodySize 截断时标记 BodyTruncated - 端口扫描改用 net.JoinHostPort，修复 IPv6 地址拼接 - HexDecode 遇到非法输入不再 log.Fatal 终止宿主进程 Web - 任务状态改由互斥量保护，finalizeTask 保证只执行一次， 修复取消扫描时并发计数被多减、队列超额准入的问题 - 扫描结束后释放 scanner，避免长驻服务累积引擎与轮询协程 同步更新中英文 SDK 文档与全部示例，新增 55 个 SDK 测试、 3 个 web 并发测试，全量包通过 -race。 --- .gitignore | 1 + README.md | 66 +- afrog.go | 1376 ----------------- cmd/sdk/main.go | 69 +- docs/SDK_Usage_Guide_English.md | 1075 +++++++------ ...7\345\215\227_\344\270\255\346\226\207.md" | 1073 +++++++------ examples/README.md | 69 + examples/async_scan/main.go | 349 ++--- examples/basic_scan/main.go | 118 +- examples/full_output/main.go | 127 ++ examples/internal/examplepath/examplepath.go | 41 + examples/oob_scan/main.go | 324 ++-- examples/port_scan/main.go | 82 +- examples/progress_scan/main.go | 221 +-- examples/sdk_portscan/main.go | 147 +- examples/vuln_scan/main.go | 114 +- pkg/config/afrogupdate.go | 96 +- pkg/config/config.go | 62 + pkg/config/oobadapter.go | 10 +- pkg/config/options.go | 62 +- pkg/config/options_legacy_oob_test.go | 1 - pkg/config/pocmerge_test.go | 104 ++ pkg/config/pocsource.go | 165 ++ pkg/config/pocsource_test.go | 258 ++++ pkg/cyberspace/cyberspace.go | 16 +- pkg/portscan/scan.go | 4 +- pkg/protocols/http/retryhttpclient/client.go | 13 + pkg/result/result.go | 4 + pkg/runner/checker.go | 17 +- pkg/runner/engine.go | 113 +- pkg/runner/oob_manager.go | 25 + pkg/runner/oob_resolver.go | 16 +- pkg/runner/oob_resolver_test.go | 6 +- pkg/runner/runner.go | 36 +- pkg/sdk/doc.go | 109 ++ pkg/sdk/errors.go | 50 + pkg/sdk/event.go | 388 +++++ pkg/sdk/event_test.go | 84 + pkg/sdk/lifecycle_test.go | 158 ++ pkg/sdk/monitor_test.go | 216 +++ pkg/sdk/options.go | 1099 +++++++++++++ pkg/sdk/options_test.go | 217 +++ pkg/sdk/scanner.go | 1345 ++++++++++++++++ pkg/sdk/sdk_test.go | 720 +++++++++ pkg/sdk/stream.go | 99 ++ pkg/sdk/targetsource_test.go | 167 ++ pkg/utils/utils.go | 13 +- pkg/web/handlers.go | 41 +- pkg/web/scans.go | 236 +-- pkg/web/scans_test.go | 97 ++ pkg/web/server_routes copy.go | 197 --- 51 files changed, 7945 insertions(+), 3551 deletions(-) delete mode 100644 afrog.go create mode 100644 examples/README.md create mode 100644 examples/full_output/main.go create mode 100644 examples/internal/examplepath/examplepath.go create mode 100644 pkg/config/pocmerge_test.go create mode 100644 pkg/config/pocsource.go create mode 100644 pkg/config/pocsource_test.go create mode 100644 pkg/sdk/doc.go create mode 100644 pkg/sdk/errors.go create mode 100644 pkg/sdk/event.go create mode 100644 pkg/sdk/event_test.go create mode 100644 pkg/sdk/lifecycle_test.go create mode 100644 pkg/sdk/monitor_test.go create mode 100644 pkg/sdk/options.go create mode 100644 pkg/sdk/options_test.go create mode 100644 pkg/sdk/scanner.go create mode 100644 pkg/sdk/sdk_test.go create mode 100644 pkg/sdk/stream.go create mode 100644 pkg/sdk/targetsource_test.go create mode 100644 pkg/web/scans_test.go delete mode 100644 pkg/web/server_routes copy.go diff --git a/.gitignore b/.gitignore index 8da7b3f22..79d5edce0 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,7 @@ urls* .vscode .idea .trae +.cursor .DS_Store .config diff --git a/README.md b/README.md index 7b93ff80f..e26d54333 100644 --- a/README.md +++ b/README.md @@ -709,20 +709,72 @@ afrog -t https://example.com -ja result.json ## As Library -### Simple Example +```go +package main + +import ( + "context" + "fmt" + "log" + + "github.com/zan8in/afrog/v3/pkg/sdk" +) + +func main() { + ctx := context.Background() + + scanner, err := sdk.New(ctx, + sdk.WithTargets("https://example.com"), + sdk.WithPocPaths("./pocs/afrog-pocs"), // file, directory or glob + ) + if err != nil { + log.Fatal(err) + } + defer scanner.Close() + + if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) + } + + for _, r := range scanner.Results() { + fmt.Printf("[%s] %s - %s\n", r.Severity, r.FullTarget, r.PocName) + + // The complete request and response of every step is available. + for _, ex := range r.Exchanges { + fmt.Println(ex.Request) + fmt.Println(ex.Response) + } + } +} +``` + +`Execute` runs synchronously; `Start` plus `Wait`/`Done` runs the scan in the +background. The SDK writes nothing to stdout or stderr, returns typed errors +that work with `errors.Is`, and `Results()` is JSON serialisable as-is. For comprehensive SDK documentation: - [SDK Usage Guide (English)](docs/SDK_Usage_Guide_English.md) - [SDK使用指南 (中文)](docs/SDK使用指南_中文.md) -### More Examples & Documentation +### Runnable Examples + +Every example resolves the bundled PoC directory automatically and accepts +`-pocs` to override it: + +```sh +go run ./examples/basic_scan +go run ./examples/full_output -json +``` -- [Basic scanner](examples/basic_scan/main.go) -- [Async scanner](examples/async_scan/main.go) -- [OOB scanner](examples/oob_scan/main.go) -- [Progress scanner](examples/progress_scan/main.go) -- [SDK PortScan (sync/async)](examples/sdk_portscan/main.go) +- [Basic scanner](examples/basic_scan/main.go) — smallest useful program +- [Full output](examples/full_output/main.go) — complete request/response data and JSON +- [Async scanner](examples/async_scan/main.go) — streaming results and progress +- [Progress scanner](examples/progress_scan/main.go) — progress bar +- [OOB scanner](examples/oob_scan/main.go) — out-of-band detection +- [SDK PortScan (sync/async)](examples/sdk_portscan/main.go) — port pre-scanning +- [Vulnerability scan](examples/vuln_scan/main.go) — CI-style streaming consumption +- [Port scan](examples/port_scan/main.go) — using the portscan package directly diff --git a/afrog.go b/afrog.go deleted file mode 100644 index b22cf928b..000000000 --- a/afrog.go +++ /dev/null @@ -1,1376 +0,0 @@ -package afrog - -import ( - "context" - "errors" - "fmt" - "os" - "strings" - "sync" - "sync/atomic" - "time" - - "github.com/zan8in/oobadapter/pkg/oobadapter" - - "github.com/zan8in/afrog/v3/pkg/catalog" - "github.com/zan8in/afrog/v3/pkg/config" - "github.com/zan8in/afrog/v3/pkg/curated/service" - "github.com/zan8in/afrog/v3/pkg/fingerprint" - "github.com/zan8in/afrog/v3/pkg/poc" - "github.com/zan8in/afrog/v3/pkg/protocols/http/retryhttpclient" - "github.com/zan8in/afrog/v3/pkg/result" - "github.com/zan8in/afrog/v3/pkg/runner" - "github.com/zan8in/afrog/v3/pkg/targets" - "github.com/zan8in/afrog/v3/pkg/utils" - "github.com/zan8in/afrog/v3/pocs" -) - -// SDKScanner SDK版本的扫描器，专为库调用优化 -type SDKScanner struct { - // runner 内部扫描引擎实例 - runner *runner.Runner - - // results 存储所有扫描结果 - results []*result.Result - - // mu 用于保护results的并发访问 - mu sync.Mutex - - openPortsMu sync.Mutex - openPorts map[string]map[int]struct{} - - // options 存储扫描配置选项 - options *config.Options - - // sdkOpts 保存原始SDK配置选项 - sdkOpts *SDKOptions - - // 实时结果回调（同步版本） - OnResult func(*result.Result) - - OnPort func(host string, port int) - - OnWebProbe func(r WebProbeResult) - - // 实时结果通道（流式版本） - ResultChan chan *result.Result - - PortChan chan PortScanResult - - HostChan chan HostDiscoveryResult - - WebProbeChan chan WebProbeResult - - PhaseProgressChan chan PhaseProgress - - ScanInfoChan chan ScanInfoUpdate - - closeChansOnce sync.Once - - runStartOnce sync.Once - runDoneOnce sync.Once - runStarted chan struct{} - runDone chan struct{} - - // 控制流式输出的context - ctx context.Context - cancel context.CancelFunc - - // 扫描统计信息 - stats *ScanStats - - phaseMu sync.Mutex - phases map[string]PhaseProgress - - lastVulnPhasePercent int32 -} - -// ScanStats 扫描统计信息 -type ScanStats struct { - StartTime time.Time - EndTime time.Time - TotalTargets int - TotalPocs int - TotalScans int - CompletedScans int32 - FoundVulns int32 -} - -type PortScanResult struct { - Host string - Port int -} - -type HostDiscoveryResult struct { - Host string -} - -type WebProbeResult struct { - URL string - Title string - Server string - PoweredBy string -} - -type PhaseProgress struct { - Phase string - Status string - Finished int64 - Total int64 - Percent int -} - -type ScanInfoUpdate struct { - TotalTargets int - Targets []string - TotalPocs int - TotalScans int - OOBEnabled bool - OOBStatus string -} - -// SDKOptions SDK扫描配置选项（优化版） -type SDKOptions struct { - // ========== 目标配置 ========== - Targets []string // 扫描目标列表 - TargetsFile string // 目标文件路径 - - // ========== POC配置 ========== - PocFile string // POC文件或目录路径（必须） - AppendPoc []string // 附加POC文件或目录路径 - Search string // POC搜索关键词 - Severity string // 严重程度过滤 - ExcludePocs []string - ExcludePocsFile string - - // ========== 性能配置 ========== - RateLimit int // 请求速率限制 (默认: 150) - ReqLimitPerTarget int - AutoReqLimit bool - Polite bool - Balanced bool - Aggressive bool - Concurrency int // 并发数 (默认: 25) - Retries int // 重试次数 (默认: 1) - Timeout int // 超时时间秒 (默认: 10) - MaxHostError int // 主机最大错误数 (默认: 3) - Smart bool - DisableFingerprint bool - EnableWebProbe bool - FingerprintFilterMode string - MaxRespBodySize int - BruteMaxRequests int - DefaultAccept bool - VulnerabilityScannerBreakpoint bool - - PortScan bool - PSPorts string - PSRateLimit int - PSTimeout int - PSRetries int - PSSkipDiscovery bool - PSS4Chunk int - - // ========== 网络配置 ========== - Proxy string // HTTP/SOCKS5代理 - Headers []string - - // ========== OOB配置 ========== - EnableOOB bool // 是否启用OOB检测 (默认: false) - OOB string // OOB适配器类型: ceyeio, dnslogcn, alphalog, xray, revsuit - OOBKey string // OOB API密钥 - OOBDomain string // OOB域名 - OOBApiUrl string // OOB API地址 - OOBHttpUrl string // OOB HTTP地址 - OOBRateLimit int - OOBConcurrency int - OOBFinalizeTimeout int - - // ========== 输出配置 ========== - EnableStream bool // 启用流式输出 - - Dingtalk bool - Wecom bool - - CuratedEnabled string - CuratedEndpoint string - CuratedTimeout int - CuratedForceUpdate bool -} - -// NewSDKOptions 创建默认配置 -func NewSDKOptions() *SDKOptions { - return &SDKOptions{ - RateLimit: 150, - Concurrency: 25, - Retries: 1, - Timeout: 50, - MaxHostError: 3, - MaxRespBodySize: 2, - BruteMaxRequests: 5000, - DefaultAccept: true, - FingerprintFilterMode: "strict", - PSPorts: "top", - PSS4Chunk: 1000, - OOBRateLimit: 25, - OOBConcurrency: 25, - OOBFinalizeTimeout: -1, - } -} - -// NewSDKScanner 创建SDK扫描器实例 -func NewSDKScanner(opts *SDKOptions) (*SDKScanner, error) { - if opts == nil { - opts = NewSDKOptions() - } - - // 转换为内部配置 - options := convertSDKOptions(opts) - - // 强制SDK模式设置 - options.Silent = true - options.DisableUpdateCheck = true - options.DisableOutputHtml = true - options.SDKMode = true - options.EnableOOB = opts.EnableOOB - - // 禁用所有文件输出 - options.Json = "" - options.JsonAll = "" - options.Output = "" - - // 使用空配置，避免读取配置文件 - // SDK 模式下，尝试读取默认配置文件以获取 OOB 配置 - cfg, err := config.NewConfig("") - if err != nil { - // 如果读取失败，则使用空配置 - cfg = &config.Config{} - } - options.Config = cfg - - if v := strings.TrimSpace(opts.CuratedEnabled); v != "" { - options.Config.Curated.Enabled = v - } - if v := strings.TrimSpace(opts.CuratedEndpoint); v != "" { - options.Config.Curated.Endpoint = v - } - if opts.CuratedTimeout > 0 { - options.Config.Curated.TimeoutSec = opts.CuratedTimeout - } - options.CuratedForceUpdate = opts.CuratedForceUpdate - if err := applyCuratedMount(options); err != nil { - return nil, err - } - - // 设置OOB配置（只有启用OOB且配置了OOB适配器才设置） - if opts.EnableOOB && opts.OOB != "" { - options.OOB = opts.OOB - options.OOBKey = opts.OOBKey - options.OOBDomain = opts.OOBDomain - options.OOBApiUrl = opts.OOBApiUrl - options.OOBHttpUrl = opts.OOBHttpUrl - - // 如果 SDK 选项中未提供 OOB 配置，尝试从配置文件中读取 - if options.OOBKey == "" && options.OOBDomain == "" && options.OOBApiUrl == "" && options.OOBHttpUrl == "" { - switch options.OOB { - case "ceyeio": - options.OOBKey = cfg.Reverse.Ceye.ApiKey - options.OOBDomain = cfg.Reverse.Ceye.Domain - case "dnslogcn": - options.OOBDomain = cfg.Reverse.Dnslogcn.Domain - case "alphalog": - options.OOBDomain = cfg.Reverse.Alphalog.Domain - options.OOBApiUrl = cfg.Reverse.Alphalog.ApiUrl - case "xray": - options.OOBKey = cfg.Reverse.Xray.XToken - options.OOBDomain = cfg.Reverse.Xray.Domain - options.OOBApiUrl = cfg.Reverse.Xray.ApiUrl - case "revsuit": - options.OOBKey = cfg.Reverse.Revsuit.Token - options.OOBDomain = cfg.Reverse.Revsuit.DnsDomain - options.OOBApiUrl = cfg.Reverse.Revsuit.ApiUrl - options.OOBHttpUrl = cfg.Reverse.Revsuit.HttpUrl - } - } - - if options.OOB == "dnslogcn" && options.OOBDomain == "" { - options.OOBDomain = "dnslog.cn" - } - } - - // 验证配置 - if err := validateSDKConfig(options); err != nil { - return nil, err - } - - // 创建runner - r, err := createSDKRunner(options) - if err != nil { - return nil, fmt.Errorf("创建扫描引擎失败: %w", err) - } - - ctx, cancel := context.WithCancel(context.Background()) - - scanner := &SDKScanner{ - runner: r, - results: make([]*result.Result, 0), - options: options, - sdkOpts: opts, - ctx: ctx, - cancel: cancel, - openPorts: make(map[string]map[int]struct{}), - runStarted: make(chan struct{}), - runDone: make(chan struct{}), - stats: &ScanStats{ - StartTime: time.Now(), - }, - phases: make(map[string]PhaseProgress), - } - atomic.StoreInt32(&scanner.lastVulnPhasePercent, -1) - - options.OnPortScanResult = func(host string, port int) { - scanner.openPortsMu.Lock() - pm, ok := scanner.openPorts[host] - if !ok { - pm = make(map[int]struct{}) - scanner.openPorts[host] = pm - } - pm[port] = struct{}{} - scanner.openPortsMu.Unlock() - - if scanner.PortChan != nil { - ch := scanner.PortChan - func() { - defer func() { _ = recover() }() - select { - case ch <- PortScanResult{Host: host, Port: port}: - case <-scanner.ctx.Done(): - return - default: - } - }() - } - - if scanner.OnPort != nil { - scanner.OnPort(host, port) - } - } - - options.OnHostDiscovered = func(host string) { - host = strings.TrimSpace(host) - if host == "" { - return - } - if scanner.HostChan != nil { - ch := scanner.HostChan - func() { - defer func() { _ = recover() }() - select { - case ch <- HostDiscoveryResult{Host: host}: - case <-scanner.ctx.Done(): - return - default: - } - }() - } - } - - scanner.runner.OnWebProbe = func(meta runner.WebMeta) { - if scanner.WebProbeChan == nil && scanner.OnWebProbe == nil { - return - } - r := WebProbeResult{ - URL: strings.TrimSpace(meta.URL), - Title: strings.TrimSpace(meta.Title), - Server: strings.TrimSpace(meta.Server), - PoweredBy: strings.TrimSpace(meta.PoweredBy), - } - if scanner.WebProbeChan != nil { - ch := scanner.WebProbeChan - func() { - defer func() { _ = recover() }() - select { - case ch <- r: - case <-scanner.ctx.Done(): - return - default: - } - }() - } - if scanner.OnWebProbe != nil { - scanner.OnWebProbe(r) - } - } - - // 如果启用流式输出，创建结果通道 - if opts.EnableStream { - scanner.ResultChan = make(chan *result.Result, 100) - scanner.PhaseProgressChan = make(chan PhaseProgress, 64) - scanner.ScanInfoChan = make(chan ScanInfoUpdate, 16) - } - - if opts.PortScan { - scanner.PortChan = make(chan PortScanResult, 100) - scanner.HostChan = make(chan HostDiscoveryResult, 256) - } - - if opts.EnableWebProbe { - scanner.WebProbeChan = make(chan WebProbeResult, 100) - } - - options.OnPhaseProgress = func(phase string, status string, finished int64, total int64, percent int) { - phase = strings.ToLower(strings.TrimSpace(phase)) - if phase == "" { - return - } - if percent < 0 { - percent = 0 - } - if percent > 100 { - percent = 100 - } - pp := PhaseProgress{ - Phase: phase, - Status: strings.ToLower(strings.TrimSpace(status)), - Finished: finished, - Total: total, - Percent: percent, - } - scanner.phaseMu.Lock() - scanner.phases[phase] = pp - scanner.phaseMu.Unlock() - if scanner.PhaseProgressChan != nil { - ch := scanner.PhaseProgressChan - func() { - defer func() { _ = recover() }() - select { - case ch <- pp: - case <-scanner.ctx.Done(): - return - default: - } - }() - } - } - - options.OnScanInfoUpdate = func(info config.ScanInfoUpdate) { - scanner.stats.TotalTargets = info.TotalTargets - scanner.stats.TotalPocs = info.TotalPocs - scanner.stats.TotalScans = info.TotalScans - if scanner.ScanInfoChan != nil { - ch := scanner.ScanInfoChan - up := ScanInfoUpdate{ - TotalTargets: info.TotalTargets, - Targets: append([]string(nil), info.Targets...), - TotalPocs: info.TotalPocs, - TotalScans: info.TotalScans, - OOBEnabled: info.OOBEnabled, - OOBStatus: info.OOBStatus, - } - func() { - defer func() { _ = recover() }() - select { - case ch <- up: - case <-scanner.ctx.Done(): - return - default: - } - }() - } - } - - // 计算扫描统计 - pocSlice := options.CreatePocList() - fingerprintPocs, pocSlice := options.FingerprintPoCs(pocSlice) - - allTargets := make([]string, 0, options.Targets.Len()) - for _, t := range options.Targets.List() { - s, ok := t.(string) - if !ok { - continue - } - s = strings.TrimSpace(s) - if s == "" { - continue - } - allTargets = append(allTargets, s) - } - idx := targets.BuildTargetIndex(allTargets) - netTargets := idx.NetTargets() - - isNetOnlyPoc := func(p poc.Poc) bool { - hasHTTP := false - hasNet := false - hasGo := false - for _, rm := range p.Rules { - t := strings.ToLower(strings.TrimSpace(rm.Value.Request.Type)) - switch t { - case "", poc.HTTP_Type, poc.HTTPS_Type: - hasHTTP = true - case poc.TCP_Type, poc.UDP_Type, poc.SSL_Type: - hasNet = true - case poc.GO_Type: - hasGo = true - default: - hasHTTP = true - } - } - if hasGo { - return false - } - return hasNet && !hasHTTP - } - - taskCount := 0 - if !options.DisableFingerprint && len(fingerprintPocs) > 0 { - taskCount += len(fingerprintPocs) * len(allTargets) - } - for _, p := range pocSlice { - if !isNetOnlyPoc(p) { - taskCount += len(allTargets) - } else { - taskCount += len(netTargets) - } - } - - pocTotal := len(pocSlice) - if !options.DisableFingerprint && len(fingerprintPocs) > 0 { - pocTotal += len(fingerprintPocs) - } - - scanner.stats.TotalTargets = len(allTargets) - scanner.stats.TotalPocs = pocTotal - scanner.stats.TotalScans = taskCount - - return scanner, nil -} - -// Run 执行扫描（同步版本） -func (s *SDKScanner) Run() error { - // 在扫描开始前输出基本信息 - s.printScanInfo() - err := s.run() - s.closeChans() - return err -} - -// RunAsync 执行扫描（异步版本） -func (s *SDKScanner) RunAsync() error { - go func() { - // 在扫描开始前输出基本信息 - s.printScanInfo() - s.run() - s.closeChans() - }() - return nil -} - -func (s *SDKScanner) closeChans() { - s.closeChansOnce.Do(func() { - if s.ResultChan != nil { - close(s.ResultChan) - } - if s.PortChan != nil { - close(s.PortChan) - } - if s.HostChan != nil { - close(s.HostChan) - } - if s.WebProbeChan != nil { - close(s.WebProbeChan) - } - if s.PhaseProgressChan != nil { - close(s.PhaseProgressChan) - } - if s.ScanInfoChan != nil { - close(s.ScanInfoChan) - } - }) -} - -// run 内部扫描执行 -func (s *SDKScanner) run() error { - s.runStartOnce.Do(func() { - close(s.runStarted) - }) - defer s.runDoneOnce.Do(func() { - close(s.runDone) - }) - - // 设置结果处理器 - s.runner.OnResult = func(r *result.Result) { - if r == nil || !r.SkipCount { - atomic.AddInt32(&s.stats.CompletedScans, 1) - } - if s.options != nil && s.options.OnPhaseProgress != nil { - total := int64(s.stats.TotalScans) - completed := int64(atomic.LoadInt32(&s.stats.CompletedScans)) - percent := 100 - if total > 0 { - percent = int(completed * 100 / total) - if percent > 100 { - percent = 100 - } - if percent < 0 { - percent = 0 - } - } - if int32(percent) != atomic.LoadInt32(&s.lastVulnPhasePercent) { - atomic.StoreInt32(&s.lastVulnPhasePercent, int32(percent)) - status := "running" - if total == 0 || (completed >= total && percent >= 100) { - status = "completed" - } - s.options.OnPhaseProgress("vuln", status, completed, total, percent) - } - } - - if r.IsVul { - s.mu.Lock() - s.results = append(s.results, r) - atomic.AddInt32(&s.stats.FoundVulns, 1) - s.mu.Unlock() - - // 同步回调 - if s.OnResult != nil { - s.OnResult(r) - } - - // 流式输出 - if s.ResultChan != nil { - ch := s.ResultChan - func() { - defer func() { _ = recover() }() - select { - case ch <- r: - case <-s.ctx.Done(): - return - default: - } - }() - } - - // 如果设置了发现漏洞即停止 - if s.options.VulnerabilityScannerBreakpoint { - return - } - } - } - - s.runner.OnFingerprint = func(targetKey string, hits []fingerprint.Hit) { - for _, hit := range hits { - sev := strings.TrimSpace(hit.Severity) - if sev == "" { - sev = "info" - } - name := strings.TrimSpace(hit.Name) - if name == "" { - name = strings.TrimSpace(hit.ID) - } - - rst := s.runner.FingerprintResult(targetKey, hit.ID) - if rst == nil { - rst = &result.Result{ - IsVul: true, - Target: targetKey, - FullTarget: targetKey, - PocInfo: &poc.Poc{ - Id: hit.ID, - Info: poc.Info{ - Name: name, - Severity: sev, - Tags: hit.Tags, - }, - }, - FingerResult: []fingerprint.Hit{hit}, - } - } else { - rst.IsVul = true - if rst.PocInfo == nil { - rst.PocInfo = &poc.Poc{Id: hit.ID} - } else { - rst.PocInfo.Id = hit.ID - } - rst.PocInfo.Info.Name = name - rst.PocInfo.Info.Severity = sev - rst.PocInfo.Info.Tags = hit.Tags - rst.FingerResult = []fingerprint.Hit{hit} - } - if strings.TrimSpace(rst.FullTarget) == "" { - rst.FullTarget = rst.Target - } - if s.runner.OnResult != nil { - s.runner.OnResult(rst) - } - } - } - - // 执行扫描 - s.runner.Execute() - - if s.options != nil && s.options.OnPhaseProgress != nil { - total := int64(s.stats.TotalScans) - completed := int64(atomic.LoadInt32(&s.stats.CompletedScans)) - percent := 100 - if total > 0 { - percent = int(completed * 100 / total) - if percent > 100 { - percent = 100 - } - if percent < 0 { - percent = 0 - } - } - status := "completed" - if s.ctx != nil && s.ctx.Err() != nil { - status = "interrupted" - } else if total > 0 && completed < total { - status = "interrupted" - } - if status == "completed" { - percent = 100 - if total > 0 { - completed = total - } - } - s.options.OnPhaseProgress("vuln", status, completed, total, percent) - } - - s.stats.EndTime = time.Now() - return nil -} - -// GetResults 获取所有扫描结果 -func (s *SDKScanner) GetResults() []*result.Result { - s.mu.Lock() - defer s.mu.Unlock() - - results := make([]*result.Result, len(s.results)) - copy(results, s.results) - return results -} - -func (s *SDKScanner) GetOpenPorts() map[string][]int { - s.openPortsMu.Lock() - defer s.openPortsMu.Unlock() - - out := make(map[string][]int, len(s.openPorts)) - for host, ports := range s.openPorts { - if len(ports) == 0 { - continue - } - dst := make([]int, 0, len(ports)) - for p := range ports { - dst = append(dst, p) - } - out[host] = dst - } - return out -} - -// GetStats 获取扫描统计信息 -func (s *SDKScanner) GetStats() ScanStats { - stats := *s.stats - stats.CompletedScans = atomic.LoadInt32(&s.stats.CompletedScans) - stats.FoundVulns = atomic.LoadInt32(&s.stats.FoundVulns) - return stats -} - -// GetProgress 获取扫描进度（0-100） -func (s *SDKScanner) GetProgress() float64 { - clip := func(v float64) float64 { - if v < 0 { - return 0 - } - if v > 100 { - return 100 - } - return v - } - - getPhasePercent := func(name string) int { - name = strings.ToLower(strings.TrimSpace(name)) - if name == "" { - return 0 - } - s.phaseMu.Lock() - pp, ok := s.phases[name] - s.phaseMu.Unlock() - if !ok { - return 0 - } - if pp.Percent < 0 { - return 0 - } - if pp.Percent > 100 { - return 100 - } - return pp.Percent - } - - vulnPercent := 100.0 - if s.stats.TotalScans > 0 { - completed := atomic.LoadInt32(&s.stats.CompletedScans) - vulnPercent = float64(completed) / float64(s.stats.TotalScans) * 100 - } - - hasPort := s.sdkOpts != nil && s.sdkOpts.PortScan - hasWebProbe := s.sdkOpts != nil && s.sdkOpts.EnableWebProbe - - portStagePercent := 100.0 - if hasPort { - hostDisc := getPhasePercent("host_discovery") - portscan := getPhasePercent("portscan") - if hostDisc == 0 && portscan == 0 { - portStagePercent = 0 - } else { - portStagePercent = (float64(hostDisc) + float64(portscan)) / 2 - } - } - - webprobePercent := 100.0 - if hasWebProbe { - wp := getPhasePercent("webprobe") - webprobePercent = float64(wp) - } - - wPort := 0.0 - if hasPort { - wPort = 0.2 - } - wWeb := 0.0 - if hasWebProbe { - wWeb = 0.2 - } - wVuln := 1.0 - wPort - wWeb - if wVuln < 0 { - wVuln = 0 - } - - return clip(wPort*portStagePercent + wWeb*webprobePercent + wVuln*clip(vulnPercent)) -} - -// Stop 停止扫描 -func (s *SDKScanner) Stop() { - s.cancel() - s.options.VulnerabilityScannerBreakpoint = true - if s.runner != nil { - s.runner.Stop() - } -} - -// Close 关闭扫描器，释放资源 -func (s *SDKScanner) Close() { - s.Stop() - if s.runStarted != nil { - select { - case <-s.runStarted: - if s.runDone != nil { - <-s.runDone - } - default: - } - } - s.closeChans() - s.results = nil -} - -// HasVulnerabilities 检查是否发现漏洞 -func (s *SDKScanner) HasVulnerabilities() bool { - return atomic.LoadInt32(&s.stats.FoundVulns) > 0 -} - -// GetVulnerabilityCount 获取漏洞数量 -func (s *SDKScanner) GetVulnerabilityCount() int { - return int(atomic.LoadInt32(&s.stats.FoundVulns)) -} - -func (s *SDKScanner) Pause() { - if s.runner != nil { - s.runner.Pause() - } -} - -func (s *SDKScanner) Resume() { - if s.runner != nil { - s.runner.Resume() - } -} - -func (s *SDKScanner) IsPaused() bool { - if s.runner == nil { - return false - } - return s.runner.IsPaused() -} - -func (s *SDKScanner) IsStopping() bool { - return s.options.VulnerabilityScannerBreakpoint -} - -// SetProxy 动态设置代理 -func (s *SDKScanner) SetProxy(proxy string) { - s.options.Proxy = proxy - retryhttpclient.Init(&retryhttpclient.Options{ - Proxy: proxy, - Timeout: s.options.Timeout, - Retries: s.options.Retries, - MaxRespBodySize: s.options.MaxRespBodySize, - ReqLimitPerTarget: s.options.ReqLimitPerTarget, - DefaultAccept: s.options.DefaultAccept, - }) -} - -// SetRateLimit 动态设置速率限制 -func (s *SDKScanner) SetRateLimit(rateLimit int) { - s.options.RateLimit = rateLimit -} - -// SetConcurrency 动态设置并发数 -func (s *SDKScanner) SetConcurrency(concurrency int) { - s.options.Concurrency = concurrency -} - -// ========== OOB检测相关函数 ========== - -// IsOOBEnabled 检查是否启用了OOB检测 -func (s *SDKScanner) IsOOBEnabled() bool { - // 首先检查SDK配置中是否明确启用了OOB - if s.sdkOpts != nil && s.sdkOpts.EnableOOB { - return true - } - // 检查是否配置了OOB相关参数 - return s.options.OOB != "" && (s.options.OOBKey != "" || s.options.OOBDomain != "") -} - -// GetOOBStatus 获取OOB状态信息 -func (s *SDKScanner) GetOOBStatus() (bool, string) { - // 首先检查是否在SDK选项中明确启用了OOB - sdkOpts := s.getSDKOptions() - if sdkOpts != nil && !sdkOpts.EnableOOB { - return false, "OOB未配置或未启用" - } - - if !s.IsOOBEnabled() { - return false, "OOB未配置或未启用" - } - - if s.runner == nil { - return false, "扫描器未初始化" - } - - // 只有在启用OOB时才进行连接检测 - if sdkOpts != nil && sdkOpts.EnableOOB { - return s.checkOOBConnection() - } - - return false, "OOB未启用" -} - -// getSDKOptions 获取SDK配置选项 -func (s *SDKScanner) getSDKOptions() *SDKOptions { - return s.sdkOpts -} - -// checkOOBConnection 检查OOB连接状态 - SDK专用检测函数 -func (s *SDKScanner) checkOOBConnection() (bool, string) { - // 从配置中获取当前OOB服务名称 - serviceName := strings.ToLower(s.options.OOB) - if serviceName == "" { - return false, "OOB未配置" - } - - // 检查配置完整性 - if s.options.OOBKey == "" && s.options.OOBDomain == "" { - return false, fmt.Sprintf("%s (配置不完整)", serviceName) - } - - // 尝试创建OOB适配器进行连接测试，这里使用和runner相同的逻辑，但在SDK中独立执行 - if oobAdapter, err := s.createOOBAdapter(); err == nil { - if oobAdapter.IsVaild() { - return true, fmt.Sprintf("%s (连接正常)", serviceName) - } else { - return false, fmt.Sprintf("%s (连接失败)", serviceName) - } - } else { - return false, fmt.Sprintf("%s (初始化失败: %v)", serviceName, err) - } -} - -// OOBAdapter 简化的OOB适配器接口 -type OOBAdapter interface { - IsVaild() bool -} - -// simpleOOBAdapter 简单的OOB适配器实现 -type simpleOOBAdapter struct { - valid bool -} - -func (s *simpleOOBAdapter) IsVaild() bool { - return s.valid -} - -// createOOBAdapter 创建OOB适配器 - SDK内部使用 -func (s *SDKScanner) createOOBAdapter() (OOBAdapter, error) { - // 检查配置完整性 - if s.options.OOBKey == "" && s.options.OOBDomain == "" { - return nil, fmt.Errorf("配置不完整") - } - - // 进行OOB反链可用性检测 - if oobAdapter, err := oobadapter.NewOOBAdapter(s.options.OOB, &oobadapter.ConnectorParams{ - Key: s.options.OOBKey, - Domain: s.options.OOBDomain, - HTTPUrl: s.options.OOBHttpUrl, - ApiUrl: s.options.OOBApiUrl, - }); err == nil { - if oobAdapter.IsVaild() { - return &simpleOOBAdapter{valid: true}, nil - } else { - return &simpleOOBAdapter{valid: false}, nil - } - } else { - return &simpleOOBAdapter{valid: false}, err - } - -} - -// printScanInfo 输出扫描基本信息 -func (s *SDKScanner) printScanInfo() { - fmt.Printf("\n========== 扫描信息 ==========\n") - fmt.Printf("目标数量: %d\n", s.stats.TotalTargets) - fmt.Printf("POC数量: %d\n", s.stats.TotalPocs) - fmt.Printf("总扫描任务: %d\n", s.stats.TotalScans) - - // 输出目标列表 - if s.stats.TotalTargets <= 5 { - fmt.Printf("扫描目标: ") - targets := s.options.Targets.List() - for i, target := range targets { - if i > 0 { - fmt.Printf(", ") - } - fmt.Printf("%v", target) - } - fmt.Printf("\n") - } else { - fmt.Printf("目标过多，仅显示前3个: ") - targets := s.options.Targets.List() - for i := 0; i < 3 && i < len(targets); i++ { - if i > 0 { - fmt.Printf(", ") - } - fmt.Printf("%v", targets[i]) - } - fmt.Printf("...\n") - } - - // OOB状态 - 只在明确启用OOB时才显示 - if s.sdkOpts != nil && s.sdkOpts.EnableOOB { - if oobEnabled, oobStatus := s.GetOOBStatus(); oobEnabled { - fmt.Printf("OOB状态: ✓ %s\n", oobStatus) - } else { - fmt.Printf("OOB状态: ✗ %s\n", oobStatus) - } - } else { - fmt.Printf("OOB状态: ✗ OOB未配置或未启用\n") - } - - fmt.Printf("=============================\n") -} - -func applyCuratedMount(options *config.Options) error { - if options == nil || options.Config == nil { - return nil - } - cur := options.Config.Curated - enabled := strings.ToLower(strings.TrimSpace(cur.Enabled)) - endpoint := strings.TrimSpace(cur.Endpoint) - if enabled == "off" || enabled == "false" || enabled == "0" || endpoint == "" { - _ = os.Setenv("AFROG_CURATED_DISABLED", "1") - _ = os.Unsetenv("AFROG_POCS_CURATED_DIR") - return nil - } - - _ = os.Unsetenv("AFROG_CURATED_DISABLED") - svc := service.New(service.Config{ - Endpoint: endpoint, - Channel: strings.TrimSpace(cur.Channel), - CuratedPocDir: "", - LicenseKey: strings.TrimSpace(cur.LicenseKey), - NoUpdate: cur.AutoUpdate != nil && !*cur.AutoUpdate && !options.CuratedForceUpdate, - ForceUpdate: options.CuratedForceUpdate, - ClientVersion: config.Version, - }) - - ctx, cancel := context.WithTimeout(context.Background(), time.Duration(cur.TimeoutSec)*time.Second) - if cur.TimeoutSec <= 0 { - ctx, cancel = context.WithCancel(context.Background()) - } - defer cancel() - - dir, err := svc.Mount(ctx) - if err != nil { - fmt.Fprintf(os.Stderr, "curated mount failed: %s\n", strings.TrimSpace(err.Error())) - return nil - } - if strings.TrimSpace(dir) != "" { - _ = os.Setenv("AFROG_POCS_CURATED_DIR", dir) - } - return nil -} - -// convertSDKOptions 转换SDK配置到内部配置 -func convertSDKOptions(opts *SDKOptions) *config.Options { - options := &config.Options{ - TargetsFile: opts.TargetsFile, - PocFile: opts.PocFile, - AppendPoc: opts.AppendPoc, - Search: opts.Search, - Severity: opts.Severity, - ExcludePocs: opts.ExcludePocs, - ExcludePocsFile: opts.ExcludePocsFile, - RateLimit: opts.RateLimit, - ReqLimitPerTarget: opts.ReqLimitPerTarget, - AutoReqLimit: opts.AutoReqLimit, - Polite: opts.Polite, - Balanced: opts.Balanced, - Aggressive: opts.Aggressive, - Concurrency: opts.Concurrency, - Retries: opts.Retries, - Timeout: opts.Timeout, - MaxHostError: opts.MaxHostError, - Proxy: opts.Proxy, - MaxRespBodySize: opts.MaxRespBodySize, - BruteMaxRequests: opts.BruteMaxRequests, - DefaultAccept: opts.DefaultAccept, - OOBRateLimit: opts.OOBRateLimit, - OOBConcurrency: opts.OOBConcurrency, - OOBFinalizeTimeout: opts.OOBFinalizeTimeout, - Smart: opts.Smart, - DisableFingerprint: opts.DisableFingerprint, - EnableWebProbe: opts.EnableWebProbe, - FingerprintFilterMode: opts.FingerprintFilterMode, - VulnerabilityScannerBreakpoint: opts.VulnerabilityScannerBreakpoint, - PortScan: opts.PortScan, - PSPorts: opts.PSPorts, - PSRateLimit: opts.PSRateLimit, - PSTimeout: opts.PSTimeout, - PSRetries: opts.PSRetries, - PSSkipDiscovery: opts.PSSkipDiscovery, - PSS4Chunk: opts.PSS4Chunk, - Dingtalk: opts.Dingtalk, - Wecom: opts.Wecom, - CuratedEnabled: opts.CuratedEnabled, - CuratedEndpoint: opts.CuratedEndpoint, - CuratedTimeout: opts.CuratedTimeout, - CuratedForceUpdate: opts.CuratedForceUpdate, - } - - if options.MaxRespBodySize <= 0 { - options.MaxRespBodySize = 2 - } - if strings.TrimSpace(options.FingerprintFilterMode) == "" { - options.FingerprintFilterMode = "strict" - } - if options.OOBRateLimit == 0 { - options.OOBRateLimit = 25 - } - if options.OOBConcurrency == 0 { - options.OOBConcurrency = 25 - } - if options.ReqLimitPerTarget == 0 { - if options.Polite { - options.ReqLimitPerTarget = 5 - } else if options.Balanced { - options.ReqLimitPerTarget = 15 - } else if options.Aggressive { - options.ReqLimitPerTarget = 50 - } else if options.AutoReqLimit { - baseRate := options.RateLimit - if baseRate <= 0 { - baseRate = 150 - } - r := baseRate / 10 - if r < 5 { - r = 5 - } - if r > 15 { - r = 15 - } - con := options.Concurrency - if con <= 0 { - con = 1 - } - if con >= 100 && r > 8 { - r = 8 - } else if con >= 50 && r > 12 { - r = 12 - } - options.ReqLimitPerTarget = r - } - } - - if len(opts.Headers) > 0 { - for _, h := range opts.Headers { - options.Header = append(options.Header, h) - } - } - - // 转换目标列表 - if len(opts.Targets) > 0 { - for _, target := range opts.Targets { - options.Target = append(options.Target, target) - } - } - - return options -} - -// validateSDKConfig SDK配置验证 -func validateSDKConfig(options *config.Options) error { - limitModeCount := 0 - if options.ReqLimitPerTarget > 0 { - limitModeCount++ - } - if options.AutoReqLimit { - limitModeCount++ - } - if options.Polite { - limitModeCount++ - } - if options.Balanced { - limitModeCount++ - } - if options.Aggressive { - limitModeCount++ - } - if limitModeCount > 1 { - return errors.New("only one of ReqLimitPerTarget/AutoReqLimit/Polite/Balanced/Aggressive can be used") - } - if options.ReqLimitPerTarget < 0 { - return errors.New("ReqLimitPerTarget must be >= 0") - } - - options.FingerprintFilterMode = strings.ToLower(strings.TrimSpace(options.FingerprintFilterMode)) - if options.FingerprintFilterMode == "" { - options.FingerprintFilterMode = "strict" - } - if options.FingerprintFilterMode != "strict" && options.FingerprintFilterMode != "opportunistic" { - options.FingerprintFilterMode = "strict" - } - - // 验证目标 - if len(options.Target) == 0 && len(options.TargetsFile) == 0 { - return errors.New("未指定扫描目标") - } - - // 验证POC文件 - if options.PocFile == "" && len(options.AppendPoc) == 0 { - // 如果PocFile为空且AppendPoc也为空，且不使用默认配置（这里允许为空，由Runner处理默认值） - // return errors.New("必须指定POC文件或目录") - } - - // 验证POC文件是否存在 - if options.PocFile != "" { - if _, err := os.Stat(options.PocFile); err != nil { - return fmt.Errorf("POC文件或目录不存在: %s", options.PocFile) - } - } - - if options.Config != nil { - tokensEmpty := func(tokens []string) bool { - for _, t := range tokens { - if strings.TrimSpace(t) != "" { - return false - } - } - return true - } - - if options.Dingtalk && tokensEmpty(options.Config.Webhook.Dingtalk.Tokens) { - return errors.New("Dingtalk webhook token is required") - } - if options.Wecom && tokensEmpty(options.Config.Webhook.Wecom.Tokens) { - return errors.New("Wecom webhook token is required") - } - } - - return nil -} - -// createSDKRunner 创建SDK专用的Runner -func createSDKRunner(options *config.Options) (*runner.Runner, error) { - // 初始化HTTP客户端 - retryhttpclient.Init(&retryhttpclient.Options{ - Proxy: options.Proxy, - Timeout: options.Timeout, - Retries: options.Retries, - MaxRespBodySize: options.MaxRespBodySize, - ReqLimitPerTarget: options.ReqLimitPerTarget, - DefaultAccept: options.DefaultAccept, - }) - - // 处理目标 - seen := make(map[string]struct{}) - - // 添加命令行目标 - if len(options.Target) > 0 { - for _, rawTarget := range options.Target { - trimmedTarget := strings.TrimSpace(rawTarget) - if _, ok := seen[trimmedTarget]; !ok { - seen[trimmedTarget] = struct{}{} - options.Targets.Append(trimmedTarget) - } - } - } - - // 从文件读取目标 - if len(options.TargetsFile) > 0 { - allTargets, err := utils.ReadFileLineByLine(options.TargetsFile) - if err != nil { - return nil, err - } - for _, rawTarget := range allTargets { - trimmedTarget := strings.TrimSpace(rawTarget) - if len(trimmedTarget) > 0 { - if _, ok := seen[trimmedTarget]; !ok { - seen[trimmedTarget] = struct{}{} - options.Targets.Append(trimmedTarget) - } - } - } - } - - // 验证目标 - if options.Targets.Len() == 0 { - return nil, errors.New("未找到有效目标") - } - - // 设置POC目录 - if options.PocFile != "" { - options.PocsDirectory.Set(options.PocFile) - } - for _, p := range options.AppendPoc { - options.PocsDirectory.Set(p) - } - - // 清空Target切片 - options.Target = nil - - // 创建Runner - r, err := runner.NewRunner(options) - if err != nil { - return nil, err - } - - // SDK模式特殊处理 - c := catalog.New(options.PocsDirectory.String()) - allPocsYamlSlice := c.GetPocsPath(options.PocsDirectory) - if len(allPocsYamlSlice) == 0 && len(pocs.EmbedFileList) == 0 { - return nil, errors.New("未找到POC文件") - } - - return r, nil -} diff --git a/cmd/sdk/main.go b/cmd/sdk/main.go index 1e89d4841..bc7a05293 100644 --- a/cmd/sdk/main.go +++ b/cmd/sdk/main.go @@ -1,36 +1,69 @@ +// Command sdk is a minimal CI security gate built on the afrog SDK. +// +// It scans a target and exits non-zero when a vulnerability is found. package main import ( + "context" + "flag" "fmt" "os" + "os/signal" - "github.com/zan8in/afrog/v3" + "github.com/zan8in/afrog/v3/pkg/sdk" ) func main() { - options := afrog.NewSDKOptions() - options.Targets = []string{"honey.scanme.sh"} - options.PocFile = "./pocs/afrog-pocs" - // options.Severity = "info" - // options.Search = "React" - // options.Proxy = "http://127.0.0.1:51024" - - scanner, err := afrog.NewSDKScanner(options) + target := flag.String("t", "https://scanme.sh", "target to scan") + pocs := flag.String("pocs", "", "PoC file, directory or glob pattern (empty = built-in PoCs)") + severity := flag.String("severity", "", "severity filter, e.g. \"high,critical\"") + proxy := flag.String("proxy", "", "HTTP/SOCKS5 proxy") + flag.Parse() + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + found, err := run(ctx, *target, *pocs, *severity, *proxy) if err != nil { + fmt.Fprintf(os.Stderr, "afrog: %v\n", err) + os.Exit(1) + } + if found { os.Exit(1) } +} + +func run(ctx context.Context, target, pocs, severity, proxy string) (bool, error) { + options := []sdk.Option{sdk.WithTargets(target)} + if pocs != "" { + options = append(options, sdk.WithPocPaths(pocs), sdk.WithPocPathsOnly()) + } + if severity != "" { + options = append(options, sdk.WithSeverity(severity)) + } + if proxy != "" { + options = append(options, sdk.WithProxy(proxy)) + } + + scanner, err := sdk.New(ctx, options...) + if err != nil { + return false, err + } defer scanner.Close() - scanner.Run() + if err := scanner.Execute(ctx); err != nil { + return false, err + } - if scanner.HasVulnerabilities() { - fmt.Println("❌ 发现安全漏洞，阻止部署") - results := scanner.GetResults() - for _, r := range results { - fmt.Printf("- %s: %s\n", r.Target, r.PocInfo.Info.Name) - } - os.Exit(1) + results := scanner.Results() + if len(results) == 0 { + fmt.Println("security check passed / 安全检查通过") + return false, nil } - fmt.Println("✅ 安全检查通过") + fmt.Printf("found %d vulnerabilities / 发现 %d 个漏洞\n", len(results), len(results)) + for _, v := range results { + fmt.Printf("- [%s] %s: %s\n", v.Severity, v.FullTarget, v.PocName) + } + return true, nil } diff --git a/docs/SDK_Usage_Guide_English.md b/docs/SDK_Usage_Guide_English.md index 7e3fd5527..5eeda26fa 100644 --- a/docs/SDK_Usage_Guide_English.md +++ b/docs/SDK_Usage_Guide_English.md @@ -2,14 +2,18 @@ ## Overview -The Afrog SDK provides a clean, efficient Go programming interface specifically designed for integrating vulnerability scanning capabilities. The SDK features the following core characteristics: +The Afrog SDK is the Go API for embedding vulnerability scanning into your own programs. The import path is `github.com/zan8in/afrog/v3/pkg/sdk`. -### 🚀 Core Features -- ✅ **Structured Returns** - Direct Go struct returns for easy program processing -- ✅ **Real-time Result Streaming** - Supports both synchronous callbacks and asynchronous streaming -- ✅ **OOB Detection Support** - Complete Out-of-Band detection configuration and management -- ✅ **Detailed Statistics** - Provides scan progress, performance, and result statistics -- ✅ **Concurrency Safe** - All APIs are thread-safe +### Key features + +- **Structured results** — plain Go structs that can be passed straight to `json.Marshal` +- **Complete data output** — the raw request and response of every scan step is available +- **Flexible PoC input** — single files, directories (searched recursively) and glob patterns +- **Synchronous and asynchronous** — `Execute` blocks; `Start` plus `Wait`/`Done` runs in the background +- **Handlers and streams** — register multiple callbacks, or subscribe to event channels on demand +- **Silent by default** — nothing is written to stdout or stderr +- **Typed errors** — failures are matched with `errors.Is` +- **Deterministic cleanup** — `Close` releases every background goroutine ## Installation @@ -17,616 +21,755 @@ The Afrog SDK provides a clean, efficient Go programming interface specifically go get -u github.com/zan8in/afrog/v3 ``` -## Quick Start - -### Basic Scan Example - -The simplest usage pattern, suitable for quick integration: +## Quick start ```go package main import ( - "fmt" - "log" - "path/filepath" - "github.com/zan8in/afrog/v3" + "context" + "fmt" + "log" + + "github.com/zan8in/afrog/v3/pkg/sdk" ) func main() { - // Create scan options - options := afrog.NewSDKOptions() - - // Set scan targets - options.Targets = []string{"https://www.example.com"} - - // Set POC path (required) - pocPath, _ := filepath.Abs("./pocs/afrog-pocs") - options.PocFile = pocPath - - // Create scanner - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - log.Fatal(err) - } - defer scanner.Close() - - // Execute scan - scanner.Run() - - // Get results - results := scanner.GetResults() - fmt.Printf("Found %d vulnerabilities\n", len(results)) + ctx := context.Background() + + scanner, err := sdk.New(ctx, + sdk.WithTargets("https://example.com"), + sdk.WithPocPaths("./pocs/afrog-pocs"), + ) + if err != nil { + log.Fatal(err) + } + defer scanner.Close() + + if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) + } + + for _, r := range scanner.Results() { + fmt.Printf("[%s] %s - %s\n", r.Severity, r.FullTarget, r.PocName) + } } ``` -## SDK Configuration Options +## PoC input -### SDKOptions Structure +`WithPocPaths` accepts three forms, which can be mixed and repeated: ```go -type SDKOptions struct { - // ========== Target Configuration ========== - Targets []string // List of scan targets - TargetsFile string // Path to targets file - - // ========== POC Configuration ========== - PocFile string // POC file or directory path (required) - Search string // POC search keywords - Severity string // Severity level filter - - // ========== Performance Configuration ========== - RateLimit int // Request rate limit (default: 150) - Concurrency int // Concurrency level (default: 25) - Retries int // Retry attempts (default: 1) - Timeout int // Timeout in seconds (default: 10) - MaxHostError int // Max errors per host (default: 3) - - // ========== PortScan Pre-scan Configuration ========== - PortScan bool // Enable port pre-scan (same as CLI -ps) - PSPorts string // Ports definition: top/full/all/80,443/1-1024 etc. (same as -p) - PSRateLimit int // Pre-scan rate limit (same as -prate) - PSTimeout int // Pre-scan timeout in milliseconds (same as -ptimeout) - PSRetries int // Pre-scan retries (same as -ptries) - PSSkipDiscovery bool // Skip host discovery (same as -Pn) - PSS4Chunk int // Chunk size when ports=full (same as --ps-s4-chunk) - - // ========== Network Configuration ========== - Proxy string // HTTP/SOCKS5 proxy - - // ========== OOB Configuration ========== - EnableOOB bool // Enable OOB detection - OOB string // OOB adapter type - OOBKey string // OOB API key - OOBDomain string // OOB domain - OOBApiUrl string // OOB API URL - OOBHttpUrl string // OOB HTTP URL - - // ========== Output Configuration ========== - EnableStream bool // Enable streaming output -} +sdk.WithPocPaths( + "/path/to/single.yaml", // a single file + "/path/to/pocs", // a directory, searched recursively + "/path/to/pocs/*.yaml", // a glob pattern +) ``` -### Configuration Options Explained - -#### Target Configuration -- `Targets`: Directly specify list of scan targets -- `TargetsFile`: Read targets from file (one per line) - -#### POC Configuration -- `PocFile`: **Required** POC file or directory path -- `Search`: Filter POCs by keywords, e.g., "tomcat,phpinfo" -- `Severity`: Filter by severity levels, e.g., "high,critical" +### Append or exclusive -#### Performance Tuning -- `Concurrency`: Number of concurrent scan threads, adjust based on target count -- `RateLimit`: Requests per second limit to avoid triggering defenses -- `Timeout`: Individual request timeout -- `Retries`: Number of retry attempts for failed requests +| Configuration | Behaviour | +|---------------|-----------| +| `WithPocPaths(...)` | **Append**: merged with the built-in, curated, my and local PoCs; explicit paths win on name conflicts | +| `WithPocPaths(...)` + `WithPocPathsOnly()` | **Exclusive**: only the listed PoCs are used | -## Core Feature Examples +### Inspecting what was loaded -### 1. Real-time Result Callbacks - -Process vulnerabilities immediately upon discovery: +You can verify the PoC selection before any network traffic: ```go -scanner.OnResult = func(r *result.Result) { - fmt.Printf("Vulnerability found: %s - %s [%s]\n", - r.Target, - r.PocInfo.Info.Name, - r.PocInfo.Info.Severity) - - // Immediate processing logic - if r.PocInfo.Info.Severity == "critical" { - sendAlert(r) - } +fmt.Printf("loaded %d pocs\n", scanner.PocCount()) + +for _, p := range scanner.Pocs() { + fmt.Println(p.Id, p.Info.Name) } -scanner.Run() +// Which PoCs were skipped, and why +for _, d := range scanner.PocDiagnostics() { + fmt.Printf("skipped %s: %s\n", d.Path, d.Reason) +} ``` -### 2. Progress Monitoring +| `PocLoadError.Reason` | Meaning | +|-----------------------|---------| +| `config.PocLoadNotFound` | Path does not exist, or the glob matched nothing | +| `config.PocLoadReadFailed` | The file could not be read | +| `config.PocLoadParseFailed` | YAML parsing failed | +| `config.PocLoadLegacyOOB` | The PoC uses the deprecated v2 OOB syntax | + +## Complete data output -Monitor scan progress in real-time: +`Results()` returns `sdk.Result` values whose `Exchanges` carry the full request and response of every step: ```go -go func() { - ticker := time.NewTicker(1 * time.Second) - defer ticker.Stop() - - for range ticker.C { - progress := scanner.GetProgress() - stats := scanner.GetStats() - fmt.Printf("Progress: %.2f%% (%d/%d) Vulnerabilities: %d\n", - progress, - stats.CompletedScans, - stats.TotalScans, - stats.FoundVulns) - } -}() +for _, r := range scanner.Results() { + fmt.Printf("%s [%s] %s\n", r.PocID, r.Severity, r.FullTarget) -scanner.Run() -``` + for _, ex := range r.Exchanges { + fmt.Printf("%s %s -> %d (%d ms)\n", ex.Method, ex.URL, ex.StatusCode, ex.LatencyMs) -### 3. Asynchronous Scanning with Streaming + fmt.Println("--- raw request ---") + fmt.Println(ex.Request) -Non-blocking scanning with real-time results: + fmt.Println("--- raw response ---") + fmt.Println(ex.Response) -```go -options.EnableStream = true -scanner, _ := afrog.NewSDKScanner(options) - -// Start async scan -scanner.RunAsync() - -// Read real-time results from channel -for result := range scanner.ResultChan { - fmt.Printf("Live discovery: %s - %s\n", - result.Target, - result.PocInfo.Info.Name) - - // Process each result in real-time - processResult(result) + if ex.BodyTruncated { + fmt.Println("warning: response body was truncated at MaxRespBodySize") + } + } } ``` -### 4. OOB (Out-of-Band) Detection Configuration +### Result -#### CEYE.io Configuration (Recommended) ```go -options.EnableOOB = true -options.OOB = "ceyeio" -options.OOBKey = "your-ceye-api-token" -options.OOBDomain = "your-subdomain.ceye.io" +type Result struct { + PocID string `json:"poc_id"` + PocName string `json:"poc_name,omitempty"` + Severity string `json:"severity,omitempty"` + Author string `json:"author,omitempty"` + Description string `json:"description,omitempty"` + Reference []string `json:"reference,omitempty"` + Tags []string `json:"tags,omitempty"` + + CveID string `json:"cve_id,omitempty"` + CweID string `json:"cwe_id,omitempty"` + CvssScore float64 `json:"cvss_score,omitempty"` + CvssMetrics string `json:"cvss_metrics,omitempty"` + + Target string `json:"target"` + FullTarget string `json:"full_target,omitempty"` + + Extractors map[string]string `json:"extractors,omitempty"` + Fingerprints []Fingerprint `json:"fingerprints,omitempty"` + Exchanges []Exchange `json:"exchanges,omitempty"` + + FoundAt time.Time `json:"found_at"` +} ``` -#### DNSLog.cn Configuration (Free) +### Exchange + ```go -options.EnableOOB = true -options.OOB = "dnslogcn" -options.OOBDomain = "your.dnslog.cn" +type Exchange struct { + Request string `json:"request,omitempty"` // raw request + Response string `json:"response,omitempty"` // raw response + + Method string `json:"method,omitempty"` + URL string `json:"url,omitempty"` + RequestHeaders map[string]string `json:"request_headers,omitempty"` + RequestBody string `json:"request_body,omitempty"` + StatusCode int `json:"status_code,omitempty"` + ResponseHeaders map[string]string `json:"response_headers,omitempty"` + ResponseBody string `json:"response_body,omitempty"` + ContentType string `json:"content_type,omitempty"` + LatencyMs int64 `json:"latency_ms,omitempty"` + + Matched bool `json:"matched"` + BodyTruncated bool `json:"body_truncated,omitempty"` + BruteTruncated bool `json:"brute_truncated,omitempty"` + BruteRequests int `json:"brute_requests,omitempty"` +} ``` -#### Other OOB Services -```go -// Alphalog -options.OOB = "alphalog" -options.OOBDomain = "your.alphalog.cn" -options.OOBApiUrl = "https://api.alphalog.cn" +Raw messages are strings rather than `[]byte`, so they serialise as readable text instead of base64: -// XRay -options.OOB = "xray" -options.OOBDomain = "your.xray.domain" -options.OOBApiUrl = "http://xray-api:8777" -options.OOBKey = "your-xray-token" +```go +data, err := json.MarshalIndent(scanner.Results(), "", " ") ``` -#### OOB Status Check +### Controlling memory usage + ```go -if oobEnabled, oobStatus := scanner.GetOOBStatus(); oobEnabled { - fmt.Printf("✓ OOB Status: %s\n", oobStatus) -} else { - fmt.Printf("✗ OOB Status: %s\n", oobStatus) -} +sdk.WithRequestResponse(false), // do not retain Exchanges +sdk.WithMaxStoredResults(1000), // accumulate at most 1000 results ``` -### 5. Port Pre-scan (PortScan) +`MaxStoredResults` only bounds internal accumulation. Handlers and streams still receive **every** result. -The SDK can perform a port pre-scan before running PoCs. Discovered open ports will be appended to the internal Targets (as `host:port`), and subsequent PoC scans will run against the updated target set. +### Response body truncation -In SDK mode, open ports are not printed by default. Consume them via callback or by retrieving the collected results. +The response body limit is `MaxRespBodySize` (2 MB by default). Anything beyond it is discarded and `Exchange.BodyTruncated` is set, so a truncated response is distinguishable from a complete one. ```go -options := afrog.NewSDKOptions() -options.Targets = []string{"1.2.3.4"} -options.PocFile = pocPath +sdk.WithMaxRespBodySize(10) // raise to 10 MB +``` -options.PortScan = true -options.PSPorts = "top" // or "full"/"all"/"80,443"/"1-1024" -options.PSSkipDiscovery = true -options.PSTimeout = 500 +## Synchronous and asynchronous execution -scanner, _ := afrog.NewSDKScanner(options) +### Synchronous -scanner.OnPort = func(host string, port int) { - fmt.Printf("open: %s:%d\n", host, port) +```go +if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) } +results := scanner.Results() +``` -scanner.Run() +### Asynchronous -open := scanner.GetOpenPorts() -_ = open +```go +if err := scanner.Start(ctx); err != nil { + log.Fatal(err) +} + +go func() { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + fmt.Printf("progress: %.1f%%\n", scanner.Progress()) + case <-scanner.Done(): + return + } + } +}() + +if err := scanner.Wait(ctx); err != nil { + log.Printf("scan error: %v", err) +} ``` -You can also consume port pre-scan results asynchronously via `PortChan`. It is initialized when `PortScan` is enabled and will be closed automatically after the scan finishes. +### Lifecycle methods + +| Method | Description | +|--------|-------------| +| `Execute(ctx)` | Runs synchronously until the scan finishes | +| `Start(ctx)` | Starts asynchronously and returns immediately | +| `Wait(ctx)` | Blocks until the scan finishes and returns its error | +| `Done()` | Channel closed when the scan finishes | +| `Err()` | The scan error, or nil while still running | +| `Stop()` | Requests a stop and returns immediately | +| `Close()` | Stops the scan, waits for goroutines to exit, releases resources | +| `Pause()` / `Resume()` / `IsPaused()` | Pause control | +| `IsStopping()` / `IsRunning()` | State queries | + +A scanner is single-use: ```go -options := afrog.NewSDKOptions() -options.Targets = []string{"1.2.3.4"} -options.PocFile = pocPath -options.PortScan = true +scanner.Execute(ctx) // first call: fine +scanner.Execute(ctx) // second call: ErrAlreadyFinished +``` -scanner, _ := afrog.NewSDKScanner(options) +Create a new scanner to scan again. `Close` is idempotent and safe to defer immediately after `New`. -_ = scanner.RunAsync() +## Handlers and streams -for r := range scanner.PortChan { - fmt.Printf("open: %s:%d\n", r.Host, r.Port) -} +### Handlers + +```go +scanner, _ := sdk.New(ctx, + sdk.WithResultHandler(saveToDatabase), + sdk.WithResultHandler(sendAlert), // may be registered several times + sdk.WithFailureHandler(func(f sdk.Failure) { + log.Printf("poc %s failed on %s: %v", f.PocID, f.Target, f.Err) + }), + sdk.WithPortHandler(func(p sdk.PortEvent) { /* ... */ }), + sdk.WithHostHandler(func(h sdk.HostEvent) { /* ... */ }), + sdk.WithWebProbeHandler(func(w sdk.WebProbeEvent) { /* ... */ }), + sdk.WithProgressHandler(func(p sdk.PhaseProgress) { /* ... */ }), + sdk.WithScanInfoHandler(func(i sdk.ScanInfo) { /* ... */ }), +) ``` -You can also run the example: `examples/sdk_portscan/`. +Handlers are invoked **concurrently** from scan workers, so implementations must synchronise their own state. -## API Method Reference +### Streams -### SDKScanner Core Methods +Streams are subscription-based: **nothing is published until the subscribe method is first called**, so an unused stream costs nothing and can never stall the scan. -| Method | Description | Return Value | -|--------|-------------|--------------| -| `NewSDKScanner(opts)` | Create scanner instance | `*SDKScanner, error` | -| `Run()` | Execute scan synchronously | `error` | -| `RunAsync()` | Execute scan asynchronously | `error` | -| `GetResults()` | Get all scan results | `[]*result.Result` | -| `GetOpenPorts()` | Get open ports discovered by pre-scan | `map[string][]int` | -| `GetStats()` | Get scan statistics | `ScanStats` | -| `GetProgress()` | Get scan progress (0-100) | `float64` | -| `GetVulnerabilityCount()` | Get vulnerability count | `int` | -| `HasVulnerabilities()` | Check if vulnerabilities exist | `bool` | -| `Stop()` | Stop scanning | - | -| `Close()` | Close scanner and release resources | - | +```go +results := scanner.ResultStream() // subscribe before Start -### Dynamic Configuration Methods +scanner.Start(ctx) -| Method | Description | -|--------|-------------| -| `SetProxy(proxy)` | Dynamically set proxy | -| `SetRateLimit(n)` | Dynamically set rate limit | -| `SetConcurrency(n)` | Dynamically set concurrency | +go func() { + for r := range results { // the channel closes when the scan finishes + fmt.Println(r.PocID, r.FullTarget) + } +}() + +scanner.Wait(ctx) +``` + +| Method | Event type | +|--------|-----------| +| `ResultStream()` | `Result` | +| `PortStream()` | `PortEvent` | +| `HostStream()` | `HostEvent` | +| `WebProbeStream()` | `WebProbeEvent` | +| `ProgressStream()` | `PhaseProgress` | +| `ScanInfoStream()` | `ScanInfo` | -### OOB Related Methods +> **Important**: once subscribed, a stream must be consumed. So that findings are never silently dropped, sends **block** when the buffer fills rather than discarding data. Cancelling the context or calling `Stop` releases any blocked send. +> +> Subscribing after the scan has finished yields an already-closed channel, so a range loop terminates immediately instead of deadlocking. -| Method | Description | Return Value | -|--------|-------------|--------------| -| `IsOOBEnabled()` | Check if OOB is enabled | `bool` | -| `GetOOBStatus()` | Get OOB status information | `bool, string` | +### Advanced: the engine's own result type -### ScanStats Structure +When you need fields the stable `Result` does not expose, such as persisting the internal structure: ```go -type ScanStats struct { - StartTime time.Time // Scan start time - EndTime time.Time // Scan end time - TotalTargets int // Total target count - TotalPocs int // Total POC count - TotalScans int // Total scan tasks - CompletedScans int32 // Completed scan count - FoundVulns int32 // Found vulnerability count -} +sdk.WithRawResultHandler(func(r *result.Result) { + _ = persist(r) +}) ``` -## Advanced Usage Examples +`result.Result` is an internal type whose shape is not covered by the SDK's stability guarantee. Prefer `WithResultHandler`. -### Batch Scanning with Result Analysis +## Error handling ```go -options := afrog.NewSDKOptions() -options.TargetsFile = "targets.txt" // Read many targets from file -options.PocFile = "/path/to/pocs" -options.Severity = "high,critical" // Only scan high-risk vulnerabilities -options.Concurrency = 50 // Increase concurrency - -scanner, _ := afrog.NewSDKScanner(options) - -// Handle different severity levels differently -scanner.OnResult = func(r *result.Result) { - switch r.PocInfo.Info.Severity { - case "critical": - sendUrgentAlert(r) - case "high": - logHighRiskVuln(r) - default: - saveToDatabase(r) - } +scanner, err := sdk.New(ctx, opts...) +switch { +case errors.Is(err, sdk.ErrNoTargets): + log.Fatal("no targets specified") +case errors.Is(err, sdk.ErrPocPathNotFound): + log.Fatal("poc path could not be resolved") +case errors.Is(err, sdk.ErrInvalidOptions): + log.Fatal("invalid options") +case err != nil: + log.Fatal(err) } +``` + +| Error | Meaning | +|-------|---------| +| `ErrNoTargets` | No scan targets specified | +| `ErrNoPocs` | No executable PoCs | +| `ErrPocPathNotFound` | PoC path resolved to no files | +| `ErrAlreadyRunning` | A scan is already in progress | +| `ErrAlreadyFinished` | The scan finished; the scanner cannot be reused | +| `ErrClosed` | The scanner has been closed | +| `ErrNotStarted` | The scan has not been started | +| `ErrInvalidOptions` | Invalid option combination | +| `ErrWebhookTokenRequired` | A webhook was enabled without a token | + +`CuratedMountError` reports that the optional curated PoC source failed to mount. It is **not fatal**: the scan continues, and the error is available from `scanner.CuratedError()`. + +## Port pre-scanning + +Open ports discovered before the PoC scan are appended to the target set as `host:port`: + +```go +scanner, _ := sdk.New(ctx, + sdk.WithTargets("192.168.1.0/24"), + sdk.WithPocPaths(pocPath), + sdk.WithPortScan(sdk.PortScanOptions{ + Ports: "top", // top|full|all|80,443|1-1024 + TimeoutMs: 500, + SkipDiscovery: true, + }), + sdk.WithPortHandler(func(p sdk.PortEvent) { + fmt.Printf("open: %s:%d\n", p.Host, p.Port) + }), +) + +scanner.Execute(ctx) -scanner.Run() -results := scanner.GetResults() -generateReport(results) +open := scanner.OpenPorts() // map[string][]int ``` -### Intelligent Scan Control +## OOB (out-of-band) detection ```go -scanner.OnResult = func(r *result.Result) { - // Stop scanning when critical vulnerability found - if r.PocInfo.Info.Severity == "critical" { - fmt.Println("Critical vulnerability found, stopping scan") - scanner.Stop() - } +scanner, _ := sdk.New(ctx, + sdk.WithTargets(target), + sdk.WithPocPaths(pocPath), + sdk.WithOOB(sdk.OOBOptions{ + Adapter: "ceyeio", + Key: "your-ceye-api-token", + Domain: "your-subdomain.ceye.io", + }), +) + +// Note: this performs a live probe against the OOB service +if enabled, status := scanner.OOBStatus(); !enabled { + log.Printf("OOB unavailable: %s", status) } +``` -// Dynamically adjust scan parameters -go func() { - time.Sleep(30 * time.Second) - // Reduce rate after 30 seconds - scanner.SetRateLimit(50) -}() +| Adapter | Required fields | +|---------|-----------------| +| `ceyeio` | `Key`, `Domain` | +| `dnslogcn` | `Domain` | +| `alphalog` | `Domain`, `ApiURL` | +| `xray` | `Key`, `Domain`, `ApiURL` | +| `revsuit` | `Key`, `Domain`, `ApiURL`, `HttpURL` | + +When not configured explicitly, the SDK reads `~/.config/afrog/afrog-config.yaml`. The file is only ever read — the SDK never creates or rewrites it. + +`OOBOptions` also tunes the polling cadence. The defaults match the CLI: + +| Field | Default | Description | +|-------|---------|-------------| +| `PollInterval` | `2` | Seconds between polls of the OOB service | +| `HitRetention` | `10` | Minutes a recorded hit stays available | +| `RateLimit` | `25` | Rate limit for the OOB stage | +| `Concurrency` | `25` | Concurrency for the OOB stage | +| `FinalizeTimeout` | `-1` | Seconds to wait for late callbacks; `-1` lets each PoC decide | + +### v3 OOB PoC syntax + +```yaml +rules: + r0: + request: + method: GET + path: /?dns=ping%20{{oob.DNS}} + expression: oobCheck(oob.ProtocolDNS, 5) +expression: r0() ``` -### Multi-target Parallel Scanning +The v2 forms `set: oob: oob()`, `{{oobDNS}}`, `oobWait(...)` and `oobCheck(oob, ...)` are deprecated. PoCs using them are skipped and reported through `PocDiagnostics()`. + +## Console output + +The SDK is **completely silent** by default. For a summary, use the structured API: ```go -targets := [][]string{ - {"https://site1.com", "https://site2.com"}, - {"https://site3.com", "https://site4.com"}, -} +info := scanner.Info() +log.Printf("%d targets, %d pocs, %d tasks", + info.TotalTargets, info.TotalPocs, info.TotalScans) +``` -var wg sync.WaitGroup -results := make(chan []*result.Result, len(targets)) - -for _, targetGroup := range targets { - wg.Add(1) - go func(targets []string) { - defer wg.Done() - - options := afrog.NewSDKOptions() - options.Targets = targets - options.PocFile = pocPath - - scanner, _ := afrog.NewSDKScanner(options) - defer scanner.Close() - - scanner.Run() - results <- scanner.GetResults() - }(targetGroup) -} +Or opt into printing with `sdk.WithVerbose()`. -wg.Wait() -close(results) +## Per-task timeout -// Aggregate all results -allResults := []*result.Result{} -for groupResults := range results { - allResults = append(allResults, groupResults...) -} +`WithTimeout` bounds a single request, but a PoC with many rules can occupy a worker for far longer than any one request. `WithTaskTimeout` bounds the whole target+PoC task: + +```go +sdk.WithTaskTimeout(sdk.TaskTimeoutOptions{ + HardSec: 120, // fixed ceiling in seconds, 0 disables + Smart: true, // derive the ceiling from the PoC's content +}) ``` -## Performance Optimization Tips +`Smart` estimates the ceiling from rule count, sleeps, brute force and payloads. When both are set the **larger** value wins, so `HardSec` acts as a floor rather than an override. -### 1. Concurrency Optimization +The estimate is capped per protocol family, with the same defaults as the CLI: `VisibleCapSec` 300 (plain HTTP), `NetCapSec` 360 (tcp/udp/ssl), `GoCapSec` 420 (go PoCs). + +## Execution monitor + +The SDK equivalent of the CLI's `-pedm`, for finding slow or stuck PoCs: ```go -targetCount := len(options.Targets) +sdk.WithExecutionMonitor(sdk.ExecutionMonitorOptions{ + SlowThresholdSec: 20, // seconds after which a task counts as slow + SummaryTop: 10, // report the N slowest PoCs at the end + SummaryBy: sdk.MonitorSummaryByMax, // or MonitorSummaryByAvg +}), +sdk.WithMonitorHandler(func(line string) { + log.Println(line) +}), +``` -// Dynamically adjust concurrency based on target count -switch { -case targetCount <= 10: - options.Concurrency = 5 -case targetCount <= 100: - options.Concurrency = 25 -case targetCount <= 1000: - options.Concurrency = 50 -default: - options.Concurrency = 100 -} +Reports go only to the handlers registered with `WithMonitorHandler`; the SDK never writes them to the console. **Without a handler the monitor still runs but its output goes nowhere**, so use the two options together. + +## Resuming a scan + +The SDK equivalent of the CLI's `-resume`. The checkpoint is read at startup to skip finished work, and rewritten periodically while the scan runs: + +```go +sdk.WithCheckpoint(sdk.CheckpointOptions{ + Path: "scan.afg", + SaveInterval: 10 * time.Second, // 0 uses the 10s default +}) ``` -### 2. Memory Optimization +Progress is keyed by PoC id and target, so **the target and PoC sets must be unchanged** between runs or the skip mapping will not line up. A missing file is treated as a fresh scan rather than an error. + +Do not confuse this with `Scanner.Resume()`, which lifts a `Pause()`. + +## Sourcing targets from a search engine + +The SDK equivalent of the CLI's `-cs` / `-q` / `-qc`. Targets can come entirely from the search, with no `WithTargets` at all: ```go -// For large-scale scans, use streaming to avoid memory accumulation -options.EnableStream = true +scanner, _ := sdk.New(ctx, + sdk.WithCyberspace(sdk.CyberspaceOptions{ + Engine: sdk.CyberspaceZoomEye, + Query: `app:"tomcat"`, + Count: 100, + }), + sdk.WithPocPaths(pocPath), +) +``` -// Process results immediately, don't accumulate -scanner.OnResult = func(r *result.Result) { - processImmediately(r) - // Don't store in slices -} +Only **ZoomEye** is implemented; any other engine name returns `ErrInvalidOptions`. The API key is read from `cyberspace.zoom_eyes` in the configuration file, and `sdk.New` fails when it is missing. A search that matches nothing returns `ErrNoTargets`. + +## Target pre-probe + +The SDK equivalent of the CLI's `-mt`. It probes each target's protocol and liveness in parallel with the scan, blacklisting hosts that exceed `MaxHostError`: + +```go +sdk.WithTargetPreProbe() ``` -### 3. Network Optimization +Despite the CLI flag being named monitor-targets, it does **not** watch the targets file for changes. + +## Configuration reference + +### Targets + +| Option | Description | +|--------|-------------| +| `WithTargets(...)` | Targets to scan | +| `WithTargetsFile(path)` | File with one target per line | +| `WithCyberspace(cfg)` | Source targets from a search engine (ZoomEye only) | +| `WithTargetPreProbe()` | Probe target protocol and liveness in parallel (CLI `-mt`) | + +### PoC + +| Option | Description | +|--------|-------------| +| `WithPocPaths(...)` | File/directory/glob, append semantics | +| `WithPocPathsOnly()` | Use only the explicit PoCs | +| `WithSearch(kw)` | Keyword filter | +| `WithSeverity(sev)` | Severity filter | +| `WithExcludePocs(...)` | Exclude specific PoCs | +| `WithExcludePocsFile(path)` | Exclusion list file | + +### Performance + +| Option | Default | +|--------|---------| +| `WithConcurrency(n)` | `25` | +| `WithRateLimit(n)` | `150` | +| `WithTimeout(sec)` | `50` | +| `WithRetries(n)` | `1` | +| `WithMaxHostError(n)` | `3` | +| `WithMaxRespBodySize(mb)` | `2` | +| `WithRequestLimitPerTarget(n)` | `0` | +| `WithPolite()` / `WithBalanced()` / `WithAggressive()` | — | +| `WithAutoRequestLimit()` | — | +| `WithSmartConcurrency()` | — | +| `WithStopOnFirstMatch()` | — | + +`WithRequestLimitPerTarget`, `WithAutoRequestLimit`, `WithPolite`, `WithBalanced` and `WithAggressive` are mutually exclusive; setting more than one returns `ErrInvalidOptions`. + +### Fingerprinting and probing + +| Option | Default | +|--------|---------| +| `WithFingerprintDisabled()` | Fingerprinting is on by default | +| `WithFingerprintFilterMode(mode)` | `"strict"` (or `"opportunistic"`) | +| `WithWebProbe()` | Off by default | + +### Network + +| Option | Description | +|--------|-------------| +| `WithProxy(p)` | HTTP/SOCKS5 proxy | +| `WithHeaders(...)` | Custom headers in `"Name: value"` form | + +### Output + +| Option | Default | +|--------|---------| +| `WithRequestResponse(b)` | `true` | +| `WithMaxStoredResults(n)` | `0` (unlimited) | +| `WithStreamBuffer(n)` | `256` | +| `WithRedactedHeaders(...)` | No redaction | +| `WithVerbose()` | Silent by default | + +### Redacting sensitive data + +`Exchange` carries the full raw request and response by default, which may +include credentials such as `Authorization`, `Cookie` and `Set-Cookie`. Enable +redaction whenever results are logged, persisted or returned over an API: ```go -// Configuration for unstable networks -options.Retries = 3 -options.Timeout = 30 -options.RateLimit = 50 // Reduce request frequency +sdk.WithRedactedHeaders() // mask the default credential headers +sdk.WithRedactedHeaders("authorization", "x-token") // mask a custom set +``` + +Redaction masks the value to `[REDACTED]` in both the raw `Exchange.Request` / +`Response` and the `RequestHeaders` / `ResponseHeaders` maps, touching only +headers and never the body. It is opt-in because the raw messages are the point +of `Exchange`, and masking everything by default would weaken debugging. + +### Other + +| Option | Description | +|--------|-------------| +| `WithOOB(cfg)` | Out-of-band detection | +| `WithPortScan(cfg)` | Port pre-scan | +| `WithCurated(cfg)` | Curated PoC source | +| `WithTaskTimeout(cfg)` | Ceiling for a single target+PoC task | +| `WithExecutionMonitor(cfg)` | PoC execution duration monitor (CLI `-pedm`) | +| `WithMonitorHandler(fn)` | Receive execution monitor reports | +| `WithCheckpoint(cfg)` | Resume an interrupted scan (CLI `-resume`) | +| `WithDingtalk()` / `WithWecom()` | Webhook notifications | +| `WithOptions(o)` | Use a fully populated `Options` | + +## API reference + +### Construction + +| Method | Returns | +|--------|---------| +| `New(ctx, options...)` | `*Scanner, error` | +| `NewOptions()` | `*Options` | + +### Results + +| Method | Returns | +|--------|---------| +| `Results()` | `[]Result` | +| `ResultCount()` | `int` | +| `HasResults()` | `bool` | +| `OpenPorts()` | `map[string][]int` | +| `Stats()` | `Stats` | +| `Progress()` | `float64` | + +### PoCs and information + +| Method | Returns | +|--------|---------| +| `Pocs()` | `[]poc.Poc` | +| `PocCount()` | `int` | +| `PocDiagnostics()` | `[]config.PocLoadError` | +| `Info()` | `ScanInfo` | +| `OOBStatus()` | `bool, string` | +| `CuratedError()` | `error` | -// Use proxy pools -proxies := []string{"proxy1:8080", "proxy2:8080"} -scanner.SetProxy(proxies[rand.Intn(len(proxies))]) +## Concurrency + +**A single scanner instance is safe for concurrent use** — its methods may be called from multiple goroutines. + +**Running several scanners concurrently in one process is not supported.** The HTTP client, rate limiter and protocol probe cache are process-global, so concurrent scanners overwrite each other's proxy, timeout and rate-limit settings. + +```go +// Correct: sequential reuse +for _, group := range targetGroups { + scanner, _ := sdk.New(ctx, sdk.WithTargets(group...), sdk.WithPocPaths(pocPath)) + if err := scanner.Execute(ctx); err != nil { + log.Print(err) + } + results = append(results, scanner.Results()...) + scanner.Close() +} ``` -## Error Handling Best Practices +## Integration examples -### Complete Error Handling +### CI security gate ```go -scanner, err := afrog.NewSDKScanner(options) +scanner, err := sdk.New(ctx, + sdk.WithTargetsFile("staging-urls.txt"), + sdk.WithPocPaths("/security/pocs"), + sdk.WithSeverity("high,critical"), +) if err != nil { - switch { - case strings.Contains(err.Error(), "POC文件"): - log.Fatal("POC configuration error:", err) - case strings.Contains(err.Error(), "目标"): - log.Fatal("Target configuration error:", err) - default: - log.Fatal("Initialization failed:", err) - } + log.Fatal(err) +} +defer scanner.Close() + +if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) } -// Scan error handling -if err := scanner.Run(); err != nil { - log.Printf("Scan exception: %v", err) - - // Even with errors, partial results can be obtained - results := scanner.GetResults() - if len(results) > 0 { - fmt.Printf("Partial results obtained: %d vulnerabilities\n", len(results)) - } +if results := scanner.Results(); len(results) > 0 { + for _, v := range results { + fmt.Printf("- [%s] %s: %s\n", v.Severity, v.FullTarget, v.PocName) + } + os.Exit(1) } ``` -### Timeout and Cancellation Handling +### Timeout control ```go ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) defer cancel() -go func() { - scanner.RunAsync() -}() - -select { -case <-ctx.Done(): - scanner.Stop() - fmt.Println("Scan timed out, stopped") -case <-scanner.ResultChan: - // Normal completion +if err := scanner.Execute(ctx); err != nil { + if errors.Is(err, context.DeadlineExceeded) { + log.Println("scan timed out") + } } -``` -## Integration Examples +results := scanner.Results() // partial results remain available +``` -### Web Service Integration +### Signal handling ```go -func scanHandler(w http.ResponseWriter, r *http.Request) { - target := r.URL.Query().Get("target") - - options := afrog.NewSDKOptions() - options.Targets = []string{target} - options.PocFile = os.Getenv("POC_PATH") - - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - http.Error(w, err.Error(), 500) - return - } - defer scanner.Close() - - scanner.Run() - results := scanner.GetResults() - - // Return JSON results - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]interface{}{ - "vulnerabilities": len(results), - "results": results, - }) -} +ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) +defer stop() + +scanner.Execute(ctx) // Ctrl+C stops the scan and returns context.Canceled ``` -### CI/CD Integration +### Web service integration ```go -func main() { - options := afrog.NewSDKOptions() - options.TargetsFile = "staging-urls.txt" - options.PocFile = "/security/pocs" - options.Severity = "high,critical" - - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - os.Exit(1) - } - defer scanner.Close() - - scanner.Run() - - if scanner.HasVulnerabilities() { - fmt.Println("❌ Security vulnerabilities found, blocking deployment") - results := scanner.GetResults() - for _, r := range results { - fmt.Printf("- %s: %s\n", r.Target, r.PocInfo.Info.Name) - } - os.Exit(1) - } - - fmt.Println("✅ Security check passed") +func scanHandler(w http.ResponseWriter, r *http.Request) { + scanner, err := sdk.New(r.Context(), + sdk.WithTargets(r.URL.Query().Get("target")), + sdk.WithPocPaths(os.Getenv("POC_PATH")), + ) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + defer scanner.Close() + + if err := scanner.Execute(r.Context()); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(scanner.Results()) } ``` -## Frequently Asked Questions +## Examples -### Q: How to handle scanning of large numbers of targets? -A: Use streaming output and appropriate concurrency control: -```go -options.EnableStream = true -options.Concurrency = 50 -scanner.OnResult = func(r *result.Result) { - // Process immediately, don't accumulate - processImmediately(r) -} -``` +Every program under `examples/` runs as-is. PoC paths resolve to the repository's `pocs/afrog-pocs` automatically and can be overridden with `-pocs`: -### Q: How to ensure OOB detection works properly? -A: Check OOB status before scanning: -```go -if enabled, status := scanner.GetOOBStatus(); !enabled { - log.Printf("OOB Warning: %s", status) -} +```bash +go run ./examples/basic_scan +go run ./examples/full_output -json +go run ./examples/async_scan +go run ./examples/progress_scan +go run ./examples/oob_scan -oob dnslogcn -oob-domain your.dnslog.cn +go run ./examples/sdk_portscan -target 127.0.0.1 +go run ./examples/vuln_scan -target https://example.com +go run ./examples/port_scan -targets 127.0.0.1 ``` -### Q: How to optimize scan performance? -A: Adjust parameters based on network and target conditions: -```go -// Internal network scanning -options.Concurrency = 100 -options.RateLimit = 500 +## FAQ -// External network scanning -options.Concurrency = 25 -options.RateLimit = 150 -options.Timeout = 15 -``` +### I set a PoC directory — why didn't the built-in PoCs run? -### Q: How to handle scan interruption? -A: Use context and signal handling: -```go -c := make(chan os.Signal, 1) -signal.Notify(c, os.Interrupt) +You probably added `WithPocPathsOnly()`. Remove it to merge with the built-in PoCs. -go func() { - <-c - scanner.Stop() - fmt.Println("Scan stopped") -}() -``` +### The scan appears to hang -## Important Notes +Check whether a stream was subscribed to but is not being consumed. A subscribed stream blocks the scan once its buffer fills, so that findings are not dropped. -1. **POC Path Must Be Specified** - SDK won't automatically download or find POCs -2. **Completely Silent Operation** - No console output, suitable for program integration -3. **No File Generation** - Won't create any report files -4. **Resource Management** - Must call `Close()` to release resources -5. **Concurrency Safe** - All methods are thread-safe -6. **OOB Configuration** - Proper configuration required for out-of-band vulnerability detection +### Why does `Wait` return `context.Canceled`? -## License +The scan was cancelled by `Stop` or by the parent context. Results discovered so far are still available from `Results()`. -MIT License +### Memory grows quickly on large scans ---- +```go +sdk.WithRequestResponse(false), +sdk.WithMaxStoredResults(1000), +``` -For more examples and detailed documentation, please refer to the example code in the `examples/` directory. +Process results in a handler instead of relying on `Results()` to accumulate them. + +## License + +MIT License diff --git "a/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" "b/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" index 469ed6408..20375b86d 100644 --- "a/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" +++ "b/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" @@ -2,14 +2,18 @@ ## 概述 -Afrog SDK 提供了一个简洁、高效的 Go 编程接口，专为集成漏洞扫描功能而设计。SDK 具有以下核心特性： +Afrog SDK 是把漏洞扫描能力嵌入自己程序的 Go 接口，包路径为 `github.com/zan8in/afrog/v3/pkg/sdk`。 -### 🚀 核心特性 -- ✅ **结构化返回** - 直接返回 Go 结构体，便于程序处理 -- ✅ **实时结果流** - 支持同步回调和异步流式输出 -- ✅ **OOB 检测支持** - 完整的带外检测配置和管理 -- ✅ **详细统计信息** - 提供扫描进度、性能和结果统计 -- ✅ **并发安全** - 所有 API 都是线程安全的 +### 核心特性 + +- **结构化返回** —— 纯 Go 结构体，可直接 `json.Marshal` +- **完整数据输出** —— 每一步扫描的原始请求与响应报文都可获取 +- **灵活的 PoC 输入** —— 单个文件、目录（递归）、glob 通配符 +- **同步与异步** —— `Execute` 同步阻塞；`Start` + `Wait`/`Done` 异步 +- **回调与流** —— 多回调注册，或按需订阅事件通道 +- **默认静默** —— 不向 stdout/stderr 输出任何内容 +- **类型化错误** —— 使用 `errors.Is` 判断失败原因 +- **资源可控** —— `Close` 释放所有后台协程，无泄漏 ## 安装 @@ -19,617 +23,750 @@ go get -u github.com/zan8in/afrog/v3 ## 快速开始 -### 基础扫描示例 - -最简单的使用方式，适合快速集成： - ```go package main import ( - "fmt" - "log" - "path/filepath" - "github.com/zan8in/afrog/v3" + "context" + "fmt" + "log" + + "github.com/zan8in/afrog/v3/pkg/sdk" ) func main() { - // 创建扫描选项 - options := afrog.NewSDKOptions() - - // 设置扫描目标 - options.Targets = []string{"https://www.example.com"} - - // 设置 POC 路径（必须） - pocPath, _ := filepath.Abs("./pocs/afrog-pocs") - options.PocFile = pocPath - - // 创建扫描器 - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - log.Fatal(err) - } - defer scanner.Close() - - // 执行扫描 - scanner.Run() - - // 获取结果 - results := scanner.GetResults() - fmt.Printf("发现 %d 个漏洞\n", len(results)) + ctx := context.Background() + + scanner, err := sdk.New(ctx, + sdk.WithTargets("https://example.com"), + sdk.WithPocPaths("./pocs/afrog-pocs"), + ) + if err != nil { + log.Fatal(err) + } + defer scanner.Close() + + if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) + } + + for _, r := range scanner.Results() { + fmt.Printf("[%s] %s - %s\n", r.Severity, r.FullTarget, r.PocName) + } } ``` -## SDK 配置选项详解 +## PoC 输入 -### SDKOptions 结构体 +`WithPocPaths` 支持三种形式，可以混用、可以传多个： ```go -type SDKOptions struct { - // ========== 目标配置 ========== - Targets []string // 扫描目标列表 - TargetsFile string // 目标文件路径 - - // ========== POC 配置 ========== - PocFile string // POC 文件或目录路径（必须） - Search string // POC 搜索关键词 - Severity string // 严重程度过滤 - - // ========== 性能配置 ========== - RateLimit int // 请求速率限制 (默认: 150) - Concurrency int // 并发数 (默认: 25) - Retries int // 重试次数 (默认: 1) - Timeout int // 超时时间秒 (默认: 10) - MaxHostError int // 主机最大错误数 (默认: 3) - - // ========== PortScan 预扫描配置 ========== - PortScan bool // 启用端口预扫描（等价于 CLI 的 -ps） - PSPorts string // 预扫描端口定义：top/full/all/80,443/1-1024 等（等价于 -p） - PSRateLimit int // 预扫描速率限制（等价于 -prate） - PSTimeout int // 预扫描超时（毫秒，等价于 -ptimeout） - PSRetries int // 预扫描重试（等价于 -ptries） - PSSkipDiscovery bool // 跳过存活探测（等价于 -Pn） - PSS4Chunk int // 全端口扫描 chunk（等价于 --ps-s4-chunk） - - // ========== 网络配置 ========== - Proxy string // HTTP/SOCKS5 代理 - - // ========== OOB 配置 ========== - EnableOOB bool // 是否启用 OOB 检测 - OOB string // OOB 适配器类型 - OOBKey string // OOB API 密钥 - OOBDomain string // OOB 域名 - OOBApiUrl string // OOB API 地址 - OOBHttpUrl string // OOB HTTP 地址 - - // ========== 输出配置 ========== - EnableStream bool // 启用流式输出 -} +sdk.WithPocPaths( + "/path/to/single.yaml", // 单个文件 + "/path/to/pocs", // 目录（递归查找 .yaml/.yml） + "/path/to/pocs/*.yaml", // glob 通配符 +) ``` -### 配置选项说明 - -#### 目标配置 -- `Targets`: 直接指定扫描目标列表 -- `TargetsFile`: 从文件读取目标列表（每行一个） - -#### POC 配置 -- `PocFile`: **必须**指定 POC 文件或目录路径 -- `Search`: 按关键词过滤 POC，如 "tomcat,phpinfo" -- `Severity`: 按严重程度过滤，如 "high,critical" +### 追加还是独占 -#### 性能调优 -- `Concurrency`: 并发扫描线程数，建议根据目标数量调整 -- `RateLimit`: 每秒请求数限制，避免触发防护 -- `Timeout`: 单个请求超时时间 -- `Retries`: 失败重试次数 - -## 核心功能示例 +| 配置 | 行为 | +|-----|-----| +| `WithPocPaths(...)` | **追加**：与内置 PoC、curated、my、local 合并，同名时以显式路径优先 | +| `WithPocPaths(...)` + `WithPocPathsOnly()` | **独占**：只使用显式指定的 PoC | -### 1. 实时结果回调 +### 检查加载结果 -在发现漏洞时立即处理： +在发起任何网络请求之前，可以先确认加载到了什么： ```go -scanner.OnResult = func(r *result.Result) { - fmt.Printf("发现漏洞: %s - %s [%s]\n", - r.Target, - r.PocInfo.Info.Name, - r.PocInfo.Info.Severity) - - // 立即处理逻辑 - if r.PocInfo.Info.Severity == "critical" { - sendAlert(r) - } +fmt.Printf("已加载 %d 个 PoC\n", scanner.PocCount()) + +for _, p := range scanner.Pocs() { + fmt.Println(p.Id, p.Info.Name) } -scanner.Run() +// 哪些 PoC 被跳过了，以及为什么 +for _, d := range scanner.PocDiagnostics() { + fmt.Printf("跳过 %s：%s\n", d.Path, d.Reason) +} ``` -### 2. 进度监控 +`PocLoadError.Reason` 的取值： + +| 常量 | 含义 | +|-----|-----| +| `config.PocLoadNotFound` | 路径不存在或通配符没匹配到文件 | +| `config.PocLoadReadFailed` | 文件读取失败 | +| `config.PocLoadParseFailed` | YAML 解析失败 | +| `config.PocLoadLegacyOOB` | 使用了已废弃的 v2 OOB 语法 | + +## 完整数据输出 -实时监控扫描进度： +`Results()` 返回 `sdk.Result`，其中 `Exchanges` 携带每一步的完整请求/响应： ```go -go func() { - ticker := time.NewTicker(1 * time.Second) - defer ticker.Stop() - - for range ticker.C { - progress := scanner.GetProgress() - stats := scanner.GetStats() - fmt.Printf("进度: %.2f%% (%d/%d) 发现漏洞: %d\n", - progress, - stats.CompletedScans, - stats.TotalScans, - stats.FoundVulns) - } -}() +for _, r := range scanner.Results() { + fmt.Printf("%s [%s] %s\n", r.PocID, r.Severity, r.FullTarget) -scanner.Run() -``` + for _, ex := range r.Exchanges { + fmt.Printf("%s %s -> %d (%d ms)\n", ex.Method, ex.URL, ex.StatusCode, ex.LatencyMs) + + fmt.Println("--- 原始请求 ---") + fmt.Println(ex.Request) + + fmt.Println("--- 原始响应 ---") + fmt.Println(ex.Response) -### 3. 异步扫描与流式输出 + if ex.BodyTruncated { + fmt.Println("警告：响应体在 MaxRespBodySize 上限处被截断") + } + } +} +``` -非阻塞扫描，实时获取结果： +### Result 结构 ```go -options.EnableStream = true -scanner, _ := afrog.NewSDKScanner(options) - -// 启动异步扫描 -scanner.RunAsync() - -// 从通道读取实时结果 -for result := range scanner.ResultChan { - fmt.Printf("实时发现: %s - %s\n", - result.Target, - result.PocInfo.Info.Name) - - // 实时处理每个结果 - processResult(result) +type Result struct { + PocID string `json:"poc_id"` + PocName string `json:"poc_name,omitempty"` + Severity string `json:"severity,omitempty"` + Author string `json:"author,omitempty"` + Description string `json:"description,omitempty"` + Reference []string `json:"reference,omitempty"` + Tags []string `json:"tags,omitempty"` + + CveID string `json:"cve_id,omitempty"` + CweID string `json:"cwe_id,omitempty"` + CvssScore float64 `json:"cvss_score,omitempty"` + CvssMetrics string `json:"cvss_metrics,omitempty"` + + Target string `json:"target"` + FullTarget string `json:"full_target,omitempty"` + + Extractors map[string]string `json:"extractors,omitempty"` + Fingerprints []Fingerprint `json:"fingerprints,omitempty"` + Exchanges []Exchange `json:"exchanges,omitempty"` + + FoundAt time.Time `json:"found_at"` } ``` -### 4. OOB（带外）检测配置 +### Exchange 结构 -#### CEYE.io 配置（推荐） ```go -options.EnableOOB = true -options.OOB = "ceyeio" -options.OOBKey = "your-ceye-api-token" -options.OOBDomain = "your-subdomain.ceye.io" +type Exchange struct { + Request string `json:"request,omitempty"` // 原始请求报文 + Response string `json:"response,omitempty"` // 原始响应报文 + + Method string `json:"method,omitempty"` + URL string `json:"url,omitempty"` + RequestHeaders map[string]string `json:"request_headers,omitempty"` + RequestBody string `json:"request_body,omitempty"` + StatusCode int `json:"status_code,omitempty"` + ResponseHeaders map[string]string `json:"response_headers,omitempty"` + ResponseBody string `json:"response_body,omitempty"` + ContentType string `json:"content_type,omitempty"` + LatencyMs int64 `json:"latency_ms,omitempty"` + + Matched bool `json:"matched"` + BodyTruncated bool `json:"body_truncated,omitempty"` + BruteTruncated bool `json:"brute_truncated,omitempty"` + BruteRequests int `json:"brute_requests,omitempty"` +} ``` -#### DNSLog.cn 配置（免费） +原始报文是字符串而不是 `[]byte`，可以直接序列化，不会变成 base64： + ```go -options.EnableOOB = true -options.OOB = "dnslogcn" -options.OOBDomain = "your.dnslog.cn" +data, err := json.MarshalIndent(scanner.Results(), "", " ") ``` -#### 其他 OOB 服务 +### 控制内存占用 + ```go -// Alphalog -options.OOB = "alphalog" -options.OOBDomain = "your.alphalog.cn" -options.OOBApiUrl = "https://api.alphalog.cn" +sdk.WithRequestResponse(false), // 不保留 Exchanges +sdk.WithMaxStoredResults(1000), // 最多累积 1000 条 +``` + +`MaxStoredResults` 只限制内部累积，**不影响回调和流**，所有结果仍会被推送出来。 + +### 响应体截断 -// XRay -options.OOB = "xray" -options.OOBDomain = "your.xray.domain" -options.OOBApiUrl = "http://xray-api:8777" -options.OOBKey = "your-xray-token" +响应体读取上限由 `MaxRespBodySize` 控制（默认 2 MB）。超出部分会被丢弃，此时 `Exchange.BodyTruncated` 为 `true`，据此可以判断拿到的是不是完整响应。 + +```go +sdk.WithMaxRespBodySize(10) // 提高到 10 MB ``` -#### OOB 状态检查 +## 同步与异步 + +### 同步 + ```go -if oobEnabled, oobStatus := scanner.GetOOBStatus(); oobEnabled { - fmt.Printf("✓ OOB 状态: %s\n", oobStatus) -} else { - fmt.Printf("✗ OOB 状态: %s\n", oobStatus) +if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) } +results := scanner.Results() ``` -### 5. 端口预扫描（PortScan） +### 异步 + +```go +if err := scanner.Start(ctx); err != nil { + log.Fatal(err) +} -SDK 支持在 PoC 扫描之前做一次端口预扫描：扫描到的开放端口会自动追加进内部 Targets（以 `host:port` 形式），后续 PoC 会按新的目标集合执行。 +go func() { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + fmt.Printf("进度: %.1f%%\n", scanner.Progress()) + case <-scanner.Done(): + return + } + } +}() -SDK 模式下不会默认把开放端口输出到控制台，可以通过回调或获取结果来消费。 +if err := scanner.Wait(ctx); err != nil { + log.Printf("扫描出错: %v", err) +} +``` -```go -options := afrog.NewSDKOptions() -options.Targets = []string{"1.2.3.4"} -options.PocFile = pocPath +### 生命周期方法 -options.PortScan = true -options.PSPorts = "top" // 或 "full"/"all"/"80,443"/"1-1024" -options.PSSkipDiscovery = true -options.PSTimeout = 500 +| 方法 | 说明 | +|-----|-----| +| `Execute(ctx)` | 同步执行，直到扫描结束才返回 | +| `Start(ctx)` | 异步启动，立即返回 | +| `Wait(ctx)` | 阻塞等待扫描结束，返回扫描错误 | +| `Done()` | 返回扫描结束时关闭的通道 | +| `Err()` | 返回扫描错误，未结束时为 nil | +| `Stop()` | 请求停止，立即返回 | +| `Close()` | 停止扫描、等待协程退出、释放全部资源 | +| `Pause()` / `Resume()` / `IsPaused()` | 暂停控制 | +| `IsStopping()` / `IsRunning()` | 状态查询 | + +扫描器是**一次性**的： -scanner, _ := afrog.NewSDKScanner(options) +```go +scanner.Execute(ctx) // 第一次：正常 +scanner.Execute(ctx) // 第二次：返回 ErrAlreadyFinished +``` -scanner.OnPort = func(host string, port int) { - fmt.Printf("open: %s:%d\n", host, port) -} +需要重新扫描请新建实例。`Close()` 是幂等的，可以在 `New` 之后立即 `defer`。 + +## 回调与流 -scanner.Run() +### 回调 -open := scanner.GetOpenPorts() -_ = open +```go +scanner, _ := sdk.New(ctx, + sdk.WithResultHandler(saveToDatabase), + sdk.WithResultHandler(sendAlert), // 可注册多个 + sdk.WithFailureHandler(func(f sdk.Failure) { + log.Printf("PoC %s 在 %s 上失败: %v", f.PocID, f.Target, f.Err) + }), + sdk.WithPortHandler(func(p sdk.PortEvent) { /* ... */ }), + sdk.WithHostHandler(func(h sdk.HostEvent) { /* ... */ }), + sdk.WithWebProbeHandler(func(w sdk.WebProbeEvent) { /* ... */ }), + sdk.WithProgressHandler(func(p sdk.PhaseProgress) { /* ... */ }), + sdk.WithScanInfoHandler(func(i sdk.ScanInfo) { /* ... */ }), +) ``` -也可以通过 `PortChan` 异步消费端口预扫描结果：启用 `PortScan` 时会自动初始化该通道，扫描结束后会自动关闭。 +回调由扫描工作协程**并发触发**，实现方需要自行保证并发安全。 + +### 流 + +流是按需订阅的：**首次调用订阅方法之前不会产生任何数据**，因此没用到的流零开销、也绝不会阻塞扫描。 ```go -options := afrog.NewSDKOptions() -options.Targets = []string{"1.2.3.4"} -options.PocFile = pocPath -options.PortScan = true +results := scanner.ResultStream() // 在 Start 之前订阅 -scanner, _ := afrog.NewSDKScanner(options) +scanner.Start(ctx) -_ = scanner.RunAsync() +go func() { + for r := range results { // 扫描结束时通道关闭，range 自然退出 + fmt.Println(r.PocID, r.FullTarget) + } +}() -for r := range scanner.PortChan { - fmt.Printf("open: %s:%d\n", r.Host, r.Port) -} +scanner.Wait(ctx) ``` -也可以直接运行示例：`examples/sdk_portscan/`。 +| 方法 | 事件类型 | +|-----|---------| +| `ResultStream()` | `Result` | +| `PortStream()` | `PortEvent` | +| `HostStream()` | `HostEvent` | +| `WebProbeStream()` | `WebProbeEvent` | +| `ProgressStream()` | `PhaseProgress` | +| `ScanInfoStream()` | `ScanInfo` | -## API 方法参考 +> **重要**：一旦订阅就必须消费。为了保证漏洞不被静默丢弃，通道写满时发送会**阻塞**而不是丢弃数据。取消 context 或调用 `Stop()` 会释放被阻塞的发送。 +> +> 扫描结束后再订阅会得到一个已关闭的通道，`range` 立即退出，不会死锁。 -### SDKScanner 核心方法 - -| 方法 | 描述 | 返回值 | -|-----|-----|-------| -| `NewSDKScanner(opts)` | 创建扫描器实例 | `*SDKScanner, error` | -| `Run()` | 同步执行扫描 | `error` | -| `RunAsync()` | 异步执行扫描 | `error` | -| `GetResults()` | 获取所有扫描结果 | `[]*result.Result` | -| `GetOpenPorts()` | 获取预扫描开放端口 | `map[string][]int` | -| `GetStats()` | 获取扫描统计信息 | `ScanStats` | -| `GetProgress()` | 获取扫描进度(0-100) | `float64` | -| `GetVulnerabilityCount()` | 获取漏洞数量 | `int` | -| `HasVulnerabilities()` | 检查是否有漏洞 | `bool` | -| `Stop()` | 停止扫描 | - | -| `Close()` | 关闭扫描器，释放资源 | - | - -### 动态配置方法 - -| 方法 | 描述 | -|-----|-----| -| `SetProxy(proxy)` | 动态设置代理 | -| `SetRateLimit(n)` | 动态设置速率限制 | -| `SetConcurrency(n)` | 动态设置并发数 | +### 高级：拿到引擎原始结果 + +需要 `Result` 未暴露的字段（例如按引擎内部结构持久化）时： -### OOB 相关方法 +```go +sdk.WithRawResultHandler(func(r *result.Result) { + _ = persist(r) +}) +``` -| 方法 | 描述 | 返回值 | -|-----|-----|-------| -| `IsOOBEnabled()` | 检查是否启用 OOB | `bool` | -| `GetOOBStatus()` | 获取 OOB 状态信息 | `bool, string` | +`result.Result` 是内部类型，其结构不在 SDK 的稳定性保证范围内，优先使用 `WithResultHandler`。 -### ScanStats 统计结构 +## 错误处理 ```go -type ScanStats struct { - StartTime time.Time // 扫描开始时间 - EndTime time.Time // 扫描结束时间 - TotalTargets int // 总目标数 - TotalPocs int // 总 POC 数 - TotalScans int // 总扫描任务数 - CompletedScans int32 // 已完成扫描数 - FoundVulns int32 // 发现的漏洞数 +scanner, err := sdk.New(ctx, opts...) +switch { +case errors.Is(err, sdk.ErrNoTargets): + log.Fatal("未指定扫描目标") +case errors.Is(err, sdk.ErrPocPathNotFound): + log.Fatal("PoC 路径无法解析") +case errors.Is(err, sdk.ErrInvalidOptions): + log.Fatal("选项配置非法") +case err != nil: + log.Fatal(err) } ``` -## 高级用法示例 +| 错误 | 含义 | +|-----|-----| +| `ErrNoTargets` | 未指定扫描目标 | +| `ErrNoPocs` | 没有可执行的 PoC | +| `ErrPocPathNotFound` | PoC 路径无法解析为任何文件 | +| `ErrAlreadyRunning` | 扫描已在进行中 | +| `ErrAlreadyFinished` | 扫描已结束，扫描器不可复用 | +| `ErrClosed` | 扫描器已关闭 | +| `ErrNotStarted` | 尚未启动扫描 | +| `ErrInvalidOptions` | 选项组合非法 | +| `ErrWebhookTokenRequired` | 启用了 webhook 但未配置 token | + +`CuratedMountError` 表示可选的 curated PoC 源挂载失败，它**不是致命错误**，扫描会继续，可通过 `scanner.CuratedError()` 查询。 -### 批量扫描与结果分析 +## 端口预扫描 + +发现的开放端口会以 `host:port` 形式追加为扫描目标： ```go -options := afrog.NewSDKOptions() -options.TargetsFile = "targets.txt" // 从文件读取大量目标 -options.PocFile = "/path/to/pocs" -options.Severity = "high,critical" // 只扫描高危漏洞 -options.Concurrency = 50 // 提高并发数 - -scanner, _ := afrog.NewSDKScanner(options) - -// 分类处理不同严重程度的漏洞 -scanner.OnResult = func(r *result.Result) { - switch r.PocInfo.Info.Severity { - case "critical": - sendUrgentAlert(r) - case "high": - logHighRiskVuln(r) - default: - saveToDatabase(r) - } -} +scanner, _ := sdk.New(ctx, + sdk.WithTargets("192.168.1.0/24"), + sdk.WithPocPaths(pocPath), + sdk.WithPortScan(sdk.PortScanOptions{ + Ports: "top", // top|full|all|80,443|1-1024 + TimeoutMs: 500, + SkipDiscovery: true, + }), + sdk.WithPortHandler(func(p sdk.PortEvent) { + fmt.Printf("open: %s:%d\n", p.Host, p.Port) + }), +) -scanner.Run() -results := scanner.GetResults() -generateReport(results) +scanner.Execute(ctx) + +open := scanner.OpenPorts() // map[string][]int ``` -### 智能扫描控制 +## OOB（带外）检测 ```go -scanner.OnResult = func(r *result.Result) { - // 发现严重漏洞时停止扫描 - if r.PocInfo.Info.Severity == "critical" { - fmt.Println("发现严重漏洞，停止扫描") - scanner.Stop() - } +scanner, _ := sdk.New(ctx, + sdk.WithTargets(target), + sdk.WithPocPaths(pocPath), + sdk.WithOOB(sdk.OOBOptions{ + Adapter: "ceyeio", + Key: "your-ceye-api-token", + Domain: "your-subdomain.ceye.io", + }), +) + +// 注意：这会对 OOB 服务发起一次真实的网络探测 +if enabled, status := scanner.OOBStatus(); !enabled { + log.Printf("OOB 不可用: %s", status) } +``` -// 动态调整扫描参数 -go func() { - time.Sleep(30 * time.Second) - // 30秒后降低速率 - scanner.SetRateLimit(50) -}() +| Adapter | 必填字段 | +|---------|---------| +| `ceyeio` | `Key`、`Domain` | +| `dnslogcn` | `Domain` | +| `alphalog` | `Domain`、`ApiURL` | +| `xray` | `Key`、`Domain`、`ApiURL` | +| `revsuit` | `Key`、`Domain`、`ApiURL`、`HttpURL` | + +未显式配置时，SDK 会尝试从 `~/.config/afrog/afrog-config.yaml` 读取。该文件**只读**，SDK 不会创建或改写它。 + +`OOBOptions` 还可以调整轮询节奏，默认值与 CLI 一致： + +| 字段 | 默认 | 说明 | +|------|------|------| +| `PollInterval` | `2` | 轮询 OOB 服务的间隔（秒） | +| `HitRetention` | `10` | 命中记录的保留时长（分钟） | +| `RateLimit` | `25` | OOB 阶段的速率限制 | +| `Concurrency` | `25` | OOB 阶段的并发 | +| `FinalizeTimeout` | `-1` | 收敛等待上限（秒），`-1` 表示由 PoC 自身决定 | + +### v3 的 OOB PoC 语法 + +```yaml +rules: + r0: + request: + method: GET + path: /?dns=ping%20{{oob.DNS}} + expression: oobCheck(oob.ProtocolDNS, 5) +expression: r0() ``` -### 多目标并行扫描 +v2 时代的 `set: oob: oob()`、`{{oobDNS}}`、`oobWait(...)`、`oobCheck(oob, ...)` 均已废弃，使用旧语法的 PoC 会被跳过并出现在 `PocDiagnostics()` 中。 + +## 控制台输出 + +SDK **默认完全静默**。需要摘要时用结构化 API： ```go -targets := [][]string{ - {"https://site1.com", "https://site2.com"}, - {"https://site3.com", "https://site4.com"}, -} +info := scanner.Info() +log.Printf("目标 %d 个，PoC %d 个，任务 %d 个", + info.TotalTargets, info.TotalPocs, info.TotalScans) +``` -var wg sync.WaitGroup -results := make(chan []*result.Result, len(targets)) - -for _, targetGroup := range targets { - wg.Add(1) - go func(targets []string) { - defer wg.Done() - - options := afrog.NewSDKOptions() - options.Targets = targets - options.PocFile = pocPath - - scanner, _ := afrog.NewSDKScanner(options) - defer scanner.Close() - - scanner.Run() - results <- scanner.GetResults() - }(targetGroup) -} +或显式开启打印：`sdk.WithVerbose()`。 -wg.Wait() -close(results) +## 任务级超时 -// 汇总所有结果 -allResults := []*result.Result{} -for groupResults := range results { - allResults = append(allResults, groupResults...) -} +单次请求的超时用 `WithTimeout`，但一个规则很多的 PoC 可能远超单次请求时长地占住 worker。`WithTaskTimeout` 给「单个目标 + 单个 PoC」这一整个任务加上限： + +```go +sdk.WithTaskTimeout(sdk.TaskTimeoutOptions{ + HardSec: 120, // 固定上限（秒），0 表示不限 + Smart: true, // 依据 PoC 内容估算上限 +}) ``` -## 性能优化建议 +`Smart` 会根据规则数量、sleep、爆破、payload 等估算超时。两者同时设置时**取较大值**，也就是 `HardSec` 起下限作用而不是覆盖估算值。 + +估算值按协议族分别设上限，默认与 CLI 一致：`VisibleCapSec` 300（普通 HTTP）、`NetCapSec` 360（tcp/udp/ssl）、`GoCapSec` 420（go 类 PoC）。 -### 1. 并发数优化 +## 执行耗时监控 + +对应 CLI 的 `-pedm`，用于定位跑得慢或卡住的 PoC： ```go -targetCount := len(options.Targets) +sdk.WithExecutionMonitor(sdk.ExecutionMonitorOptions{ + SlowThresholdSec: 20, // 超过多少秒算慢任务 + SummaryTop: 10, // 结束时列出最慢的 N 个 PoC + SummaryBy: sdk.MonitorSummaryByMax, // 或 MonitorSummaryByAvg +}), +sdk.WithMonitorHandler(func(line string) { + log.Println(line) +}), +``` -// 根据目标数量动态调整并发数 -switch { -case targetCount <= 10: - options.Concurrency = 5 -case targetCount <= 100: - options.Concurrency = 25 -case targetCount <= 1000: - options.Concurrency = 50 -default: - options.Concurrency = 100 -} +监控内容只会送到 `WithMonitorHandler` 注册的回调，SDK 不会打印到控制台。**不注册回调时监控照常运行但输出无处可去**，所以这两个选项应当配套使用。 + +## 断点续扫 + +对应 CLI 的 `-resume`。启动时读取检查点跳过已完成的任务，扫描过程中周期性回写： + +```go +sdk.WithCheckpoint(sdk.CheckpointOptions{ + Path: "scan.afg", + SaveInterval: 10 * time.Second, // 0 表示使用默认的 10 秒 +}) ``` -### 2. 内存优化 +进度以「PoC id + 目标」为键记录，因此**续扫时目标集与 PoC 集必须与中断前一致**，否则跳过关系会错位。文件不存在时视为全新扫描，不报错。 + +注意与 `Scanner.Resume()` 区分：后者是解除 `Pause()` 的暂停，与断点续扫无关。 + +## 从空间测绘获取目标 + +对应 CLI 的 `-cs` / `-q` / `-qc`。目标可以完全来自搜索，不必再传 `WithTargets`： ```go -// 对于大规模扫描，使用流式输出避免内存积累 -options.EnableStream = true +scanner, _ := sdk.New(ctx, + sdk.WithCyberspace(sdk.CyberspaceOptions{ + Engine: sdk.CyberspaceZoomEye, + Query: `app:"tomcat"`, + Count: 100, + }), + sdk.WithPocPaths(pocPath), +) +``` -// 及时处理结果，不要积累 -scanner.OnResult = func(r *result.Result) { - processImmediately(r) - // 不要存储到切片中 -} +目前**只实现了 ZoomEye**，传其他引擎名会返回 `ErrInvalidOptions`。API Key 从配置文件的 `cyberspace.zoom_eyes` 读取，缺失时 `sdk.New` 返回错误。搜索命中为 0 时返回 `ErrNoTargets`。 + +## 目标预探测 + +对应 CLI 的 `-mt`。它会在扫描的同时并发探测每个目标的协议与存活情况，错误次数超过 `MaxHostError` 的主机会被拉黑： + +```go +sdk.WithTargetPreProbe() ``` -### 3. 网络优化 +尽管 CLI 的参数名叫 monitor-targets，它并**不会**监视目标文件的变化。 + +## 配置选项完整列表 + +### 目标 + +| 选项 | 说明 | +|-----|-----| +| `WithTargets(...)` | 扫描目标列表 | +| `WithTargetsFile(path)` | 目标文件，每行一个 | +| `WithCyberspace(cfg)` | 从空间测绘搜索获取目标（目前仅 ZoomEye） | +| `WithTargetPreProbe()` | 并发预探测目标协议与存活（CLI `-mt`） | + +### PoC + +| 选项 | 说明 | +|-----|-----| +| `WithPocPaths(...)` | 文件/目录/glob，追加语义 | +| `WithPocPathsOnly()` | 只用显式指定的 PoC | +| `WithSearch(kw)` | 关键词过滤 | +| `WithSeverity(sev)` | 严重程度过滤 | +| `WithExcludePocs(...)` | 排除指定 PoC | +| `WithExcludePocsFile(path)` | 排除列表文件 | + +### 性能 + +| 选项 | 默认值 | +|-----|-------| +| `WithConcurrency(n)` | `25` | +| `WithRateLimit(n)` | `150` | +| `WithTimeout(sec)` | `50` | +| `WithRetries(n)` | `1` | +| `WithMaxHostError(n)` | `3` | +| `WithMaxRespBodySize(mb)` | `2` | +| `WithRequestLimitPerTarget(n)` | `0` | +| `WithPolite()` / `WithBalanced()` / `WithAggressive()` | — | +| `WithAutoRequestLimit()` | — | +| `WithSmartConcurrency()` | — | +| `WithStopOnFirstMatch()` | — | + +`WithRequestLimitPerTarget`、`WithAutoRequestLimit`、`WithPolite`、`WithBalanced`、`WithAggressive` 五者互斥，同时设置多个会返回 `ErrInvalidOptions`。 + +### 指纹与探测 + +| 选项 | 默认值 | +|-----|-------| +| `WithFingerprintDisabled()` | 指纹默认开启 | +| `WithFingerprintFilterMode(mode)` | `"strict"`（可选 `"opportunistic"`） | +| `WithWebProbe()` | 默认关闭 | + +### 网络 + +| 选项 | 说明 | +|-----|-----| +| `WithProxy(p)` | HTTP/SOCKS5 代理 | +| `WithHeaders(...)` | 自定义请求头，格式 `"Name: value"` | + +### 输出 + +| 选项 | 默认值 | +|-----|-------| +| `WithRequestResponse(b)` | `true` | +| `WithMaxStoredResults(n)` | `0`（不限） | +| `WithStreamBuffer(n)` | `256` | +| `WithRedactedHeaders(...)` | 默认不脱敏 | +| `WithVerbose()` | 默认静默 | + +### 敏感信息脱敏 + +`Exchange` 默认携带完整的原始请求/响应，其中可能包含 `Authorization`、`Cookie`、`Set-Cookie` 等凭证。如果结果会写日志、落库或经 API 返回，应开启脱敏： ```go -// 网络不稳定时的配置 -options.Retries = 3 -options.Timeout = 30 -options.RateLimit = 50 // 降低请求频率 +sdk.WithRedactedHeaders() // 脱敏默认的凭证类头 +sdk.WithRedactedHeaders("authorization", "x-token") // 自定义要脱敏的头 +``` + +脱敏会同时作用于 `Exchange.Request`/`Response` 原始报文和 `RequestHeaders`/`ResponseHeaders`，把对应值替换为 `[REDACTED]`，只影响头部、不触碰响应体。脱敏是**可选**的，因为原始报文正是 `Exchange` 的核心价值，默认全脱敏会削弱调试能力。 + +### 其他 + +| 选项 | 说明 | +|-----|-----| +| `WithOOB(cfg)` | 带外检测 | +| `WithPortScan(cfg)` | 端口预扫描 | +| `WithCurated(cfg)` | curated PoC 源 | +| `WithTaskTimeout(cfg)` | 单个目标+PoC 任务的超时上限 | +| `WithExecutionMonitor(cfg)` | PoC 执行耗时监控（CLI `-pedm`） | +| `WithMonitorHandler(fn)` | 接收耗时监控输出 | +| `WithCheckpoint(cfg)` | 断点续扫（CLI `-resume`） | +| `WithDingtalk()` / `WithWecom()` | webhook 通知 | +| `WithOptions(o)` | 直接使用一份完整 `Options` | + +## API 方法参考 + +### 构造 + +| 方法 | 返回值 | +|-----|-------| +| `New(ctx, options...)` | `*Scanner, error` | +| `NewOptions()` | `*Options` | + +### 结果 + +| 方法 | 返回值 | +|-----|-------| +| `Results()` | `[]Result` | +| `ResultCount()` | `int` | +| `HasResults()` | `bool` | +| `OpenPorts()` | `map[string][]int` | +| `Stats()` | `Stats` | +| `Progress()` | `float64` | + +### PoC 与信息 -// 使用代理池 -proxies := []string{"proxy1:8080", "proxy2:8080"} -scanner.SetProxy(proxies[rand.Intn(len(proxies))]) +| 方法 | 返回值 | +|-----|-------| +| `Pocs()` | `[]poc.Poc` | +| `PocCount()` | `int` | +| `PocDiagnostics()` | `[]config.PocLoadError` | +| `Info()` | `ScanInfo` | +| `OOBStatus()` | `bool, string` | +| `CuratedError()` | `error` | + +## 并发限制 + +**单个扫描器实例是并发安全的**，可以从多个协程调用它的方法。 + +**但同一进程内不支持多个扫描器并行运行。** HTTP 客户端、限速器和协议探测缓存都是进程级全局状态，并行的扫描器会互相覆盖代理、超时和速率配置。 + +```go +// 正确：串行复用 +for _, group := range targetGroups { + scanner, _ := sdk.New(ctx, sdk.WithTargets(group...), sdk.WithPocPaths(pocPath)) + if err := scanner.Execute(ctx); err != nil { + log.Print(err) + } + results = append(results, scanner.Results()...) + scanner.Close() +} ``` -## 错误处理最佳实践 +## 集成示例 -### 完整的错误处理 +### CI 安全门禁 ```go -scanner, err := afrog.NewSDKScanner(options) +scanner, err := sdk.New(ctx, + sdk.WithTargetsFile("staging-urls.txt"), + sdk.WithPocPaths("/security/pocs"), + sdk.WithSeverity("high,critical"), +) if err != nil { - switch { - case strings.Contains(err.Error(), "POC文件"): - log.Fatal("POC 配置错误:", err) - case strings.Contains(err.Error(), "目标"): - log.Fatal("目标配置错误:", err) - default: - log.Fatal("初始化失败:", err) - } + log.Fatal(err) +} +defer scanner.Close() + +if err := scanner.Execute(ctx); err != nil { + log.Fatal(err) } -// 扫描错误处理 -if err := scanner.Run(); err != nil { - log.Printf("扫描异常: %v", err) - - // 即使出错也可以获取部分结果 - results := scanner.GetResults() - if len(results) > 0 { - fmt.Printf("获得部分结果: %d 个漏洞\n", len(results)) - } +if results := scanner.Results(); len(results) > 0 { + for _, v := range results { + fmt.Printf("- [%s] %s: %s\n", v.Severity, v.FullTarget, v.PocName) + } + os.Exit(1) } ``` -### 超时和取消处理 +### 超时控制 ```go ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) defer cancel() -go func() { - scanner.RunAsync() -}() - -select { -case <-ctx.Done(): - scanner.Stop() - fmt.Println("扫描超时，已停止") -case <-scanner.ResultChan: - // 正常完成 +if err := scanner.Execute(ctx); err != nil { + if errors.Is(err, context.DeadlineExceeded) { + log.Println("扫描超时") + } } + +results := scanner.Results() // 超时后仍可获取部分结果 ``` -## 集成示例 +### 信号处理 + +```go +ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) +defer stop() + +scanner.Execute(ctx) // Ctrl+C 会停止扫描并返回 context.Canceled +``` ### Web 服务集成 ```go func scanHandler(w http.ResponseWriter, r *http.Request) { - target := r.URL.Query().Get("target") - - options := afrog.NewSDKOptions() - options.Targets = []string{target} - options.PocFile = os.Getenv("POC_PATH") - - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - http.Error(w, err.Error(), 500) - return - } - defer scanner.Close() - - scanner.Run() - results := scanner.GetResults() - - // 返回 JSON 结果 - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]interface{}{ - "vulnerabilities": len(results), - "results": results, - }) + scanner, err := sdk.New(r.Context(), + sdk.WithTargets(r.URL.Query().Get("target")), + sdk.WithPocPaths(os.Getenv("POC_PATH")), + ) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + defer scanner.Close() + + if err := scanner.Execute(r.Context()); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(scanner.Results()) } ``` -### CI/CD 集成 +## 示例程序 -```go -func main() { - options := afrog.NewSDKOptions() - options.TargetsFile = "staging-urls.txt" - options.PocFile = "/security/pocs" - options.Severity = "high,critical" - - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - os.Exit(1) - } - defer scanner.Close() - - scanner.Run() - - if scanner.HasVulnerabilities() { - fmt.Println("❌ 发现安全漏洞，阻止部署") - results := scanner.GetResults() - for _, r := range results { - fmt.Printf("- %s: %s\n", r.Target, r.PocInfo.Info.Name) - } - os.Exit(1) - } - - fmt.Println("✅ 安全检查通过") -} +`examples/` 下的示例均可直接运行，PoC 路径会自动定位到仓库内的 `pocs/afrog-pocs`，也可用 `-pocs` 覆盖： + +```bash +go run ./examples/basic_scan +go run ./examples/full_output -json +go run ./examples/async_scan +go run ./examples/progress_scan +go run ./examples/oob_scan -oob dnslogcn -oob-domain your.dnslog.cn +go run ./examples/sdk_portscan -target 127.0.0.1 +go run ./examples/vuln_scan -target https://example.com +go run ./examples/port_scan -targets 127.0.0.1 ``` -## 常见问题解答 +## 常见问题 -### Q: 如何让弱口令/默认口令 PoC 仅在命中指纹后执行？ -A: 在 PoC 的 `info` 中使用 `requires` 与 `requires-mode` 声明指纹依赖，并使用 `requires-mode: strict` 实现“先指纹后执行”。完整用法与排障请参考：[requires 指纹门控：用法教程与问题答疑](requires-gating-guide.md) +### 指定了 PoC 目录，为什么内置 PoC 没有执行？ -### Q: 如何处理大量目标的扫描？ -A: 使用流式输出和适当的并发控制： -```go -options.EnableStream = true -options.Concurrency = 50 -scanner.OnResult = func(r *result.Result) { - // 立即处理，不要积累 - processImmediately(r) -} -``` +你可能加了 `WithPocPathsOnly()`。去掉它即可与内置 PoC 合并。 -### Q: 如何确保 OOB 检测正常工作？ -A: 在扫描前检查 OOB 状态： -```go -if enabled, status := scanner.GetOOBStatus(); !enabled { - log.Printf("OOB 警告: %s", status) -} -``` +### 扫描卡住不动了？ -### Q: 如何优化扫描性能？ -A: 根据网络和目标情况调整参数： -```go -// 内网扫描 -options.Concurrency = 100 -options.RateLimit = 500 +检查是否订阅了某个流却没有消费。订阅后的流写满会阻塞扫描，这是为了不丢弃漏洞。 -// 外网扫描 -options.Concurrency = 25 -options.RateLimit = 150 -options.Timeout = 15 -``` +### 为什么 `Wait` 返回 `context.Canceled`？ -### Q: 如何处理扫描中断？ -A: 使用 context 和信号处理： -```go -c := make(chan os.Signal, 1) -signal.Notify(c, os.Interrupt) +扫描被 `Stop()` 或外部 context 取消了。此时仍可通过 `Results()` 获取已发现的结果。 -go func() { - <-c - scanner.Stop() - fmt.Println("扫描已停止") -}() -``` +### 大规模扫描内存增长过快？ -## 注意事项 +```go +sdk.WithRequestResponse(false), +sdk.WithMaxStoredResults(1000), +``` -1. **POC 路径必须指定** - SDK 不会自动下载或查找 POC -2. **完全静默运行** - 不会有控制台输出，适合程序集成 -3. **无文件生成** - 不会创建任何报告文件 -4. **资源管理** - 必须调用 `Close()` 释放资源 -5. **并发安全** - 所有方法都是并发安全的 -6. **OOB 配置** - 需要正确配置才能检测带外漏洞 +配合回调实时处理结果，不要依赖 `Results()` 累积。 ## 许可证 MIT License - ---- - -更多示例和详细文档，请参考 `examples/` 目录中的示例代码。 diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 000000000..45eaa17c0 --- /dev/null +++ b/examples/README.md @@ -0,0 +1,69 @@ +# afrog SDK Examples / SDK 示例 + +Every example is runnable as-is. The PoC directory is resolved relative to the +repository, not to your current working directory, so these commands work from +anywhere: + +每个示例都可以直接运行。PoC 目录是相对仓库定位的，而不是相对当前工作目录， +所以下面的命令在任何目录下都能执行： + +```sh +go run ./examples/basic_scan +go run ./examples/full_output -json +go run ./examples/async_scan +go run ./examples/progress_scan +go run ./examples/sdk_portscan -target 127.0.0.1 +go run ./examples/vuln_scan -target https://example.com +go run ./examples/port_scan -targets 127.0.0.1 +go run ./examples/oob_scan -oob dnslogcn -oob-domain your.dnslog.cn +``` + +Use `-h` on any example to see its flags. `-pocs` overrides the PoC source and +accepts a file, a directory or a glob pattern: + +任何示例都可以用 `-h` 查看参数。`-pocs` 用于覆盖 PoC 来源， +支持单个文件、目录和 glob 通配符： + +```sh +go run ./examples/basic_scan -pocs /path/to/poc.yaml +go run ./examples/basic_scan -pocs /path/to/pocs +go run ./examples/basic_scan -pocs '/path/to/pocs/*.yaml' +``` + +## What each example shows / 各示例演示内容 + +| Example | Shows | 演示内容 | +|---------|-------|---------| +| [basic_scan](basic_scan) | Minimal scan, PoC introspection | 最小可用程序、PoC 加载检查 | +| [full_output](full_output) | Raw request/response, JSON output, failure callback | 原始请求响应、JSON 输出、失败回调 | +| [async_scan](async_scan) | `Start`/`Wait`/`Done`, result streaming, progress | 异步执行、流式结果、进度 | +| [progress_scan](progress_scan) | Progress bar driven by `GetProgress` | 基于 `GetProgress` 的进度条 | +| [oob_scan](oob_scan) | OOB adapters, identifying OOB findings | OOB 适配器配置、OOB 漏洞识别 | +| [sdk_portscan](sdk_portscan) | Port pre-scan feeding the PoC scan | 端口预扫描并作为 PoC 扫描目标 | +| [vuln_scan](vuln_scan) | CI gate with streaming and non-zero exit | CI 门禁：流式消费 + 非零退出 | +| [port_scan](port_scan) | The `portscan` package standalone | 独立使用 `portscan` 包 | + +## Targets / 扫描目标 + +Examples default to `https://scanme.sh`, a host that permits scanning. **Only +scan systems you are authorised to test.** + +示例默认目标是 `https://scanme.sh`，这是一个允许被扫描的测试主机。 +**请只扫描你有授权测试的系统。** + +## Local test lab / 本地靶场 + +[vulnweb](vulnweb) contains static pages with matching PoCs in +`vulnweb/pocs`, useful for exercising the scanner without touching any external +host: + +[vulnweb](vulnweb) 是一套静态靶场页面，配套的 PoC 在 `vulnweb/pocs`， +可以在不接触任何外部主机的情况下验证扫描器： + +```sh +# Serve the lab / 启动靶场 +cd examples/vulnweb && python3 -m http.server 8080 + +# Scan it / 扫描靶场 +go run ./examples/basic_scan -target http://127.0.0.1:8080 -pocs ./examples/vulnweb/pocs +``` diff --git a/examples/async_scan/main.go b/examples/async_scan/main.go index 10a8198c1..3232cff36 100644 --- a/examples/async_scan/main.go +++ b/examples/async_scan/main.go @@ -1,292 +1,159 @@ +// Async Scan Example / 异步扫描示例 +// +// Demonstrates asynchronous scanning with real-time result streaming. +// +// The scan runs in the background while the main goroutine consumes results as +// they are discovered. Completion is detected with Wait and Done rather than by +// draining the result stream, so each stream has exactly one consumer. +// +// 演示异步扫描与实时结果流。扫描在后台运行，主协程实时消费结果。 +// 扫描结束通过 Wait / Done 判断，而不是靠空转结果通道， +// 因此每个通道只有一个消费者，不会出现结果被瓜分的问题。 +// +// Run / 运行: +// +// go run ./examples/async_scan package main import ( "context" + "flag" "fmt" "log" - "path/filepath" + "os" + "os/signal" "sync" "time" - "github.com/zan8in/afrog/v3" - "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" ) -// Async Scan Example / 异步扫描示例 -// -// This example demonstrates asynchronous scanning with real-time result streaming. -// It shows how to receive scan results as they are discovered and handle them -// concurrently while the scan is still running. -// -// 此示例演示异步扫描和实时结果流。 -// 它展示如何在扫描仍在运行时接收发现的扫描结果， -// 并同时处理它们。 +// stats aggregates findings across goroutines. +// stats 跨协程汇总统计。 +type stats struct { + mu sync.Mutex + total int + bySev map[string]int + byTarget map[string]int +} -func main() { - // Create SDK scan options / 创建 SDK 扫描选项 - options := afrog.NewSDKOptions() +func newStats() *stats { + return &stats{bySev: map[string]int{}, byTarget: map[string]int{}} +} - // Set multiple scan targets for better async demonstration - // 设置多个扫描目标以更好地演示异步功能 - options.Targets = []string{ - "https://www.example.com", - } +func (s *stats) add(r sdk.Result) int { + s.mu.Lock() + defer s.mu.Unlock() + s.total++ + s.bySev[r.Severity]++ + s.byTarget[r.Target]++ + return s.total +} - // Set POC path (required) / 设置 POC 路径（必需） - pocPath, err := filepath.Abs("../pocs/afrog-pocs") - if err != nil { - log.Fatalf("Failed to get POC path / 获取 POC 路径失败: %v", err) +func (s *stats) snapshot() (int, map[string]int, map[string]int) { + s.mu.Lock() + defer s.mu.Unlock() + sev := make(map[string]int, len(s.bySev)) + for k, v := range s.bySev { + sev[k] = v } - options.PocFile = pocPath - - // Configuration for async scanning / 异步扫描配置 - options.Concurrency = 8 // Higher concurrency for async / 异步使用更高并发 - options.RateLimit = 30 // Moderate rate limit / 适中的速率限制 - options.Timeout = 12 // Reasonable timeout / 合理的超时时间 - options.Search = "react" // Search fingerprint POCs / 搜索指纹识别 POC - // options.Severity = "info,low,medium" // Multiple severity levels / 多个严重级别 - options.EnableStream = true // Enable streaming for async results / 启用流式输出获取异步结果 - - fmt.Println("Creating SDK scanner for async scanning... / 创建异步扫描的 SDK 扫描器...") + tgt := make(map[string]int, len(s.byTarget)) + for k, v := range s.byTarget { + tgt[k] = v + } + return s.total, sev, tgt +} - // Create scanner instance / 创建扫描器实例 - scanner, err := afrog.NewSDKScanner(options) +func main() { + target := flag.String("target", "https://scanme.sh", "target to scan") + pocs := examplepath.PocsFlag() + flag.Parse() + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + agg := newStats() + + scanner, err := sdk.New(ctx, + sdk.WithTargets(*target), + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithConcurrency(8), + sdk.WithRateLimit(30), + sdk.WithTimeout(12), + ) if err != nil { - log.Fatalf("Failed to create scanner / 创建扫描器失败: %v", err) + log.Fatalf("create scanner / 创建扫描器失败: %v", err) } - defer scanner.Close() // Always close the scanner / 始终关闭扫描器 + defer scanner.Close() - // Context for controlling goroutines / 用于控制协程的上下文 - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() + // Subscribe before starting so that no early result is missed. A stream + // publishes nothing until it is subscribed to, so an unused one is free. + // 在启动前订阅，避免漏掉早期结果。未被订阅的流不会产生任何开销。 + resultCh := scanner.ResultStream() - // Channels for communication / 用于通信的通道 - resultChan := make(chan *result.Result, 100) - doneChan := make(chan bool, 1) + if err := scanner.Start(ctx); err != nil { + log.Fatalf("start scan / 启动扫描失败: %v", err) + } - // WaitGroup for synchronization / 用于同步的等待组 var wg sync.WaitGroup - // Statistics tracking / 统计跟踪 - var stats struct { - sync.Mutex - totalVulns int - severityCount map[string]int - targetVulns map[string]int - startTime time.Time - lastVulnTime time.Time - } - stats.severityCount = make(map[string]int) - stats.targetVulns = make(map[string]int) - stats.startTime = time.Now() - - // Goroutine 1: Real-time result processing / 协程1：实时结果处理 + // Consumer: the sole reader of the result stream. + // 消费者：结果流的唯一读取方。 wg.Add(1) go func() { defer wg.Done() - - fmt.Println("Starting real-time result processor... / 启动实时结果处理器...") - - for { - select { - case result := <-scanner.ResultChan: - if result == nil { - fmt.Println("Result channel closed / 结果通道关闭") - return - } - - // Process result immediately / 立即处理结果 - processResult(result, &stats) - - // Forward to result channel for other processors / 转发到结果通道供其他处理器使用 - select { - case resultChan <- result: - default: - // Channel full, skip / 通道满了，跳过 - } - - case <-ctx.Done(): - return - } + for r := range resultCh { + n := agg.add(r) + fmt.Printf("\n[%d] %s %s [%s]\n", n, r.Target, r.PocName, r.Severity) } }() - // Goroutine 2: Progress monitoring / 协程2：进度监控 + // Progress reporter, terminated by Done. + // 进度输出，由 Done 通道终止。 wg.Add(1) go func() { defer wg.Done() - - ticker := time.NewTicker(1 * time.Second) + ticker := time.NewTicker(time.Second) defer ticker.Stop() - - fmt.Println("Starting progress monitor... / 启动进度监控器...") - for { select { case <-ticker.C: - progress := scanner.GetProgress() - scanStats := scanner.GetStats() - - stats.Lock() - elapsed := time.Since(stats.startTime) - avgSpeed := float64(scanStats.CompletedScans) / elapsed.Seconds() - stats.Unlock() - - // Create dynamic progress display / 创建动态进度显示 - fmt.Printf("\r[Progress / 进度] %.1f%% | Completed / 完成: %d/%d | Speed / 速度: %.1f/s | Vulns / 漏洞: %d", - progress, - scanStats.CompletedScans, - scanStats.TotalScans, - avgSpeed, - scanStats.FoundVulns) - - case <-ctx.Done(): - return - case <-doneChan: + st := scanner.Stats() + fmt.Printf("\r[progress / 进度] %.1f%% | %d/%d | vulns %d", + scanner.Progress(), st.CompletedScans, st.TotalScans, st.FoundVulns) + case <-scanner.Done(): return } } }() - // Goroutine 3: Result analyzer / 协程3：结果分析器 - wg.Add(1) - go func() { - defer wg.Done() - - fmt.Println("Starting result analyzer... / 启动结果分析器...") - - for { - select { - case result := <-resultChan: - if result == nil { - return - } - - // Perform detailed analysis / 执行详细分析 - analyzeResult(result) - - case <-ctx.Done(): - return - } - } - }() - - fmt.Println("Starting async scan... / 开始异步扫描...") - - // Start async scan / 开始异步扫描 - err = scanner.RunAsync() - if err != nil { - log.Printf("Failed to start async scan / 启动异步扫描失败: %v", err) - cancel() - return + // Wait returns the real scan error instead of always nil. + // Wait 返回真实的扫描错误，而不是恒为 nil。 + if err := scanner.Wait(ctx); err != nil { + log.Printf("\nscan finished with error / 扫描出错: %v", err) } - - // Simulate some other work while scanning / 在扫描时模拟其他工作 - go func() { - for i := 0; i < 10; i++ { - time.Sleep(2 * time.Second) - fmt.Printf("\n[Background Task / 后台任务] Processing other work... Step %d/10\n", i+1) - } - }() - - // Wait for scan completion by monitoring the result channel / 通过监控结果通道等待扫描完成 - go func() { - // Wait for result channel to close (scan finished) - // 等待结果通道关闭（扫描完成） - for range scanner.ResultChan { - // Channel is still open, scan is running - // 通道仍然开放，扫描正在运行 - } - doneChan <- true - }() - - // Wait for scan completion / 等待扫描完成 - <-doneChan - fmt.Printf("\n\nScan completed! Cleaning up... / 扫描完成！正在清理...\n") - - // Stop all goroutines / 停止所有协程 - cancel() - close(resultChan) - - // Wait for all goroutines to finish / 等待所有协程完成 wg.Wait() - // Get final results / 获取最终结果 - results := scanner.GetResults() - finalStats := scanner.GetStats() - - // Print comprehensive results / 打印综合结果 - fmt.Printf("\n========== Async Scan Results / 异步扫描结果 ==========\n") - fmt.Printf("Total vulnerabilities found / 发现漏洞总数: %d\n", len(results)) - fmt.Printf("Total scans completed / 完成扫描总数: %d\n", finalStats.CompletedScans) - fmt.Printf("Scan duration / 扫描耗时: %v\n", finalStats.EndTime.Sub(finalStats.StartTime)) + total, bySev, byTarget := agg.snapshot() + final := scanner.Stats() - stats.Lock() - fmt.Printf("Average scan speed / 平均扫描速度: %.2f scans/sec\n", - float64(finalStats.CompletedScans)/finalStats.EndTime.Sub(finalStats.StartTime).Seconds()) + fmt.Printf("\n\n========== Async Scan Results / 异步扫描结果 ==========\n") + fmt.Printf("vulnerabilities / 漏洞总数: %d\n", total) + fmt.Printf("completed scans / 完成任务: %d\n", final.CompletedScans) + fmt.Printf("duration / 耗时: %v\n", final.Duration()) - if len(stats.severityCount) > 0 { - fmt.Println("\nVulnerability distribution by severity / 按严重程度分布的漏洞:") - for severity, count := range stats.severityCount { - fmt.Printf(" %s: %d\n", severity, count) + if len(bySev) > 0 { + fmt.Println("\nby severity / 按严重程度:") + for sev, n := range bySev { + fmt.Printf(" %s: %d\n", sev, n) } } - - if len(stats.targetVulns) > 0 { - fmt.Println("\nVulnerability distribution by target / 按目标分布的漏洞:") - for target, count := range stats.targetVulns { - fmt.Printf(" %s: %d\n", target, count) + if len(byTarget) > 0 { + fmt.Println("\nby target / 按目标:") + for t, n := range byTarget { + fmt.Printf(" %s: %d\n", t, n) } } - stats.Unlock() - - fmt.Println("\n========== Async Scanning Benefits / 异步扫描的优势 ==========") - fmt.Println("✓ Real-time result processing / 实时结果处理") - fmt.Println("✓ Concurrent analysis while scanning / 扫描时并发分析") - fmt.Println("✓ Non-blocking operation / 非阻塞操作") - fmt.Println("✓ Better resource utilization / 更好的资源利用") - fmt.Println("✓ Immediate response to findings / 对发现的立即响应") - - fmt.Println("\nAsync scan completed successfully! / 异步扫描成功完成!") -} - -// processResult handles each result as it arrives / 处理每个到达的结果 -func processResult(result *result.Result, stats *struct { - sync.Mutex - totalVulns int - severityCount map[string]int - targetVulns map[string]int - startTime time.Time - lastVulnTime time.Time -}) { - stats.Lock() - defer stats.Unlock() - - stats.totalVulns++ - stats.lastVulnTime = time.Now() - stats.severityCount[result.PocInfo.Info.Severity]++ - stats.targetVulns[result.Target]++ - - // Real-time notification / 实时通知 - fmt.Printf("\n🚨 [LIVE] Vulnerability #%d found / 发现漏洞 #%d:\n", stats.totalVulns, stats.totalVulns) - fmt.Printf(" Target / 目标: %s\n", result.Target) - fmt.Printf(" POC / POC: %s\n", result.PocInfo.Info.Name) - fmt.Printf(" Severity / 严重程度: %s\n", result.PocInfo.Info.Severity) - fmt.Printf(" Time / 时间: %s\n", stats.lastVulnTime.Format("15:04:05")) -} - -// analyzeResult performs detailed analysis on each result / 对每个结果执行详细分析 -func analyzeResult(result *result.Result) { - // Simulate some analysis work / 模拟一些分析工作 - time.Sleep(100 * time.Millisecond) - - // Example: Check for specific vulnerability patterns / 示例：检查特定的漏洞模式 - if result.PocInfo.Info.Severity == "high" || result.PocInfo.Info.Severity == "critical" { - fmt.Printf("\n⚠️ [ALERT] High-priority vulnerability requires immediate attention! / 高优先级漏洞需要立即关注!\n") - fmt.Printf(" Target / 目标: %s\n", result.Target) - fmt.Printf(" POC / POC: %s\n", result.PocInfo.Info.Name) - - // Here you could trigger alerts, send notifications, etc. - // 这里您可以触发警报、发送通知等 - } } diff --git a/examples/basic_scan/main.go b/examples/basic_scan/main.go index 1d82ae1d7..5d200e919 100644 --- a/examples/basic_scan/main.go +++ b/examples/basic_scan/main.go @@ -1,82 +1,82 @@ +// Basic Scan Example / 基础扫描示例 +// +// Demonstrates the smallest useful afrog SDK program: configure a target, +// point at a PoC directory, run the scan and read the results. +// +// 演示 afrog SDK 最小可用程序：配置目标、指定 PoC 目录、执行扫描并读取结果。 +// +// Run / 运行: +// +// go run ./examples/basic_scan +// go run ./examples/basic_scan -target https://example.com -pocs /path/to/pocs package main import ( + "context" + "flag" "fmt" "log" - "path/filepath" + "os" + "os/signal" - "github.com/zan8in/afrog/v3" + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" ) -// Basic Scan Example / 基础扫描示例 -// -// This example demonstrates the most basic usage of the Afrog SDK. -// It performs a simple vulnerability scan on a target URL. -// -// 此示例演示了 Afrog SDK 的最基本用法。 -// 它对目标 URL 执行简单的漏洞扫描。 - func main() { - // Create SDK scan options / 创建 SDK 扫描选项 - options := afrog.NewSDKOptions() + target := flag.String("target", "https://scanme.sh", "target to scan") + pocs := examplepath.PocsFlag() + severity := flag.String("severity", "info", "severity filter, e.g. \"high,critical\"") + flag.Parse() - // Set scan targets / 设置扫描目标 - options.Targets = []string{"https://www.example.com"} + // Ctrl+C cancels the context, which stops the scan cleanly. + // Ctrl+C 取消 context，扫描会干净地停止。 + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() - // Set POC path (required) / 设置 POC 路径（必需） - pocPath, err := filepath.Abs("../pocs/afrog-pocs") + // WithPocPaths accepts a single file, a directory, or a glob pattern. + // WithPocPaths 支持单个文件、目录，以及 glob 通配符。 + scanner, err := sdk.New(ctx, + sdk.WithTargets(*target), + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithSeverity(*severity), + sdk.WithConcurrency(10), + sdk.WithRateLimit(50), + sdk.WithTimeout(10), + ) if err != nil { - log.Fatalf("Failed to get POC path / 获取 POC 路径失败: %v", err) + log.Fatalf("create scanner / 创建扫描器失败: %v", err) } - options.PocFile = pocPath - - // Basic configuration / 基础配置 - options.Concurrency = 10 // Concurrent threads / 并发线程数 - options.RateLimit = 50 // Request rate limit / 请求速率限制 - options.Timeout = 10 // Timeout in seconds / 超时时间（秒） - options.Search = "info" // Search for info-level POCs / 搜索信息级别的 POC - options.Severity = "info" // Only scan info severity / 只扫描信息严重级别 + defer scanner.Close() - fmt.Println("Creating SDK scanner... / 创建 SDK 扫描器...") - - // Create scanner instance / 创建扫描器实例 - scanner, err := afrog.NewSDKScanner(options) - if err != nil { - log.Fatalf("Failed to create scanner / 创建扫描器失败: %v", err) + // Inspect what was loaded before spending time on the network. + // 在真正发起网络请求之前，先确认加载到了哪些 PoC。 + fmt.Printf("loaded %d pocs / 已加载 %d 个 PoC\n", scanner.PocCount(), scanner.PocCount()) + for _, d := range scanner.PocDiagnostics() { + log.Printf("skipped poc / 跳过 PoC: %v", d) } - defer scanner.Close() // Always close the scanner / 始终关闭扫描器 - fmt.Println("Starting scan... / 开始扫描...") - - // Execute scan (synchronous) / 执行扫描（同步） - err = scanner.Run() - if err != nil { - log.Printf("Scan error occurred / 扫描出现错误: %v", err) + fmt.Println("scanning... / 扫描中...") + if err := scanner.Execute(ctx); err != nil { + log.Printf("scan finished with error / 扫描出错: %v", err) } - // Get scan results / 获取扫描结果 - results := scanner.GetResults() - stats := scanner.GetStats() + results := scanner.Results() + stats := scanner.Stats() - // Print results / 打印结果 - fmt.Printf("\n========== Scan Results / 扫描结果 ==========\n") - fmt.Printf("Vulnerabilities found / 发现漏洞: %d\n", len(results)) - fmt.Printf("Scan progress / 扫描进度: %.1f%%\n", scanner.GetProgress()) - fmt.Printf("Scan duration / 扫描耗时: %v\n", stats.EndTime.Sub(stats.StartTime)) + fmt.Printf("\n========== Results / 扫描结果 ==========\n") + fmt.Printf("vulnerabilities / 漏洞数: %d\n", len(results)) + fmt.Printf("duration / 耗时: %v\n", stats.Duration()) - // Display vulnerability details / 显示漏洞详情 - if len(results) > 0 { - fmt.Printf("\n========== Vulnerability Details / 漏洞详情 ==========\n") - for i, result := range results { - fmt.Printf("%d. Target / 目标: %s\n", i+1, result.Target) - fmt.Printf(" POC Name / POC 名称: %s\n", result.PocInfo.Info.Name) - fmt.Printf(" Severity / 严重程度: %s\n", result.PocInfo.Info.Severity) - fmt.Printf(" Description / 描述: %s\n", result.PocInfo.Info.Description) - fmt.Println(" ---") + for i, v := range results { + fmt.Printf("%d. [%s] %s\n", i+1, v.Severity, v.FullTarget) + fmt.Printf(" poc: %s (%s)\n", v.PocName, v.PocID) + if v.Description != "" { + fmt.Printf(" description / 描述: %s\n", v.Description) } - } else { - fmt.Println("No vulnerabilities found / 未发现漏洞") } - - fmt.Println("Scan completed! / 扫描完成!") + if len(results) == 0 { + fmt.Println("no vulnerabilities found / 未发现漏洞") + } } diff --git a/examples/full_output/main.go b/examples/full_output/main.go new file mode 100644 index 000000000..62588da8f --- /dev/null +++ b/examples/full_output/main.go @@ -0,0 +1,127 @@ +// Full Output Example / 完整数据输出示例 +// +// Demonstrates how to obtain the complete request and response of every scan +// step, and how to serialise results to JSON. +// +// 演示如何获取每一步扫描的完整请求/响应报文，以及如何把结果序列化为 JSON。 +// +// Run / 运行: +// +// go run ./examples/full_output +// go run ./examples/full_output -target https://example.com -json +package main + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "log" + "os" + "os/signal" + "strings" + + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" +) + +func main() { + target := flag.String("target", "https://scanme.sh", "target to scan") + pocs := examplepath.PocsFlag() + asJSON := flag.Bool("json", false, "print results as JSON") + flag.Parse() + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + scanner, err := sdk.New(ctx, + sdk.WithTargets(*target), + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithConcurrency(10), + sdk.WithTimeout(10), + + // Enabled by default; shown here to make the knob discoverable. + // Turn it off on large scans to keep memory bounded. + // 默认即为开启，这里显式写出以便发现该开关。 + // 大规模扫描时可关闭以控制内存占用。 + sdk.WithRequestResponse(true), + + // Failures used to be invisible; now they can be observed. + // 执行失败以前完全不可见，现在可以被观察到。 + sdk.WithFailureHandler(func(f sdk.Failure) { + log.Printf("poc %s failed on %s: %v", f.PocID, f.Target, f.Err) + }), + ) + if err != nil { + log.Fatalf("create scanner / 创建扫描器失败: %v", err) + } + defer scanner.Close() + + if err := scanner.Execute(ctx); err != nil { + log.Printf("scan finished with error / 扫描出错: %v", err) + } + + results := scanner.Results() + if len(results) == 0 { + fmt.Println("no vulnerabilities found / 未发现漏洞") + return + } + + // sdk.Result is JSON serialisable: raw request/response are plain strings, + // not protobuf []byte fields that would be base64 encoded. + // sdk.Result 可直接序列化：原始请求/响应是普通字符串， + // 而不是会被 base64 编码的 protobuf []byte 字段。 + if *asJSON { + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + if err := enc.Encode(results); err != nil { + log.Fatalf("encode json: %v", err) + } + return + } + + for i, v := range results { + fmt.Printf("\n========== [%d] %s ==========\n", i+1, v.PocID) + fmt.Printf("name / 名称: %s\n", v.PocName) + fmt.Printf("severity / 等级: %s\n", v.Severity) + fmt.Printf("target / 目标: %s\n", v.FullTarget) + if v.CveID != "" { + fmt.Printf("cve: %s (cvss %.1f)\n", v.CveID, v.CvssScore) + } + for k, val := range v.Extractors { + fmt.Printf("extractor: %s = %s\n", k, val) + } + + for j, ex := range v.Exchanges { + fmt.Printf("\n--- step %d/%d (matched=%v) ---\n", j+1, len(v.Exchanges), ex.Matched) + fmt.Printf("%s %s -> %d (%d ms)\n", ex.Method, ex.URL, ex.StatusCode, ex.LatencyMs) + + fmt.Println("\n>>> REQUEST / 请求") + fmt.Println(indent(ex.Request)) + + fmt.Println("<<< RESPONSE / 响应") + fmt.Println(indent(ex.Response)) + + if ex.BodyTruncated { + fmt.Println("!! response body was truncated at MaxRespBodySize") + fmt.Println("!! 响应体在 MaxRespBodySize 上限处被截断，不是完整响应") + } + if ex.BruteTruncated { + fmt.Printf("!! brute force stopped at %d requests / 爆破在 %d 个请求处截断\n", + ex.BruteRequests, ex.BruteRequests) + } + } + } +} + +func indent(s string) string { + if strings.TrimSpace(s) == "" { + return " (empty)" + } + lines := strings.Split(strings.TrimRight(s, "\n"), "\n") + for i, l := range lines { + lines[i] = " " + l + } + return strings.Join(lines, "\n") +} diff --git a/examples/internal/examplepath/examplepath.go b/examples/internal/examplepath/examplepath.go new file mode 100644 index 000000000..5c6d16c7e --- /dev/null +++ b/examples/internal/examplepath/examplepath.go @@ -0,0 +1,41 @@ +// Package examplepath resolves paths inside the afrog repository for the +// runnable examples. +// +// The examples used to hardcode relative paths such as "../pocs/afrog-pocs", +// which resolve against the caller's working directory and therefore broke +// depending on where `go run` was invoked from. Resolving against this source +// file's location instead makes the examples work from any directory. +package examplepath + +import ( + "flag" + "path/filepath" + "runtime" +) + +// RepoRoot returns the absolute path of the afrog repository root. +func RepoRoot() string { + _, thisFile, _, ok := runtime.Caller(0) + if !ok { + return "." + } + // /examples/internal/examplepath/examplepath.go -> + return filepath.Clean(filepath.Join(filepath.Dir(thisFile), "..", "..", "..")) +} + +// DefaultPocs returns the absolute path of the bundled PoC directory. +func DefaultPocs() string { + return filepath.Join(RepoRoot(), "pocs", "afrog-pocs") +} + +// VulnwebPocs returns the absolute path of the local test-lab PoC directory +// that pairs with examples/vulnweb. +func VulnwebPocs() string { + return filepath.Join(RepoRoot(), "examples", "vulnweb", "pocs") +} + +// PocsFlag registers a -pocs flag defaulting to the bundled PoC directory and +// returns a pointer to its value. Call flag.Parse before dereferencing it. +func PocsFlag() *string { + return flag.String("pocs", DefaultPocs(), "PoC file, directory or glob pattern") +} diff --git a/examples/oob_scan/main.go b/examples/oob_scan/main.go index 2b0b673ab..1eaba2ff5 100644 --- a/examples/oob_scan/main.go +++ b/examples/oob_scan/main.go @@ -1,249 +1,137 @@ +// OOB (Out-of-Band) Scan Example / OOB（带外）扫描示例 +// +// Demonstrates configuring out-of-band detection and telling OOB findings +// apart from ordinary ones. +// +// 演示如何配置带外检测，以及如何区分 OOB 漏洞与普通漏洞。 +// +// afrog v3 expresses OOB with the {{oob.DNS}} / {{oob.HTTP}} placeholders and +// the oobCheck() expression, for example: +// +// rules: +// r0: +// request: +// method: GET +// path: /?dns=ping%20{{oob.DNS}} +// expression: oobCheck(oob.ProtocolDNS, 5) +// +// afrog v3 使用 {{oob.DNS}} / {{oob.HTTP}} 占位符配合 oobCheck() 表达式描述 OOB。 +// v2 时代的 `set: oob/reverse` 与 oobWait() 写法均已废弃。 +// +// Run / 运行: +// +// go run ./examples/oob_scan -oob dnslogcn -oob-domain your.dnslog.cn +// go run ./examples/oob_scan -oob ceyeio -oob-key TOKEN -oob-domain xxx.ceye.io package main import ( + "context" + "flag" "fmt" "log" - "path/filepath" + "os" + "os/signal" "strings" + "sync/atomic" - "github.com/zan8in/afrog/v3" - "github.com/zan8in/afrog/v3/pkg/poc" - "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" ) -// OOB (Out-of-Band) Scan Example / OOB（带外）扫描示例 -// -// This example demonstrates how to configure and use OOB detection -// with different OOB adapters like ceyeio, dnslogcn, alphalog, etc. -// -// 此示例演示如何配置和使用 OOB 检测， -// 支持不同的 OOB 适配器，如 ceyeio、dnslogcn、alphalog 等。 - func main() { - // Create SDK scan options / 创建 SDK 扫描选项 - options := afrog.NewSDKOptions() - - // Set scan targets / 设置扫描目标 - options.Targets = []string{ - "https://www.example.com", - } - - // Set POC path (required) / 设置 POC 路径（必需） - pocPath, err := filepath.Abs("../pocs/afrog-pocs") - if err != nil { - log.Fatalf("Failed to get POC path / 获取 POC 路径失败: %v", err) - } - options.PocFile = pocPath - - // Basic configuration / 基础配置 - options.Concurrency = 10 - options.RateLimit = 50 - options.Timeout = 15 - options.Severity = "info,low,medium,high,critical" // All severity levels / 所有严重级别 - - // ========== OOB Configuration / OOB 配置 ========== - // Enable OOB detection / 启用 OOB 检测 - options.EnableOOB = true - - // Method 1: Configure CEYE.io (recommended) / 方法1：配置 CEYE.io（推荐） - // Register at http://ceye.io/ to get your token and domain - // 在 http://ceye.io/ 注册以获取您的令牌和域名 - options.OOB = "ceyeio" - options.OOBKey = "your-ceye-api-token" // Replace with your CEYE API token / 替换为您的 CEYE API 令牌 - options.OOBDomain = "your-subdomain.ceye.io" // Replace with your CEYE domain / 替换为您的 CEYE 域名 - - // Method 2: Configure DNSLog.cn (free, no registration required) - // 方法2：配置 DNSLog.cn（免费，无需注册） - // Uncomment the following lines to use DNSLog.cn instead: - // 取消注释以下行以使用 DNSLog.cn： - // options.OOB = "dnslogcn" - // options.OOBDomain = "your.dnslog.cn" // Get from http://dnslog.cn/ - - // Method 3: Configure Alphalog - // 方法3：配置 Alphalog - // options.OOB = "alphalog" - // options.OOBDomain = "your.alphalog.cn" - // options.OOBApiUrl = "https://api.alphalog.cn" - - // Method 4: Configure XRay - // 方法4：配置 XRay - // options.OOB = "xray" - // options.OOBDomain = "your.xray.domain" - // options.OOBApiUrl = "http://xray-api:8777" - // options.OOBKey = "your-xray-token" - - // Method 5: Configure RevSuit - // 方法5：配置 RevSuit - // options.OOB = "revsuit" - // options.OOBKey = "your-revsuit-key" - // options.OOBDomain = "your.revsuit.domain" - // options.OOBHttpUrl = "http://your.revsuit.domain" - // options.OOBApiUrl = "http://your.revsuit.domain:8080" - - fmt.Println("Creating SDK scanner with OOB configuration... / 创建带 OOB 配置的 SDK 扫描器...") - - // Create scanner instance / 创建扫描器实例 - scanner, err := afrog.NewSDKScanner(options) + target := flag.String("target", "https://scanme.sh", "target to scan") + pocs := examplepath.PocsFlag() + adapter := flag.String("oob", "dnslogcn", "oob adapter: ceyeio|dnslogcn|alphalog|xray|revsuit") + key := flag.String("oob-key", "", "oob api key / token") + domain := flag.String("oob-domain", "", "oob domain") + apiURL := flag.String("oob-api-url", "", "oob api url") + httpURL := flag.String("oob-http-url", "", "oob http url") + flag.Parse() + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + + // Handlers run on scan workers, so the counters must be atomic. + // 回调运行在扫描工作协程上，计数器必须使用原子操作。 + var oobVulns, normalVulns atomic.Int64 + + scanner, err := sdk.New(ctx, + sdk.WithTargets(*target), + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithConcurrency(10), + sdk.WithTimeout(15), + sdk.WithOOB(sdk.OOBOptions{ + Adapter: *adapter, + Key: *key, + Domain: *domain, + ApiURL: *apiURL, + HttpURL: *httpURL, + }), + sdk.WithResultHandler(func(r sdk.Result) { + if isOOBFinding(r) { + oobVulns.Add(1) + fmt.Printf("\n[OOB vulnerability / OOB 漏洞]\n") + } else { + normalVulns.Add(1) + fmt.Printf("\n[standard vulnerability / 标准漏洞]\n") + } + fmt.Printf(" target / 目标: %s\n", r.FullTarget) + fmt.Printf(" poc: %s (%s)\n", r.PocName, r.Severity) + }), + ) if err != nil { - log.Fatalf("Failed to create scanner / 创建扫描器失败: %v", err) + log.Fatalf("create scanner / 创建扫描器失败: %v", err) } - defer scanner.Close() // Always close the scanner / 始终关闭扫描器 + defer scanner.Close() - // Check OOB status before scanning / 扫描前检查 OOB 状态 - if oobEnabled, oobStatus := scanner.GetOOBStatus(); oobEnabled { - fmt.Printf("✓ OOB Status / OOB 状态: %s\n", oobStatus) + // OOBStatus performs a live connectivity probe against the OOB service. + // OOBStatus 会对 OOB 服务发起一次真实的连通性探测。 + if enabled, status := scanner.OOBStatus(); enabled { + fmt.Printf("OOB status / OOB 状态: %s\n", status) } else { - fmt.Printf("✗ OOB Status / OOB 状态: %s\n", oobStatus) - fmt.Println("Warning: OOB is not properly configured. Some POCs may not work correctly.") - fmt.Println("警告：OOB 未正确配置。某些 POC 可能无法正常工作。") - - // You can choose to continue without OOB or exit - // 您可以选择在没有 OOB 的情况下继续或退出 - // return + fmt.Printf("OOB status / OOB 状态: %s\n", status) + fmt.Println("OOB PoCs will not produce findings / OOB 类 PoC 将无法产出结果") } - // Set up real-time result callback / 设置实时结果回调 - var oobVulnCount, normalVulnCount int - scanner.OnResult = func(r *result.Result) { - // Check if this is an OOB-related vulnerability / 检查是否为 OOB 相关漏洞 - isOOBVuln := r.PocInfo != nil && pocUsesOOB(r.PocInfo) - - if isOOBVuln { - oobVulnCount++ - fmt.Printf("\n[OOB Vulnerability Found / 发现 OOB 漏洞] 🚨\n") - } else { - normalVulnCount++ - fmt.Printf("\n[Standard Vulnerability Found / 发现标准漏洞] ⚠️\n") - } - - fmt.Printf(" Target / 目标: %s\n", r.Target) - fmt.Printf(" POC Name / POC 名称: %s\n", r.PocInfo.Info.Name) - fmt.Printf(" Severity / 严重程度: %s\n", r.PocInfo.Info.Severity) - fmt.Printf(" Author / 作者: %s\n", r.PocInfo.Info.Author) - if r.PocInfo.Info.Description != "" { - fmt.Printf(" Description / 描述: %s\n", r.PocInfo.Info.Description) - } - fmt.Println(" " + strings.Repeat("-", 50)) + if err := scanner.Execute(ctx); err != nil { + log.Printf("scan finished with error / 扫描出错: %v", err) } - fmt.Println("Starting OOB-enabled scan... / 开始启用 OOB 的扫描...") - - // Execute scan (synchronous) / 执行扫描（同步） - err = scanner.Run() - if err != nil { - log.Printf("Scan error occurred / 扫描出现错误: %v", err) - } - - // Get scan results / 获取扫描结果 - results := scanner.GetResults() - stats := scanner.GetStats() - - // Print comprehensive results / 打印综合结果 + results := scanner.Results() fmt.Printf("\n========== OOB Scan Results / OOB 扫描结果 ==========\n") - fmt.Printf("Total vulnerabilities found / 发现漏洞总数: %d\n", len(results)) - fmt.Printf(" - OOB vulnerabilities / OOB 漏洞: %d\n", oobVulnCount) - fmt.Printf(" - Standard vulnerabilities / 标准漏洞: %d\n", normalVulnCount) - fmt.Printf("Scan progress / 扫描进度: %.1f%%\n", scanner.GetProgress()) - fmt.Printf("Scan duration / 扫描耗时: %v\n", stats.EndTime.Sub(stats.StartTime)) - - // Analyze POC types used / 分析使用的 POC 类型 - if len(results) > 0 { - fmt.Printf("\n========== Vulnerability Analysis / 漏洞分析 ==========\n") - - severityCount := make(map[string]int) - pocTypeCount := make(map[string]int) - - for _, result := range results { - severityCount[result.PocInfo.Info.Severity]++ - - // Analyze POC type / 分析 POC 类型 - isOOB := false - for _, rule := range result.PocInfo.Set { - if key, ok := rule.Key.(string); ok && (key == "oob" || key == "reverse") { - isOOB = true - break - } - } + fmt.Printf("total / 总数: %d\n", len(results)) + fmt.Printf("oob / OOB 漏洞: %d\n", oobVulns.Load()) + fmt.Printf("standard / 普通: %d\n", normalVulns.Load()) - if isOOB { - pocTypeCount["OOB"]++ - } else { - pocTypeCount["Standard"]++ - } - } - - fmt.Println("By Severity / 按严重程度:") - for severity, count := range severityCount { - fmt.Printf(" %s: %d\n", severity, count) - } - - fmt.Println("\nBy POC Type / 按 POC 类型:") - for pocType, count := range pocTypeCount { - fmt.Printf(" %s: %d\n", pocType, count) - } - } else { - fmt.Println("No vulnerabilities found / 未发现漏洞") - fmt.Println("This might be because:") - fmt.Println("这可能是因为:") - fmt.Println("1. The targets are secure / 目标是安全的") - fmt.Println("2. OOB configuration is incorrect / OOB 配置不正确") - fmt.Println("3. Network connectivity issues / 网络连接问题") + bySeverity := map[string]int{} + for _, v := range results { + bySeverity[v.Severity]++ + } + for sev, n := range bySeverity { + fmt.Printf(" %s: %d\n", sev, n) } - - fmt.Println("\n========== OOB Configuration Tips / OOB 配置提示 ==========") - fmt.Println("For best results with OOB detection:") - fmt.Println("为了获得 OOB 检测的最佳结果:") - fmt.Println("1. Use CEYE.io for most reliable results / 使用 CEYE.io 获得最可靠的结果") - fmt.Println("2. Ensure your OOB service is accessible / 确保您的 OOB 服务可访问") - fmt.Println("3. Check firewall settings / 检查防火墙设置") - fmt.Println("4. Verify API tokens and domains / 验证 API 令牌和域名") - - fmt.Println("\nOOB scan completed! / OOB 扫描完成!") } -func pocUsesOOB(p *poc.Poc) bool { - if p == nil { - return false - } - if containsOOBToken(p.Expression) { - return true - } - for _, it := range p.Set { - if s, ok := it.Value.(string); ok && containsOOBToken(s) { +// isOOBFinding reports whether a finding came from out-of-band detection. +// +// An OOB PoC embeds an {{oob.*}} placeholder in the request, so the evidence +// is visible in the raw request that the SDK returns. +// +// isOOBFinding 判断结果是否来自带外检测。OOB 类 PoC 会在请求中嵌入 +// {{oob.*}} 占位符，因此可以直接在 SDK 返回的原始请求里看到痕迹。 +func isOOBFinding(r sdk.Result) bool { + for _, ex := range r.Exchanges { + if strings.Contains(strings.ToLower(ex.Request), ".oob.") || + strings.Contains(strings.ToLower(ex.Request), "dnslog") || + strings.Contains(strings.ToLower(ex.Request), "ceye.io") { return true } } - for _, rm := range p.Rules { - r := rm.Value - if containsOOBToken(r.Expression) { - return true - } - for _, e := range r.Expressions { - if containsOOBToken(e) { - return true - } - } - req := r.Request - if containsOOBToken(req.Path) || containsOOBToken(req.Host) || containsOOBToken(req.Body) || containsOOBToken(req.Raw) || containsOOBToken(req.Data) { + for k := range r.Extractors { + if strings.HasPrefix(strings.ToLower(k), "oob") { return true } - for _, hv := range req.Headers { - if containsOOBToken(hv) { - return true - } - } } return false } - -func containsOOBToken(s string) bool { - if s == "" { - return false - } - l := strings.ToLower(s) - return strings.Contains(l, "oobwait(") || - strings.Contains(l, "{{oob") || - strings.Contains(l, "{{ oob") || - strings.Contains(l, "oob_") || - strings.Contains(l, "oob.") -} diff --git a/examples/port_scan/main.go b/examples/port_scan/main.go index 81fd90d64..a6c85eaab 100644 --- a/examples/port_scan/main.go +++ b/examples/port_scan/main.go @@ -1,48 +1,80 @@ +// Port Scan Example / 端口扫描示例 +// +// Uses the standalone portscan package directly, without the scanner SDK. +// For port pre-scanning as part of a vulnerability scan, see examples/sdk_portscan. +// +// 直接使用独立的 portscan 包，不经过扫描器 SDK。 +// 如果需要在漏洞扫描前做端口预扫描，请参考 examples/sdk_portscan。 +// +// Run / 运行: +// +// go run ./examples/port_scan -targets 127.0.0.1 +// go run ./examples/port_scan -targets 127.0.0.1 -ports 22,80,443 package main import ( "context" + "flag" "fmt" + "os" + "strings" + "sync" "time" "github.com/zan8in/afrog/v3/pkg/portscan" ) func main() { - // 1. Setup Options + targets := flag.String("targets", "127.0.0.1", "comma separated hosts / CIDRs to scan") + ports := flag.String("ports", "top", "ports: top|full|all|80,443|1-1024") + skipDiscovery := flag.Bool("Pn", true, "skip host discovery") + flag.Parse() + opts := portscan.DefaultOptions() - opts.Targets = []string{"8.152.216.157", "60.247.152.241"} + opts.Targets = splitAndTrim(*targets) + opts.Ports = *ports opts.DiscoveryMethod = "auto" - opts.Ports = "top" // Test prioritized full scan - opts.SkipDiscovery = true - // opts.RateLimit = 300 - // opts.Timeout = 1000 * time.Millisecond - // opts.Retries = 2 - opts.Debug = true - // portscan.ApplyQuickestStrategy(opts) - - // 2. Setup Callback - opts.OnResult = func(result *portscan.ScanResult) { - fmt.Printf("%s:%d\n", result.Host, result.Port) - // if result.Banner != "" { - // fmt.Printf(" Banner: %s\n", result.Banner) - // } + opts.SkipDiscovery = *skipDiscovery + + if len(opts.Targets) == 0 { + fmt.Fprintln(os.Stderr, "no targets provided / 未指定目标") + os.Exit(1) + } + + // OnResult is invoked concurrently from scan workers, so writes to shared + // state (including stdout) need synchronising. + // OnResult 由扫描工作协程并发调用，访问共享状态（包括 stdout）需要加锁。 + var mu sync.Mutex + opts.OnResult = func(r *portscan.ScanResult) { + mu.Lock() + defer mu.Unlock() + fmt.Printf("%s:%d\n", r.Host, r.Port) } - // 3. Create Scanner scanner, err := portscan.NewScanner(opts) if err != nil { - panic(err) + fmt.Fprintf(os.Stderr, "create scanner / 创建扫描器失败: %v\n", err) + os.Exit(1) } - fmt.Println("Starting Port Scan...") - startTime := time.Now() + fmt.Println("starting port scan... / 开始端口扫描...") + start := time.Now() - // 4. Run Scan - err = scanner.Scan(context.Background()) - if err != nil { - fmt.Printf("Scan failed: %v\n", err) + if err := scanner.Scan(context.Background()); err != nil { + fmt.Fprintf(os.Stderr, "scan failed / 扫描失败: %v\n", err) + os.Exit(1) } - fmt.Printf("Scan completed in %v\n", time.Since(startTime)) + fmt.Printf("completed in %v / 耗时 %v\n", time.Since(start), time.Since(start)) +} + +func splitAndTrim(s string) []string { + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + if v := strings.TrimSpace(p); v != "" { + out = append(out, v) + } + } + return out } diff --git a/examples/progress_scan/main.go b/examples/progress_scan/main.go index f6774518c..f2ea14bb3 100644 --- a/examples/progress_scan/main.go +++ b/examples/progress_scan/main.go @@ -1,177 +1,122 @@ +// Progress Scan Example / 带进度条的扫描示例 +// +// Demonstrates monitoring scan progress in real time while the scan runs in +// the background. +// +// 演示在扫描后台运行的同时实时监控进度。 +// +// Run / 运行: +// +// go run ./examples/progress_scan package main import ( + "context" + "flag" "fmt" "log" - "path/filepath" + "os" + "os/signal" + "strings" "sync" "time" - "github.com/zan8in/afrog/v3" - "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" ) -// Progress Scan Example / 带进度条的扫描示例 -// -// This example demonstrates how to monitor scan progress in real-time -// and display a progress bar during the scanning process. -// -// 此示例演示如何实时监控扫描进度， -// 并在扫描过程中显示进度条。 - func main() { - // Create SDK scan options / 创建 SDK 扫描选项 - options := afrog.NewSDKOptions() - - // Set multiple scan targets for better progress demonstration - // 设置多个扫描目标以更好地演示进度 - options.Targets = []string{ - "https://www.example.com", - } + target := flag.String("target", "https://scanme.sh", "target to scan") + pocs := examplepath.PocsFlag() + flag.Parse() - // Set POC path (required) / 设置 POC 路径（必需） - pocPath, err := filepath.Abs("../pocs/afrog-pocs") - if err != nil { - log.Fatalf("Failed to get POC path / 获取 POC 路径失败: %v", err) - } - options.PocFile = pocPath + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() - // Configuration for better progress visibility / 配置以更好地显示进度 - options.Concurrency = 5 // Lower concurrency for visible progress / 较低并发以显示进度 - options.RateLimit = 20 // Lower rate limit / 较低速率限制 - options.Timeout = 15 // Longer timeout / 更长超时时间 - options.Search = "fingerprint" // Search fingerprint POCs / 搜索指纹识别 POC - options.Severity = "info,low" // Multiple severity levels / 多个严重级别 - - fmt.Println("Creating SDK scanner... / 创建 SDK 扫描器...") + // Handlers are invoked concurrently from scan workers, so shared state + // (here: stdout) needs a lock. + // 回调由扫描工作协程并发触发，共享状态（这里是 stdout）需要加锁。 + var mu sync.Mutex - // Create scanner instance / 创建扫描器实例 - scanner, err := afrog.NewSDKScanner(options) + scanner, err := sdk.New(ctx, + sdk.WithTargets(*target), + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithConcurrency(5), + sdk.WithRateLimit(20), + sdk.WithTimeout(15), + sdk.WithResultHandler(func(r sdk.Result) { + mu.Lock() + defer mu.Unlock() + fmt.Printf("\n[found / 发现] %s - %s [%s]\n", r.Target, r.PocName, r.Severity) + }), + ) if err != nil { - log.Fatalf("Failed to create scanner / 创建扫描器失败: %v", err) + log.Fatalf("create scanner / 创建扫描器失败: %v", err) } - defer scanner.Close() // Always close the scanner / 始终关闭扫描器 + defer scanner.Close() - // Real-time result callback / 实时结果回调 - var vulnCount int - var mu sync.Mutex - scanner.OnResult = func(r *result.Result) { - mu.Lock() - vulnCount++ - fmt.Printf("\n[Real-time Discovery / 实时发现] %s - %s [%s]\n", - r.Target, - r.PocInfo.Info.Name, - r.PocInfo.Info.Severity) - mu.Unlock() + start := time.Now() + if err := scanner.Start(ctx); err != nil { + log.Fatalf("start scan / 启动扫描失败: %v", err) } - // Start progress monitoring goroutine / 启动进度监控协程 - done := make(chan bool) + // The progress goroutine terminates on Done, so no sleep-based + // synchronisation is needed. + // 进度协程由 Done 通道终止，不需要用 sleep 做同步。 + var wg sync.WaitGroup + wg.Add(1) go func() { - ticker := time.NewTicker(500 * time.Millisecond) // Update every 0.5 seconds / 每0.5秒更新 + defer wg.Done() + ticker := time.NewTicker(500 * time.Millisecond) defer ticker.Stop() - for { select { case <-ticker.C: - progress := scanner.GetProgress() - stats := scanner.GetStats() - - // Create progress bar / 创建进度条 - progressBar := createProgressBar(progress, 50) - - // Clear line and print progress / 清除行并打印进度 - fmt.Printf("\r[Progress / 进度] %s %.2f%% (%d/%d) Found / 发现: %d", - progressBar, - progress, - stats.CompletedScans, - stats.TotalScans, - stats.FoundVulns) - - case <-done: + st := scanner.Stats() + percent := scanner.Progress() + mu.Lock() + fmt.Printf("\r[progress / 进度] %s %.2f%% (%d/%d) found / 发现: %d", + progressBar(percent, 40), percent, st.CompletedScans, st.TotalScans, st.FoundVulns) + mu.Unlock() + case <-scanner.Done(): return } } }() - fmt.Println("Starting scan with progress monitoring... / 开始带进度监控的扫描...") - - // Execute scan (synchronous) / 执行扫描（同步） - start := time.Now() - err = scanner.Run() - if err != nil { - log.Printf("Scan error occurred / 扫描出现错误: %v", err) + if err := scanner.Wait(ctx); err != nil { + log.Printf("\nscan finished with error / 扫描出错: %v", err) } + wg.Wait() - // Stop progress monitoring / 停止进度监控 - done <- true - time.Sleep(100 * time.Millisecond) // Wait for goroutine to finish / 等待协程结束 - - // Get final results / 获取最终结果 - results := scanner.GetResults() - stats := scanner.GetStats() - duration := time.Since(start) + results := scanner.Results() + st := scanner.Stats() - // Print final results / 打印最终结果 fmt.Printf("\n\n========== Scan Completed / 扫描完成 ==========\n") - fmt.Printf("Total targets / 总目标数: %d\n", stats.TotalTargets) - fmt.Printf("Total POCs / 总 POC 数: %d\n", stats.TotalPocs) - fmt.Printf("Total scans / 总扫描数: %d\n", stats.TotalScans) - fmt.Printf("Completed scans / 完成扫描数: %d\n", stats.CompletedScans) - fmt.Printf("Vulnerabilities found / 发现漏洞: %d\n", len(results)) - fmt.Printf("Scan duration / 扫描耗时: %v\n", duration) - fmt.Printf("Average speed / 平均速度: %.2f scans/sec\n", - float64(stats.CompletedScans)/duration.Seconds()) - - // Display vulnerability summary / 显示漏洞摘要 - if len(results) > 0 { - fmt.Printf("\n========== Vulnerability Summary / 漏洞摘要 ==========\n") - severityCount := make(map[string]int) - - for _, result := range results { - severityCount[result.PocInfo.Info.Severity]++ - } - - for severity, count := range severityCount { - fmt.Printf(" %s: %d\n", severity, count) - } - - fmt.Printf("\n========== Vulnerability Details / 漏洞详情 ==========\n") - for i, result := range results { - fmt.Printf("%d. [%s] %s\n", i+1, result.PocInfo.Info.Severity, result.Target) - fmt.Printf(" POC: %s\n", result.PocInfo.Info.Name) - if result.PocInfo.Info.Description != "" { - fmt.Printf(" Description / 描述: %s\n", result.PocInfo.Info.Description) - } - fmt.Println(" ---") - } - } else { - fmt.Println("No vulnerabilities found / 未发现漏洞") + fmt.Printf("targets / 目标数: %d\n", st.TotalTargets) + fmt.Printf("pocs / PoC 数: %d\n", st.TotalPocs) + fmt.Printf("tasks / 任务数: %d/%d\n", st.CompletedScans, st.TotalScans) + fmt.Printf("vulnerabilities: %d\n", len(results)) + fmt.Printf("duration / 耗时: %v\n", time.Since(start)) + + bySeverity := map[string]int{} + for _, v := range results { + bySeverity[v.Severity]++ + } + for sev, n := range bySeverity { + fmt.Printf(" %s: %d\n", sev, n) } - - fmt.Println("Scan completed successfully! / 扫描成功完成!") } -// createProgressBar creates a visual progress bar / 创建可视化进度条 -func createProgressBar(progress float64, width int) string { - if progress > 100 { - progress = 100 +// progressBar renders a textual progress bar of the given width. +func progressBar(percent float64, width int) string { + if percent < 0 { + percent = 0 } - if progress < 0 { - progress = 0 + if percent > 100 { + percent = 100 } - - filled := int(progress * float64(width) / 100) - bar := "[" - - for i := 0; i < width; i++ { - if i < filled { - bar += "█" - } else { - bar += "░" - } - } - - bar += "]" - return bar + filled := int(percent * float64(width) / 100) + return "[" + strings.Repeat("█", filled) + strings.Repeat("░", width-filled) + "]" } diff --git a/examples/sdk_portscan/main.go b/examples/sdk_portscan/main.go index afc661e0d..c61ad7b3c 100644 --- a/examples/sdk_portscan/main.go +++ b/examples/sdk_portscan/main.go @@ -1,86 +1,114 @@ +// SDK Port Pre-Scan Example / SDK 端口预扫描示例 +// +// Demonstrates running a port pre-scan before the PoC scan. Discovered open +// ports are appended to the target set as host:port, so subsequent PoCs run +// against the expanded target list. +// +// 演示在 PoC 扫描前先做端口预扫描。发现的开放端口会以 host:port 形式 +// 追加进目标集合，后续 PoC 按扩展后的目标列表执行。 +// +// Run / 运行: +// +// go run ./examples/sdk_portscan -target 127.0.0.1 +// go run ./examples/sdk_portscan -target 127.0.0.1 -async package main import ( + "context" "flag" "fmt" "log" - "path/filepath" + "os" + "os/signal" "sort" + "sync" - "github.com/zan8in/afrog/v3" + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" ) func main() { - var ( - target string - targetFile string - pocsPath string - ports string - enablePS bool - async bool - search string - severity string - ) - - flag.StringVar(&target, "t", "", "target") - flag.StringVar(&targetFile, "T", "", "targets file") - flag.StringVar(&pocsPath, "pocs", "", "pocs directory") - flag.StringVar(&ports, "p", "top", "ports definition for pre-scan") - flag.BoolVar(&enablePS, "ps", true, "enable pre-scan port scanning") - flag.BoolVar(&async, "async", false, "use async scan and consume PortChan") - flag.StringVar(&search, "s", "__no_such_poc__", "poc search keyword") - flag.StringVar(&severity, "S", "", "poc severity filter") + target := flag.String("target", "127.0.0.1", "target to scan") + targetFile := flag.String("target-file", "", "file with one target per line") + pocs := examplepath.PocsFlag() + ports := flag.String("ports", "top", "ports: top|full|all|80,443|1-1024") + async := flag.Bool("async", false, "run asynchronously and consume the port stream") + search := flag.String("search", "", "poc search keyword") + severity := flag.String("severity", "", "severity filter") flag.Parse() - if pocsPath == "" { - abs, err := filepath.Abs("./pocs/afrog-pocs") - if err != nil { - log.Fatalf("pocs path: %v", err) - } - pocsPath = abs - } + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() - opts := afrog.NewSDKOptions() - opts.PocFile = pocsPath - opts.Search = search - opts.Severity = severity + // Guards stdout, which handlers write to from scan workers. + // 保护 stdout：回调会从扫描工作协程并发写入。 + var mu sync.Mutex - if targetFile != "" { - opts.TargetsFile = targetFile - } else if target != "" { - opts.Targets = []string{target} + options := []sdk.Option{ + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithPortScan(sdk.PortScanOptions{ + Ports: *ports, + TimeoutMs: 500, + }), + } + if *targetFile != "" { + options = append(options, sdk.WithTargetsFile(*targetFile)) } else { - opts.Targets = []string{"127.0.0.1"} + options = append(options, sdk.WithTargets(*target)) + } + if *search != "" { + options = append(options, sdk.WithSearch(*search)) + } + if *severity != "" { + options = append(options, sdk.WithSeverity(*severity)) + } + if !*async { + options = append(options, sdk.WithPortHandler(func(p sdk.PortEvent) { + mu.Lock() + defer mu.Unlock() + fmt.Printf("[open] %s:%d\n", p.Host, p.Port) + })) } - opts.PortScan = enablePS - opts.PSPorts = ports - opts.PSSkipDiscovery = false - opts.PSTimeout = 500 - - sc, err := afrog.NewSDKScanner(opts) + scanner, err := sdk.New(ctx, options...) if err != nil { - log.Fatalf("NewSDKScanner: %v", err) + log.Fatalf("create scanner / 创建扫描器失败: %v", err) } - defer sc.Close() + defer scanner.Close() - if async { - if err := sc.RunAsync(); err != nil { - log.Printf("RunAsync: %v", err) + if *async { + // Subscribe before Start. Once subscribed the stream must be consumed: + // sends block when the buffer fills so that no open port is dropped. + // 在 Start 之前订阅。订阅后必须消费：缓冲写满时发送会阻塞， + // 以保证开放端口不被静默丢弃。 + ports := scanner.PortStream() + + if err := scanner.Start(ctx); err != nil { + log.Fatalf("start scan / 启动扫描失败: %v", err) } - for r := range sc.PortChan { - fmt.Printf("[open] %s:%d\n", r.Host, r.Port) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for p := range ports { + fmt.Printf("[open] %s:%d\n", p.Host, p.Port) + } + }() + + if err := scanner.Wait(ctx); err != nil { + log.Printf("scan finished with error / 扫描出错: %v", err) } + wg.Wait() } else { - sc.OnPort = func(host string, port int) { - fmt.Printf("[open] %s:%d\n", host, port) - } - if err := sc.Run(); err != nil { - log.Printf("Run: %v", err) + if err := scanner.Execute(ctx); err != nil { + log.Printf("scan finished with error / 扫描出错: %v", err) } } - open := sc.GetOpenPorts() + fmt.Println("\n========== Open Ports / 开放端口 ==========") + open := scanner.OpenPorts() hosts := make([]string, 0, len(open)) for h := range open { hosts = append(hosts, h) @@ -94,4 +122,11 @@ func main() { fmt.Printf("%s:%d\n", h, p) } } + + if results := scanner.Results(); len(results) > 0 { + fmt.Println("\n========== Vulnerabilities / 漏洞 ==========") + for _, v := range results { + fmt.Printf("[%s] %s - %s\n", v.Severity, v.FullTarget, v.PocName) + } + } } diff --git a/examples/vuln_scan/main.go b/examples/vuln_scan/main.go index b0e5cd401..b9577a5fc 100644 --- a/examples/vuln_scan/main.go +++ b/examples/vuln_scan/main.go @@ -1,70 +1,96 @@ +// Vulnerability Scan Example / 漏洞扫描示例 +// +// Demonstrates streaming results while the scan runs, and exiting non-zero when +// something is found — the shape a CI security gate usually needs. +// +// 演示扫描过程中实时消费结果，并在发现漏洞时以非零状态码退出， +// 这是 CI 安全门禁常见的用法。 +// +// Run / 运行: +// +// go run ./examples/vuln_scan -target https://example.com -search CVE-2024-1234 package main import ( + "context" + "flag" "fmt" "log" - "path/filepath" + "os" + "os/signal" + "sync" "time" - "github.com/zan8in/afrog/v3" + "github.com/zan8in/afrog/v3/examples/internal/examplepath" + "github.com/zan8in/afrog/v3/pkg/sdk" ) func main() { - // Create SDK scan options - options := afrog.NewSDKOptions() + target := flag.String("target", "https://scanme.sh", "target to scan") + pocs := examplepath.PocsFlag() + search := flag.String("search", "", "poc search keyword") + severity := flag.String("severity", "", "severity filter, e.g. \"high,critical\"") + failOnFind := flag.Bool("fail-on-find", false, "exit 1 when a vulnerability is found") + flag.Parse() - // Set scan target - options.Targets = []string{ - "https://mmw.keshvacredit.com", - } + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() - // Set POC path (required) - pocPath, err := filepath.Abs("./pocs/afrog-pocs") // Adjust path as needed - if err != nil { - log.Fatalf("Failed to get POC path: %v", err) + options := []sdk.Option{ + sdk.WithTargets(*target), + sdk.WithPocPaths(*pocs), + sdk.WithPocPathsOnly(), + sdk.WithConcurrency(8), + sdk.WithRateLimit(30), + sdk.WithTimeout(12), + } + if *search != "" { + options = append(options, sdk.WithSearch(*search)) + } + if *severity != "" { + options = append(options, sdk.WithSeverity(*severity)) } - options.PocFile = pocPath - - // Configuration for scanning - options.Concurrency = 8 - options.RateLimit = 30 - options.Timeout = 12 - options.Search = "CVE-2025-55182" // Search for specific POC - options.EnableStream = true - - fmt.Println("Creating SDK scanner for vulnerability scanning...") - // Create scanner instance - scanner, err := afrog.NewSDKScanner(options) + scanner, err := sdk.New(ctx, options...) if err != nil { - log.Fatalf("Failed to create scanner: %v", err) + log.Fatalf("create scanner / 创建扫描器失败: %v", err) } defer scanner.Close() - // Start async scan - err = scanner.RunAsync() - if err != nil { - log.Printf("Failed to start async scan: %v", err) - return + // Subscribe before Start so that no finding is missed. + // 在 Start 之前订阅，避免漏掉任何结果。 + results := scanner.ResultStream() + + if err := scanner.Start(ctx); err != nil { + log.Fatalf("start scan / 启动扫描失败: %v", err) } - // Process results in real-time - startTime := time.Now() - var totalVulns int + start := time.Now() + found := 0 - for res := range scanner.ResultChan { - if res == nil { - break + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // The stream closes when the scan finishes, so range terminates. + // 扫描结束时流会关闭，range 自然退出。 + for r := range results { + found++ + fmt.Printf("\n[%d] %s\n", found, r.FullTarget) + fmt.Printf(" poc: %s (%s)\n", r.PocName, r.PocID) + fmt.Printf(" severity: %s\n", r.Severity) } + }() - totalVulns++ - fmt.Printf("\nVulnerability found:\n") - fmt.Printf(" Target: %s\n", res.Target) - fmt.Printf(" POC: %s\n", res.PocInfo.Info.Name) - fmt.Printf(" Severity: %s\n", res.PocInfo.Info.Severity) + if err := scanner.Wait(ctx); err != nil { + log.Printf("scan finished with error / 扫描出错: %v", err) } + wg.Wait() - // Scan completed - fmt.Printf("\nScan completed! Total vulnerabilities: %d\n", totalVulns) - fmt.Printf("Duration: %v\n", time.Since(startTime)) + fmt.Printf("\nscan completed / 扫描完成: %d vulnerabilities in %v\n", found, time.Since(start)) + + if *failOnFind && found > 0 { + fmt.Println("vulnerabilities found, failing / 发现漏洞，返回失败状态") + os.Exit(1) + } } diff --git a/pkg/config/afrogupdate.go b/pkg/config/afrogupdate.go index 7635ed5e0..2f3a874c4 100644 --- a/pkg/config/afrogupdate.go +++ b/pkg/config/afrogupdate.go @@ -1,18 +1,18 @@ package config import ( - "errors" - "fmt" - "io" - "net/http" - "os" - "path/filepath" - "strings" - - "github.com/cavaliergopher/grab/v3" - "github.com/zan8in/afrog/v3/pkg/poc" - "github.com/zan8in/afrog/v3/pkg/utils" - "github.com/zan8in/gologger" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + + "github.com/cavaliergopher/grab/v3" + "github.com/zan8in/afrog/v3/pkg/poc" + "github.com/zan8in/afrog/v3/pkg/utils" + "github.com/zan8in/gologger" ) type AfrogUpdate struct { @@ -25,10 +25,10 @@ type AfrogUpdate struct { } const ( - upHost = "https://gitee.com/zanbin/afrog/raw/main/pocs/v" - upPath = "/afrog-pocs.zip" - upRemoteVersion = "/version" - afrogVersion = "/afrog.version" + upHost = "https://gitee.com/zanbin/afrog/raw/main/pocs/v" + upPath = "/afrog-pocs.zip" + upRemoteVersion = "/version" + afrogVersion = "/afrog.version" ) func NewAfrogUpdate(updatePoc bool) (*AfrogUpdate, error) { @@ -107,48 +107,48 @@ func (u *AfrogUpdate) AfrogUpdatePocs() (string, error) { } func (u *AfrogUpdate) Download() error { - if err := os.RemoveAll(u.HomeDir + upPath); err != nil { - return err - } + if err := os.RemoveAll(u.HomeDir + upPath); err != nil { + return err + } - resp, err := grab.Get(u.HomeDir, upHost+upPath) - if err != nil { - return fmt.Errorf("%s", err.Error()) - } + resp, err := grab.Get(u.HomeDir, upHost+upPath) + if err != nil { + return fmt.Errorf("%s", err.Error()) + } - afHome := filepath.Join(u.HomeDir, ".config", "afrog") - _ = os.MkdirAll(afHome, 0755) - _ = os.RemoveAll(filepath.Join(afHome, "pocs")) + afHome := filepath.Join(u.HomeDir, ".config", "afrog") + _ = os.MkdirAll(afHome, 0755) + _ = os.RemoveAll(filepath.Join(afHome, "pocs")) utils.RandSleep(1000) - u.Unzip(resp.Filename) + u.Unzip(resp.Filename) utils.RandSleep(1000) - u.LastestVersion = u.RemoteVersion + u.LastestVersion = u.RemoteVersion - return os.Remove(resp.Filename) + return os.Remove(resp.Filename) } func (u *AfrogUpdate) Unzip(src string) error { - uz := utils.NewUnzip() - afHome := filepath.Join(u.HomeDir, ".config", "afrog") - if _, err := uz.Extract(src, afHome); err != nil { - return fmt.Errorf("afrog-poc decompression failed. %s", err.Error()) - } - - oldDir := filepath.Join(afHome, "afrog-pocs") - newDir := filepath.Join(afHome, "pocs") - if _, err := os.Stat(oldDir); err == nil { - _ = os.RemoveAll(newDir) - _ = os.Rename(oldDir, newDir) - } - - if len(u.RemoteVersion) > 0 { - u.CurrVersion = u.RemoteVersion - } - gologger.Print().Msgf("Successfully installed pocs at %s\n", strings.ReplaceAll(newDir, "\\", "/")) - - return nil + uz := utils.NewUnzip() + afHome := filepath.Join(u.HomeDir, ".config", "afrog") + if _, err := uz.Extract(src, afHome); err != nil { + return fmt.Errorf("afrog-poc decompression failed. %s", err.Error()) + } + + oldDir := filepath.Join(afHome, "afrog-pocs") + newDir := filepath.Join(afHome, "pocs") + if _, err := os.Stat(oldDir); err == nil { + _ = os.RemoveAll(newDir) + _ = os.Rename(oldDir, newDir) + } + + if len(u.RemoteVersion) > 0 { + u.CurrVersion = u.RemoteVersion + } + gologger.Print().Msgf("Successfully installed pocs at %s\n", strings.ReplaceAll(newDir, "\\", "/")) + + return nil } diff --git a/pkg/config/config.go b/pkg/config/config.go index 16b93a6ff..9251965e3 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -205,6 +205,68 @@ func NewConfig(configFile string) (*Config, error) { return ReadConfiguration(configFile) } +// LoadConfigReadOnly loads the afrog configuration without touching the filesystem. +// +// Unlike [NewConfig], it never creates ~/.config/afrog, never writes +// afrog-config.yaml, and never rewrites an existing user config to inject new +// sections. Library consumers must not have side effects on the host's home +// directory merely by constructing a scanner, so the SDK uses this instead. +// +// When no configuration file exists, it returns a Config populated with the +// same defaults NewConfig would have written. +func LoadConfigReadOnly(configFile string) (*Config, error) { + if len(configFile) > 0 && !strings.HasSuffix(configFile, ".yml") && !strings.HasSuffix(configFile, ".yaml") { + return nil, errors.New("afrog config file must be yaml format") + } + + path := configFile + if path == "" { + homeDir, err := os.UserHomeDir() + if err != nil { + return defaultConfig(), nil + } + path = filepath.Join(homeDir, ".config", "afrog", afrogConfigFilename) + } + + file, err := os.Open(path) + if err != nil { + if configFile != "" { + return nil, err + } + return defaultConfig(), nil + } + defer file.Close() + + config := &Config{} + if err := yaml.NewDecoder(file).Decode(config); err != nil { + return nil, err + } + normalizeCuratedDefaults(config) + normalizeInteractshDefaults(config) + return config, nil +} + +// defaultConfig returns the built-in configuration defaults. +func defaultConfig() *Config { + c := &Config{ServerAddress: ":16868"} + + c.Reverse.Dnslogcn.Domain = "dnslog.cn" + c.Reverse.Xray.ApiUrl = "http://x.x.x.x:8777" + c.Reverse.Interactsh.Server = "oast.pro" + + c.Webhook.Dingtalk.Range = "high,critical" + c.Webhook.Wecom.Range = "high,critical" + c.Webhook.Wecom.Markdown = true + + autoUpdate := true + c.Curated.Enabled = "auto" + c.Curated.AutoUpdate = &autoUpdate + c.Curated.TimeoutSec = 10 + c.Curated.Channel = "stable" + + return c +} + func isExistConfigFile(configFile string) error { if len(configFile) > 0 { if utils.Exists(configFile) { diff --git a/pkg/config/oobadapter.go b/pkg/config/oobadapter.go index b0985d575..a20e38246 100644 --- a/pkg/config/oobadapter.go +++ b/pkg/config/oobadapter.go @@ -1,11 +1,11 @@ package config var ( - OOBCeyeio = "ceyeio" - OOBDnslogcn = "dnslogcn" - OOBAlphalog = "alphalog" - OOBXray = "xray" - OOBRevsuit = "revsuit" + OOBCeyeio = "ceyeio" + OOBDnslogcn = "dnslogcn" + OOBAlphalog = "alphalog" + OOBXray = "xray" + OOBRevsuit = "revsuit" OOBInteractsh = "interactsh" ) diff --git a/pkg/config/options.go b/pkg/config/options.go index 1017b9aed..b0a5cbb58 100644 --- a/pkg/config/options.go +++ b/pkg/config/options.go @@ -11,7 +11,6 @@ import ( "sync" "github.com/rs/xid" - "github.com/zan8in/afrog/v3/pkg/catalog" "github.com/zan8in/afrog/v3/pkg/db/sqlite" "github.com/zan8in/afrog/v3/pkg/log" "github.com/zan8in/afrog/v3/pkg/output" @@ -46,9 +45,18 @@ type Options struct { // list of target URLs/hosts to scan (one per line) TargetsFile string - // PoC file or directory to scan + // PoC file or directory to scan. + // Deprecated: 语义为“独占”（屏蔽内置与用户目录 PoC），新代码请使用 PocPaths + PocPathsOnly。 PocFile string + // PocPaths 是追加式的 PoC 输入，支持文件、目录和 glob 通配符（如 "dir/*.yaml"）。 + // 与内置、curated、my、local 来源合并，同名时以 PocPaths 优先。 + PocPaths []string + + // PocPathsOnly 为 true 时只加载 PocPaths/AppendPoc/PocFile 指定的 PoC， + // 屏蔽内置与用户目录来源。 + PocPathsOnly bool + // Append PoC file or directory to scan AppendPoc goflags.StringSlice @@ -1080,23 +1088,22 @@ func detectLegacyOOBReasons(yamlText string) []string { return reasons } +// CreatePocList 返回本次扫描要执行的 PoC 列表。 +// 加载过程中被跳过的 PoC 会被丢弃，如需诊断信息请使用 CreatePocListWithDiagnostics。 func (o *Options) CreatePocList() []poc.Poc { + pocList, _ := o.CreatePocListWithDiagnostics() + return pocList +} + +// CreatePocListWithDiagnostics 在返回 PoC 列表的同时，返回所有被跳过的 PoC 及其原因。 +func (o *Options) CreatePocListWithDiagnostics() ([]poc.Poc, []PocLoadError) { type legacyItem struct { ID string Path string Reasons []string } - pathItems := []pocsrepo.PathItem{} - if strings.TrimSpace(o.PocFile) != "" { - c := catalog.New(o.PocFile) - paths, _ := c.GetPocPath(o.PocFile) - for _, pth := range paths { - pathItems = append(pathItems, pocsrepo.PathItem{Path: pth, Source: pocsrepo.SourceLocal}) - } - } else { - pathItems, _ = pocsrepo.CollectOrderedPocPaths(o.AppendPoc) - } + pathItems, diagnostics := o.resolvePocPathItems() newPocSlice := make([]poc.Poc, 0, len(pathItems)) legacy := make([]legacyItem, 0) @@ -1116,13 +1123,15 @@ func (o *Options) CreatePocList() []poc.Poc { raw, err = os.ReadFile(srcPath) } if err != nil { - gologger.Error().Msgf("Invalid POC format, discard: %s, error: %v", srcPath, err) + diagnostics = append(diagnostics, PocLoadError{Path: it.Path, Reason: PocLoadReadFailed, Err: err}) + o.logPocDiscarded(srcPath, err) continue } pm := poc.PocMeta{} if e := yaml.Unmarshal(raw, &pm); e != nil { - gologger.Error().Msgf("Invalid POC format, discard: %s, error: %v", srcPath, e) + diagnostics = append(diagnostics, PocLoadError{Path: it.Path, Reason: PocLoadParseFailed, Err: e}) + o.logPocDiscarded(srcPath, e) continue } id := strings.TrimSpace(pm.Id) @@ -1137,12 +1146,19 @@ func (o *Options) CreatePocList() []poc.Poc { reasons := detectLegacyOOBReasons(string(raw)) if len(reasons) > 0 { legacy = append(legacy, legacyItem{ID: id, Path: srcPath, Reasons: reasons}) + diagnostics = append(diagnostics, PocLoadError{ + Path: it.Path, + ID: id, + Reason: PocLoadLegacyOOB, + Detail: strings.Join(reasons, "; "), + }) continue } pp := poc.Poc{} if e := yaml.Unmarshal(raw, &pp); e != nil { - gologger.Error().Msgf("Invalid POC format, discard: %s, error: %v", srcPath, e) + diagnostics = append(diagnostics, PocLoadError{Path: it.Path, ID: id, Reason: PocLoadParseFailed, Err: e}) + o.logPocDiscarded(srcPath, e) continue } pp.Id = strings.TrimSpace(pp.Id) @@ -1161,7 +1177,7 @@ func (o *Options) CreatePocList() []poc.Poc { newPocSlice = append(newPocSlice, pp) } - if len(legacy) > 0 { + if len(legacy) > 0 && !o.consoleQuiet() { total := len(legacy) const limit = 20 gologger.Print().Msgf("检测到旧OOB POC（已跳过）：%d 个", total) @@ -1210,7 +1226,19 @@ func (o *Options) CreatePocList() []poc.Poc { sort.Sort(POCSlices(finalPocSlice)) } - return finalPocSlice + return finalPocSlice, diagnostics +} + +// consoleQuiet 表示当前不应该向控制台输出内容（SDK 模式或显式静默）。 +func (o *Options) consoleQuiet() bool { + return o != nil && (o.SDKMode || o.Silent) +} + +func (o *Options) logPocDiscarded(path string, err error) { + if o.consoleQuiet() { + return + } + gologger.Error().Msgf("Invalid POC format, discard: %s, error: %v", path, err) } // 定义包含 POC 结构的切片 diff --git a/pkg/config/options_legacy_oob_test.go b/pkg/config/options_legacy_oob_test.go index 6c66e0859..3b7d74572 100644 --- a/pkg/config/options_legacy_oob_test.go +++ b/pkg/config/options_legacy_oob_test.go @@ -76,4 +76,3 @@ expression: r0() }) } } - diff --git a/pkg/config/pocmerge_test.go b/pkg/config/pocmerge_test.go new file mode 100644 index 000000000..313f9152f --- /dev/null +++ b/pkg/config/pocmerge_test.go @@ -0,0 +1,104 @@ +package config + +import ( + "os" + "path/filepath" + "sort" + "testing" +) + +// Distinct PoCs that happen to share a filename must all survive: silently +// dropping one would make the user scan less than they asked for, with no +// diagnostic to explain it. +func TestCreatePocList_KeepsSameNamedPocsFromDifferentDirs(t *testing.T) { + dirA := t.TempDir() + dirB := t.TempDir() + writePoc(t, dirA, "shared.yaml", "poc-from-dir-a") + writePoc(t, dirB, "shared.yaml", "poc-from-dir-b") + + opt := &Options{ + PocPaths: []string{dirA, dirB}, + PocPathsOnly: true, + SDKMode: true, + } + + got := opt.CreatePocList() + ids := make([]string, 0, len(got)) + for _, p := range got { + ids = append(ids, p.Id) + } + sort.Strings(ids) + + want := []string{"poc-from-dir-a", "poc-from-dir-b"} + if len(ids) != len(want) { + t.Fatalf("loaded %v, want %v", ids, want) + } + for i := range want { + if ids[i] != want[i] { + t.Fatalf("loaded %v, want %v", ids, want) + } + } +} + +func TestCreatePocList_ExplicitPathOverridesBuiltinOfSameName(t *testing.T) { + // A PoC named after a built-in one should replace it rather than run twice. + dir := t.TempDir() + writePoc(t, dir, "afrog-demo.yaml", "my-override") + + opt := &Options{PocPaths: []string{dir}, SDKMode: true} + got := opt.CreatePocList() + + var overrides, builtins int + for _, p := range got { + switch p.Id { + case "my-override": + overrides++ + case "afrog-demo": + builtins++ + } + } + if overrides != 1 { + t.Errorf("explicit poc appeared %d times, want 1", overrides) + } + if builtins != 0 { + t.Errorf("built-in poc of the same filename appeared %d times, want 0", builtins) + } +} + +func TestResolvePocInputs_DeduplicatesByAbsolutePath(t *testing.T) { + dir := t.TempDir() + p := writePoc(t, dir, "one.yaml", "one") + + // The same file reachable two ways must resolve to a single entry. + paths, diags := ResolvePocInputs([]string{p, dir, filepath.Join(dir, "*.yaml")}) + if len(diags) != 0 { + t.Fatalf("unexpected diagnostics: %v", diags) + } + if len(paths) != 1 { + t.Fatalf("resolved %d paths, want 1: %v", len(paths), paths) + } +} + +func TestLoadConfigReadOnly_MissingHomeConfigReturnsDefaults(t *testing.T) { + cfg, err := LoadConfigReadOnly(filepath.Join(t.TempDir(), "nope.yaml")) + if err == nil { + t.Fatal("an explicitly requested missing file should error") + } + if cfg != nil { + t.Fatal("no config should be returned on error") + } + + // The default path falls back to in-memory defaults without touching disk. + home := t.TempDir() + t.Setenv("HOME", home) + cfg, err = LoadConfigReadOnly("") + if err != nil { + t.Fatalf("LoadConfigReadOnly: %v", err) + } + if cfg == nil || cfg.ServerAddress == "" { + t.Fatal("expected populated defaults when no config file exists") + } + if _, err := os.Stat(filepath.Join(home, ".config", "afrog")); !os.IsNotExist(err) { + t.Error("LoadConfigReadOnly must not create the config directory") + } +} diff --git a/pkg/config/pocsource.go b/pkg/config/pocsource.go new file mode 100644 index 000000000..84585f5ae --- /dev/null +++ b/pkg/config/pocsource.go @@ -0,0 +1,165 @@ +package config + +import ( + "errors" + "fmt" + "path/filepath" + "strings" + + "github.com/zan8in/afrog/v3/pkg/catalog" + "github.com/zan8in/afrog/v3/pkg/pocsrepo" +) + +// POC 加载失败的原因分类，便于调用方做机器判断而不是匹配错误字符串。 +const ( + PocLoadNotFound = "not_found" + PocLoadReadFailed = "read_failed" + PocLoadParseFailed = "parse_failed" + PocLoadLegacyOOB = "legacy_oob" +) + +// ErrNoPocMatched 表示路径本身存在，但没有匹配到任何 .yaml/.yml 文件。 +var ErrNoPocMatched = errors.New("no poc file matched") + +// PocLoadError 描述一个在加载阶段被跳过的 POC。 +// 以前这些信息只会被打印到控制台，SDK 调用方无从得知，现在通过诊断列表返回。 +type PocLoadError struct { + // Path 是 POC 文件路径；内置 POC 带 "embedded:" 前缀。 + Path string + // ID 是解析出的 POC id，解析失败时为空。 + ID string + // Reason 是 PocLoad* 常量之一。 + Reason string + // Detail 提供人类可读的补充说明，例如旧 OOB 语法的具体命中项。 + Detail string + // Err 是底层错误，可能为 nil。 + Err error +} + +func (e PocLoadError) Error() string { + msg := fmt.Sprintf("poc %s: %s", e.Path, e.Reason) + if e.Detail != "" { + msg += " (" + e.Detail + ")" + } + if e.Err != nil { + msg += ": " + e.Err.Error() + } + return msg +} + +func (e PocLoadError) Unwrap() error { return e.Err } + +// pocInputs 汇总用户显式指定的 POC 输入。 +// PocFile、PocPaths、AppendPoc 三者会被合并，这修复了以前 PocFile 一旦设置 +// AppendPoc 就被静默丢弃的问题。 +func (o *Options) pocInputs() []string { + inputs := make([]string, 0, 1+len(o.PocPaths)+len(o.AppendPoc)) + if v := strings.TrimSpace(o.PocFile); v != "" { + inputs = append(inputs, v) + } + for _, p := range o.PocPaths { + if v := strings.TrimSpace(p); v != "" { + inputs = append(inputs, v) + } + } + for _, p := range o.AppendPoc { + if v := strings.TrimSpace(p); v != "" { + inputs = append(inputs, v) + } + } + return inputs +} + +// pocInputsExclusive 表示是否只使用显式指定的 POC，屏蔽内置/curated/my/local 来源。 +// PocFile 保留历史语义（独占），PocPaths 则是追加语义。 +func (o *Options) pocInputsExclusive() bool { + return o.PocPathsOnly || strings.TrimSpace(o.PocFile) != "" +} + +// ResolvePocInputs 把文件、目录、glob 通配符统一解析成具体的 POC 文件列表。 +func ResolvePocInputs(inputs []string) ([]string, []PocLoadError) { + out := make([]string, 0, len(inputs)) + diags := make([]PocLoadError, 0) + seen := make(map[string]struct{}) + + for _, in := range inputs { + in = strings.TrimSpace(in) + if in == "" { + continue + } + paths, err := catalog.New(in).GetPocPath(in) + if err != nil { + diags = append(diags, PocLoadError{Path: in, Reason: PocLoadNotFound, Err: err}) + continue + } + if len(paths) == 0 { + diags = append(diags, PocLoadError{Path: in, Reason: PocLoadNotFound, Err: ErrNoPocMatched}) + continue + } + for _, p := range paths { + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + out = append(out, p) + } + } + return out, diags +} + +// ValidatePocInputs 在扫描前校验 POC 输入是否可用。 +// 相比直接 os.Stat，它能正确处理 glob 通配符和目录。 +func ValidatePocInputs(inputs []string) error { + if len(inputs) == 0 { + return nil + } + matched, diags := ResolvePocInputs(inputs) + if len(matched) > 0 { + return nil + } + if len(diags) > 0 { + return diags[0] + } + return ErrNoPocMatched +} + +// resolvePocPathItems 返回本次扫描最终要加载的 POC 路径集合。 +func (o *Options) resolvePocPathItems() ([]pocsrepo.PathItem, []PocLoadError) { + explicit, diags := ResolvePocInputs(o.pocInputs()) + + // 用户显式指定的 POC 全部保留。ResolvePocInputs 已按绝对路径去重， + // 这里不能再按文件名去重：不同目录下的同名 POC 是两个不同的 POC， + // 静默丢弃其中一个会让用户少扫内容且无从察觉。 + items := make([]pocsrepo.PathItem, 0, len(explicit)) + explicitNames := make(map[string]struct{}, len(explicit)) + source := pocsrepo.SourceAppend + if o.pocInputsExclusive() { + source = pocsrepo.SourceLocal + } + for _, p := range explicit { + items = append(items, pocsrepo.PathItem{Path: p, Source: source}) + explicitNames[pocBaseName(p)] = struct{}{} + } + + if o.pocInputsExclusive() { + return items, diags + } + + // 合并内置/curated/my/local 来源。同名时以显式指定的为准， + // 这是"用自己的版本覆盖内置版本"的预期语义。 + base, _ := pocsrepo.CollectOrderedPocPaths(nil) + for _, it := range base { + if _, overridden := explicitNames[pocBaseName(it.Path)]; overridden { + continue + } + items = append(items, it) + } + return items, diags +} + +// pocBaseName 返回 POC 文件名（小写、去扩展名），用于同名覆盖判断。 +func pocBaseName(path string) string { + name := filepath.Base(strings.ReplaceAll(path, "\\", "/")) + name = strings.TrimSuffix(strings.TrimSuffix(name, ".yaml"), ".yml") + return strings.ToLower(name) +} diff --git a/pkg/config/pocsource_test.go b/pkg/config/pocsource_test.go new file mode 100644 index 000000000..75f4a770e --- /dev/null +++ b/pkg/config/pocsource_test.go @@ -0,0 +1,258 @@ +package config + +import ( + "errors" + "os" + "path/filepath" + "sort" + "testing" +) + +// writePoc creates a minimal valid PoC file and returns its path. +func writePoc(t *testing.T, dir, name, id string) string { + t.Helper() + path := filepath.Join(dir, name) + body := "id: " + id + ` +info: + name: ` + id + ` + author: test + severity: info +rules: + r0: + request: + method: GET + path: / + expression: response.status == 200 +expression: r0() +` + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + t.Fatalf("write poc %s: %v", path, err) + } + return path +} + +func pocIDs(pocs []struct{ ID string }) []string { + out := make([]string, 0, len(pocs)) + for _, p := range pocs { + out = append(out, p.ID) + } + sort.Strings(out) + return out +} + +func TestResolvePocInputs(t *testing.T) { + dir := t.TempDir() + single := writePoc(t, dir, "single.yaml", "single") + writePoc(t, dir, "second.yml", "second") + + nested := filepath.Join(dir, "nested") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + writePoc(t, nested, "deep.yaml", "deep") + + tests := []struct { + name string + inputs []string + wantCount int + wantDiags int + }{ + {name: "single file", inputs: []string{single}, wantCount: 1}, + {name: "directory is walked recursively", inputs: []string{dir}, wantCount: 3}, + {name: "glob pattern", inputs: []string{filepath.Join(dir, "*.yaml")}, wantCount: 1}, + {name: "glob matches both extensions", inputs: []string{filepath.Join(dir, "*.y*ml")}, wantCount: 2}, + {name: "duplicates are removed", inputs: []string{single, single}, wantCount: 1}, + {name: "missing path is reported", inputs: []string{filepath.Join(dir, "nope.yaml")}, wantCount: 0, wantDiags: 1}, + {name: "empty input is ignored", inputs: []string{" "}, wantCount: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, diags := ResolvePocInputs(tt.inputs) + if len(got) != tt.wantCount { + t.Errorf("resolved %d paths, want %d (%v)", len(got), tt.wantCount, got) + } + if len(diags) != tt.wantDiags { + t.Errorf("got %d diagnostics, want %d (%v)", len(diags), tt.wantDiags, diags) + } + }) + } +} + +func TestValidatePocInputs(t *testing.T) { + dir := t.TempDir() + writePoc(t, dir, "a.yaml", "a") + + tests := []struct { + name string + inputs []string + wantErr bool + }{ + {name: "no input is valid", inputs: nil}, + {name: "existing directory", inputs: []string{dir}}, + {name: "glob that matches", inputs: []string{filepath.Join(dir, "*.yaml")}}, + {name: "glob that matches nothing", inputs: []string{filepath.Join(dir, "*.json")}, wantErr: true}, + {name: "missing path", inputs: []string{filepath.Join(dir, "missing")}, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidatePocInputs(tt.inputs) + if tt.wantErr && err == nil { + t.Fatal("expected an error, got nil") + } + if !tt.wantErr && err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + } +} + +// AppendPoc used to be silently dropped whenever PocFile was set. +func TestCreatePocList_MergesPocFileAndAppendPoc(t *testing.T) { + dirA := t.TempDir() + dirB := t.TempDir() + fromPocFile := writePoc(t, dirA, "a.yaml", "from-pocfile") + writePoc(t, dirB, "b.yaml", "from-appendpoc") + + opt := &Options{ + PocFile: fromPocFile, + AppendPoc: []string{dirB}, + SDKMode: true, + } + + got := opt.CreatePocList() + ids := make([]string, 0, len(got)) + for _, p := range got { + ids = append(ids, p.Id) + } + sort.Strings(ids) + + want := []string{"from-appendpoc", "from-pocfile"} + if len(ids) != len(want) { + t.Fatalf("loaded %v, want %v", ids, want) + } + for i := range want { + if ids[i] != want[i] { + t.Fatalf("loaded %v, want %v", ids, want) + } + } +} + +func TestCreatePocList_PocPathsSupportsGlob(t *testing.T) { + dir := t.TempDir() + writePoc(t, dir, "one.yaml", "one") + writePoc(t, dir, "two.yaml", "two") + writePoc(t, dir, "skip.txt", "skip") + + opt := &Options{ + PocPaths: []string{filepath.Join(dir, "*.yaml")}, + PocPathsOnly: true, + SDKMode: true, + } + + got := opt.CreatePocList() + if len(got) != 2 { + ids := make([]string, 0, len(got)) + for _, p := range got { + ids = append(ids, p.Id) + } + t.Fatalf("loaded %d pocs %v, want 2", len(got), ids) + } +} + +func TestCreatePocListWithDiagnostics_ReportsSkippedPocs(t *testing.T) { + dir := t.TempDir() + writePoc(t, dir, "good.yaml", "good") + + // Legacy v2 OOB syntax is skipped; the caller should be able to find out why. + legacy := filepath.Join(dir, "legacy.yaml") + legacyBody := `id: legacy-oob +info: + name: legacy + author: test + severity: info +set: + oob: oob() +rules: + r0: + request: + method: GET + path: /?x={{oobDNS}} + expression: oobCheck(oob, 5) +expression: r0() +` + if err := os.WriteFile(legacy, []byte(legacyBody), 0o644); err != nil { + t.Fatalf("write legacy poc: %v", err) + } + + opt := &Options{PocPaths: []string{dir}, PocPathsOnly: true, SDKMode: true} + pocs, diags := opt.CreatePocListWithDiagnostics() + + if len(pocs) != 1 || pocs[0].Id != "good" { + t.Fatalf("expected only the valid poc, got %d", len(pocs)) + } + + var found bool + for _, d := range diags { + if d.Reason == PocLoadLegacyOOB { + found = true + if d.Detail == "" { + t.Error("legacy diagnostic should explain which syntax was matched") + } + } + } + if !found { + t.Fatalf("expected a %s diagnostic, got %v", PocLoadLegacyOOB, diags) + } +} + +func TestPocLoadError_UnwrapsUnderlyingError(t *testing.T) { + sentinel := errors.New("boom") + err := PocLoadError{Path: "/tmp/x.yaml", Reason: PocLoadReadFailed, Err: sentinel} + + if !errors.Is(err, sentinel) { + t.Fatal("PocLoadError should unwrap to the underlying error") + } + if err.Error() == "" { + t.Fatal("PocLoadError should render a message") + } +} + +func TestLoadConfigReadOnly_DoesNotWriteAnything(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "afrog-config.yaml") + + // Missing file must not be created. + if _, err := LoadConfigReadOnly(path); err == nil { + t.Fatal("expected an error for an explicitly requested missing config file") + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatal("LoadConfigReadOnly must not create the config file") + } + + // Existing file must be read back without being rewritten. + if err := os.WriteFile(path, []byte("server: \":9999\"\n"), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + before, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read config: %v", err) + } + + cfg, err := LoadConfigReadOnly(path) + if err != nil { + t.Fatalf("LoadConfigReadOnly: %v", err) + } + if cfg.ServerAddress != ":9999" { + t.Errorf("ServerAddress = %q, want \":9999\"", cfg.ServerAddress) + } + + after, err := os.ReadFile(path) + if err != nil { + t.Fatalf("re-read config: %v", err) + } + if string(before) != string(after) { + t.Error("LoadConfigReadOnly must not rewrite an existing config file") + } +} diff --git a/pkg/cyberspace/cyberspace.go b/pkg/cyberspace/cyberspace.go index ea82049e7..34bbcec51 100644 --- a/pkg/cyberspace/cyberspace.go +++ b/pkg/cyberspace/cyberspace.go @@ -14,6 +14,9 @@ type Cyberspace struct { Engine string Query string QueryCount int + // Quiet suppresses the search progress output. Library embedders such as + // the SDK must not have this written to their process stdout. + Quiet bool } func New(config *config.Config, engine, query string, queryCount int) (*Cyberspace, error) { @@ -107,13 +110,16 @@ func (c *Cyberspace) GetTargets() ([]string, error) { break } } - fmt.Printf("\rZoomEye Searching... Total: %d, Query Count: %d, Current: %d", total, c.QueryCount, currentTotal) + if !c.Quiet { + fmt.Printf("\rZoomEye Searching... Total: %d, Query Count: %d, Current: %d", total, c.QueryCount, currentTotal) + } } - fmt.Println("") - - if currentTotal == 0 { - gologger.Info().Msg("no result found") + if !c.Quiet { + fmt.Println("") + if currentTotal == 0 { + gologger.Info().Msg("no result found") + } } return results, nil diff --git a/pkg/portscan/scan.go b/pkg/portscan/scan.go index a45613604..bb7c68b1a 100644 --- a/pkg/portscan/scan.go +++ b/pkg/portscan/scan.go @@ -450,7 +450,9 @@ func (s *Scanner) maybeAdaptiveDelay() { // Returns the open connection (if successful) or error. // Caller is responsible for closing the connection. func (s *Scanner) checkPortOpen(host string, port int) (net.Conn, error) { - address := fmt.Sprintf("%s:%d", host, port) + // net.JoinHostPort brackets IPv6 literals; fmt.Sprintf("%s:%d") does not + // and produces an address net.Dial cannot parse. + address := net.JoinHostPort(host, strconv.Itoa(port)) var conn net.Conn var err error diff --git a/pkg/protocols/http/retryhttpclient/client.go b/pkg/protocols/http/retryhttpclient/client.go index 364d174f1..3c9fe8dd7 100644 --- a/pkg/protocols/http/retryhttpclient/client.go +++ b/pkg/protocols/http/retryhttpclient/client.go @@ -46,6 +46,10 @@ var ( taskGateWaitCount int64 ) +// VarResponseBodyTruncated 是写入 variableMap 的内部键（不参与 CEL 求值）， +// 用于把“响应体被 MaxRespBodySize 截断”这一事实传递给上层结果构造。 +const VarResponseBodyTruncated = "__response_body_truncated" + type Options struct { Proxy string Timeout int @@ -693,8 +697,17 @@ func Request(target string, header []string, rule poc.Rule, variableMap map[stri return err } } + // 读满上限时再探一个字节，用来区分“正好等于上限”和“确实被截断”。 + bodyTruncated := false + if lr.N <= 0 { + var probe [1]byte + if n, _ := resp.Body.Read(probe[:]); n > 0 { + bodyTruncated = true + } + } resp.Body.Close() respBody := buf.Bytes() + variableMap[VarResponseBodyTruncated] = bodyTruncated responseText := "" if len(respBody) > 0 { diff --git a/pkg/result/result.go b/pkg/result/result.go index 95e0e75eb..617fe2ad1 100644 --- a/pkg/result/result.go +++ b/pkg/result/result.go @@ -33,6 +33,9 @@ type PocResult struct { IsVul bool BruteTruncated bool BruteRequests int + // BodyTruncated 表示响应体在 MaxRespBodySize 上限处被截断， + // 此时 ResultResponse.Body 不是完整的服务端响应。 + BodyTruncated bool } func DebugDumpRequestText(pr *PocResult) string { @@ -102,6 +105,7 @@ func (pr *PocResult) Snapshot() *PocResult { IsVul: pr.IsVul, BruteTruncated: pr.BruteTruncated, BruteRequests: pr.BruteRequests, + BodyTruncated: pr.BodyTruncated, } } diff --git a/pkg/runner/checker.go b/pkg/runner/checker.go index 900487074..f8a02d60f 100644 --- a/pkg/runner/checker.go +++ b/pkg/runner/checker.go @@ -148,6 +148,15 @@ func celSafeIdent(s string) string { return string(b) } +// responseBodyTruncated 读取最近一次 HTTP 响应是否因 MaxRespBodySize 被截断。 +func (c *Checker) responseBodyTruncated() bool { + if c == nil || c.VariableMap == nil { + return false + } + t, _ := c.VariableMap[retryhttpclient.VarResponseBodyTruncated].(bool) + return t +} + func (c *Checker) Check(target string, pocItem *poc.Poc) (err error) { defer func() { if r := recover(); r != nil { @@ -438,9 +447,15 @@ func (c *Checker) Check(target string, pocItem *poc.Poc) (err error) { last.IsVul = isMatch last.BruteTruncated = bruteTruncated last.BruteRequests = bruteRequests + last.BodyTruncated = c.responseBodyTruncated() c.Result.AllPocResult = append(c.Result.AllPocResult, stepResults...) } else { - pocRstTemp := result.PocResult{IsVul: isMatch, BruteTruncated: bruteTruncated, BruteRequests: bruteRequests} + pocRstTemp := result.PocResult{ + IsVul: isMatch, + BruteTruncated: bruteTruncated, + BruteRequests: bruteRequests, + BodyTruncated: c.responseBodyTruncated(), + } if c.VariableMap["response"] != nil { pocRstTemp.ResultResponse = c.VariableMap["response"].(*proto.Response) } diff --git a/pkg/runner/engine.go b/pkg/runner/engine.go index cfcf7faae..b47b68d8e 100644 --- a/pkg/runner/engine.go +++ b/pkg/runner/engine.go @@ -462,9 +462,9 @@ func (e *Engine) AcquireChecker() *Checker { c.Result.Output = e.options.Output c.OOBAdapter = e.oobAdapter c.OOBAlive = e.oobAlive - c.OOBMgr = e.oobMgr + c.OOBMgr = e.OOBMgr() if c.CustomLib != nil { - c.CustomLib.SetOOBManager(e.oobMgr) + c.CustomLib.SetOOBManager(e.OOBMgr()) } return c } @@ -481,8 +481,10 @@ func (e *Engine) ReleaseChecker(c *Checker) { } type Engine struct { - options *config.Options - ticker *time.Ticker + options *config.Options + // ticker paces task scheduling. It is created per stage and stopped from + // Stop, which runs on a different goroutine, so access must be atomic. + ticker atomic.Pointer[time.Ticker] mu sync.Mutex paused uint32 stopped uint32 @@ -500,7 +502,43 @@ type Engine struct { pedmStop chan struct{} oobAdapter *oobadapter.OOBAdapter oobAlive bool - oobMgr *OOBManager + // oobMgr is swapped between scans and read lock-free from the scan hot + // path, so it must be an atomic pointer rather than a plain field. + oobMgr atomic.Pointer[OOBManager] +} + +// setTicker installs the scheduling ticker, stopping any previous one. +func (e *Engine) setTicker(t *time.Ticker) { + if old := e.ticker.Swap(t); old != nil { + old.Stop() + } +} + +// stopTicker stops and clears the scheduling ticker. It is safe to call when +// no ticker is installed and safe to call concurrently. +func (e *Engine) stopTicker() { + if old := e.ticker.Swap(nil); old != nil { + old.Stop() + } +} + +// OOBMgr returns the active out-of-band manager, or nil when OOB is inactive. +func (e *Engine) OOBMgr() *OOBManager { + if e == nil { + return nil + } + return e.oobMgr.Load() +} + +// setOOBManager installs a manager, stopping any previous one so its poll +// goroutine cannot outlive the swap. +func (e *Engine) setOOBManager(m *OOBManager) { + if e == nil { + return + } + if old := e.oobMgr.Swap(m); old != nil { + old.Stop() + } } func NewEngine(options *config.Options) *Engine { @@ -925,9 +963,11 @@ func (runner *Runner) Execute() { } } if runner.engine != nil { + // Stop any manager left over from a previous run before dropping the + // reference, otherwise its poll goroutine would leak. + runner.engine.stopOOBManager() runner.engine.oobAdapter = nil runner.engine.oobAlive = false - runner.engine.oobMgr = nil } pocSlice := options.CreatePocList() @@ -973,7 +1013,11 @@ func (runner *Runner) Execute() { if runner.engine != nil && runner.engine.oobAlive && runner.engine.oobAdapter != nil { pollInterval := time.Duration(options.OOBPollInterval) * time.Second hitRetention := time.Duration(options.OOBHitRetention) * time.Minute - runner.engine.oobMgr = NewOOBManager(runner.ctx, runner.engine.oobAdapter, pollInterval, hitRetention) + runner.engine.setOOBManager(NewOOBManager(runner.ctx, runner.engine.oobAdapter, pollInterval, hitRetention)) + // The poll loop must not outlive the scan. Without this the goroutine + // survives every completed scan, because runner.ctx is only cancelled + // by an explicit Stop. + defer runner.engine.stopOOBManager() } runner.startOOBResolver() defer runner.stopOOBResolver() @@ -1453,13 +1497,14 @@ func (runner *Runner) Execute() { concurrency = 1 } - runner.engine.ticker = time.NewTicker(time.Second / time.Duration(rate)) - defer func() { - if runner.engine.ticker != nil { - runner.engine.ticker.Stop() - runner.engine.ticker = nil - } - }() + interval := time.Second / time.Duration(rate) + if interval <= 0 { + // time.NewTicker panics on a non-positive interval, which a very + // high rate limit would otherwise produce. + interval = time.Nanosecond + } + runner.engine.setTicker(time.NewTicker(interval)) + defer runner.engine.stopTicker() type stageTask struct { tap *TransData @@ -1692,6 +1737,18 @@ func (runner *Runner) exec(tap *TransData) { } } +// emitFailure 把单个 PoC 的执行失败上报给调用方，不影响扫描继续进行。 +func (runner *Runner) emitFailure(target string, p *poc.Poc, err error) { + if runner == nil || runner.OnFailure == nil || err == nil { + return + } + pocID := "" + if p != nil { + pocID = p.Id + } + runner.OnFailure(target, pocID, err) +} + func (runner *Runner) executeExpression(ctx context.Context, target string, poc *poc.Poc) { c := runner.engine.AcquireChecker() defer runner.engine.ReleaseChecker(c) @@ -1703,6 +1760,7 @@ func (runner *Runner) executeExpression(ctx context.Context, target string, poc // https://github.com/zan8in/afrog/v3/issues/7 if r := recover(); r != nil { c.Result.IsVul = false + runner.emitFailure(target, poc, fmt.Errorf("panic: %v", r)) runner.OnResult(c.Result) } }() @@ -1710,7 +1768,9 @@ func (runner *Runner) executeExpression(ctx context.Context, target string, poc if ctx != nil { c.VariableMap[retryhttpclient.ContextVarKey] = ctx } - c.Check(target, poc) + if err := c.Check(target, poc); err != nil { + runner.emitFailure(target, poc, err) + } if c.Result != nil { c.Result.FingerResult = runner.fingerprintForTarget(c.Result.Target) } @@ -1865,12 +1925,13 @@ func (runner *Runner) fingerprintForTarget(target string) []fingerprint.Hit { } func (e *Engine) waitTick() { - if e.ticker == nil { + ticker := e.ticker.Load() + if ticker == nil { return } start := time.Now() select { - case <-e.ticker.C: + case <-ticker.C: case <-e.quit: return } @@ -1903,15 +1964,23 @@ func (e *Engine) Stop() { if !atomic.CompareAndSwapUint32(&e.stopped, 0, 1) { return } - e.mu.Lock() - if e.ticker != nil { - e.ticker.Stop() - } - e.mu.Unlock() + e.stopTicker() close(e.quit) + e.stopOOBManager() gologger.Debug().Msgf("engine stopped: ticker stopped and scheduling halted") } +// stopOOBManager terminates the OOB poll goroutine and clears the reference. +// It is safe to call when no manager is running. +func (e *Engine) stopOOBManager() { + if e == nil { + return + } + if mgr := e.oobMgr.Swap(nil); mgr != nil { + mgr.Stop() + } +} + func (runner *Runner) NotVulCallback() { runner.OnResult(&result.Result{IsVul: false}) } diff --git a/pkg/runner/oob_manager.go b/pkg/runner/oob_manager.go index aadb902a1..ac12f2b72 100644 --- a/pkg/runner/oob_manager.go +++ b/pkg/runner/oob_manager.go @@ -26,6 +26,15 @@ type OOBManager struct { maxSeen int lastPolledAt map[string]time.Time lastPollError map[string]time.Time + + // stop terminates the poll loop independently of the parent context. + // Without it the poller would outlive a completed scan, because the + // runner context is only cancelled on an explicit Stop. + stop chan struct{} + stopOnce sync.Once + // done is closed once the poll loop has exited, so callers can wait for + // the goroutine to be gone instead of merely signalled. + done chan struct{} } type OOBHitSnapshot struct { @@ -84,11 +93,23 @@ func NewOOBManager(ctx context.Context, adapter *oobadapter.OOBAdapter, pollInte maxSeen: 200, lastPolledAt: make(map[string]time.Time), lastPollError: make(map[string]time.Time), + stop: make(chan struct{}), + done: make(chan struct{}), } go m.loop(ctx) return m } +// Stop terminates the poll loop and waits for it to exit. +// It is safe to call more than once. +func (m *OOBManager) Stop() { + if m == nil { + return + } + m.stopOnce.Do(func() { close(m.stop) }) + <-m.done +} + func (m *OOBManager) Watch(filter string, filterType string) { if m == nil || m.adapter == nil || strings.TrimSpace(filter) == "" { return @@ -318,6 +339,8 @@ func waitClosed(ch <-chan struct{}, timeout time.Duration) bool { } func (m *OOBManager) loop(ctx context.Context) { + defer close(m.done) + ticker := time.NewTicker(m.pollInterval) defer ticker.Stop() @@ -325,6 +348,8 @@ func (m *OOBManager) loop(ctx context.Context) { select { case <-ctx.Done(): return + case <-m.stop: + return case <-ticker.C: } diff --git a/pkg/runner/oob_resolver.go b/pkg/runner/oob_resolver.go index 5ee2d8c35..035e4c048 100644 --- a/pkg/runner/oob_resolver.go +++ b/pkg/runner/oob_resolver.go @@ -25,7 +25,7 @@ func (r *Runner) registerOOBPendings(res *result.Result, pendings []OOBPending) if r == nil || res == nil || len(pendings) == 0 { return } - if r.engine == nil || r.engine.oobMgr == nil { + if r.engine == nil || r.engine.OOBMgr() == nil { return } for _, p := range pendings { @@ -77,12 +77,12 @@ func (r *Runner) registerOOBPendings(res *result.Result, pendings []OOBPending) } r.oobPendingMu.Unlock() - r.engine.oobMgr.Watch(filter, filterType) + r.engine.OOBMgr().Watch(filter, filterType) } } func (r *Runner) startOOBResolver() { - if r == nil || r.engine == nil || r.engine.oobMgr == nil { + if r == nil || r.engine == nil || r.engine.OOBMgr() == nil { return } if r.oobResolverStop != nil || r.oobResolverDone != nil { @@ -127,7 +127,7 @@ func (r *Runner) stopOOBResolver() { } func (r *Runner) resolveOOBPendingsOnce() int { - if r == nil || r.engine == nil || r.engine.oobMgr == nil || r.OnResult == nil { + if r == nil || r.engine == nil || r.engine.OOBMgr() == nil || r.OnResult == nil { return 0 } @@ -150,15 +150,15 @@ func (r *Runner) resolveOOBPendingsOnce() int { resolved := 0 for _, it := range items { ent := it.ent - _, ok := r.engine.oobMgr.HitSnapshot(ent.filter, ent.filterType) + _, ok := r.engine.OOBMgr().HitSnapshot(ent.filter, ent.filterType) if !ok { continue } - if ent.token != "" && !r.engine.oobMgr.TokenMatches(ent.filter, ent.filterType, ent.token) { + if ent.token != "" && !r.engine.OOBMgr().TokenMatches(ent.filter, ent.filterType, ent.token) { continue } - ev := r.engine.oobMgr.Evidence(ent.filter, ent.filterType, 5) + ev := r.engine.OOBMgr().Evidence(ent.filter, ent.filterType, 5) if strings.TrimSpace(ev) == "" { continue } @@ -241,7 +241,7 @@ func (r *Runner) finalizeOOBPendings() { if r == nil { return } - if r.engine == nil || r.engine.oobMgr == nil { + if r.engine == nil || r.engine.OOBMgr() == nil { if r.options != nil && r.options.OnPhaseProgress != nil { r.options.OnPhaseProgress("oob_finalize", "skipped", 0, 0, 100) } diff --git a/pkg/runner/oob_resolver_test.go b/pkg/runner/oob_resolver_test.go index 87851f591..d18285414 100644 --- a/pkg/runner/oob_resolver_test.go +++ b/pkg/runner/oob_resolver_test.go @@ -37,10 +37,10 @@ func TestResolveOOBPendingsOncePreservesRequestResponse(t *testing.T) { } var got *result.Result + engine := &Engine{} + engine.oobMgr.Store(mgr) r := &Runner{ - engine: &Engine{ - oobMgr: mgr, - }, + engine: engine, OnResult: func(rst *result.Result) { got = rst }, diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index ea169d1a5..7c0e6596c 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -27,13 +27,16 @@ import ( type OnResult func(*result.Result) type Runner struct { - options *config.Options - catalog *catalog.Catalog - Report *report.Report - JsonReport *report.JsonReport - OnResult OnResult - OnFingerprint func(targetKey string, hits []fingerprint.Hit) - OnWebProbe func(meta WebMeta) + options *config.Options + catalog *catalog.Catalog + Report *report.Report + JsonReport *report.JsonReport + OnResult OnResult + OnFingerprint func(targetKey string, hits []fingerprint.Hit) + OnWebProbe func(meta WebMeta) + // OnFailure 在单个 PoC 执行失败时触发（请求错误、表达式异常、panic）。 + // 以前这些失败会被完全吞掉，调用方无从得知。 + OnFailure func(target string, pocID string, err error) PocsYaml utils.StringSlice PocsEmbedYaml utils.StringSlice engine *Engine @@ -134,6 +137,7 @@ func NewRunner(options *config.Options) (*Runner, error) { if err != nil { return nil, err } + cyberspace.Quiet = options.SDKMode || options.Silent runner.Cyberspace = cyberspace } @@ -300,6 +304,24 @@ func (r *Runner) Stop() { } } +// Release stops the runner and frees every resource it owns: the scan engine, +// the OOB poll loop, the OOB resolver and the progress store. +// +// Stop only signals the scan to halt; Release additionally guarantees that no +// background goroutine survives. Embedders must call it, otherwise a +// long-running host process accumulates one poller per completed scan. +func (r *Runner) Release() { + if r == nil { + return + } + r.Stop() + r.stopOOBResolver() + if r.engine != nil { + r.engine.stopOOBManager() + r.engine.pedmStopMonitor() + } +} + func (r *Runner) LiveStatsSuffix() string { if r == nil || r.engine == nil { return "" diff --git a/pkg/sdk/doc.go b/pkg/sdk/doc.go new file mode 100644 index 000000000..2d75d3721 --- /dev/null +++ b/pkg/sdk/doc.go @@ -0,0 +1,109 @@ +/* +Package sdk is the Go API of the afrog vulnerability scanner. + +# Getting started + +A scan is built from functional options and driven by a context: + + scanner, err := sdk.New(ctx, + sdk.WithTargets("https://example.com"), + sdk.WithPocPaths("./pocs/afrog-pocs"), + ) + if err != nil { + return err + } + defer scanner.Close() + + if err := scanner.Execute(ctx); err != nil { + return err + } + for _, r := range scanner.Results() { + fmt.Println(r.PocID, r.FullTarget) + } + +# Synchronous and asynchronous execution + +Execute runs the scan and returns when it finishes. For asynchronous use, Start +returns immediately and completion is observed with Wait or Done: + + if err := scanner.Start(ctx); err != nil { + return err + } + for { + select { + case <-ticker.C: + log.Printf("%.1f%%", scanner.Progress()) + case <-scanner.Done(): + return scanner.Err() + } + } + +A Scanner is single-use. Once a scan finishes, Start and Execute return +[ErrAlreadyFinished]; create a new Scanner to scan again. + +# PoC input + +WithPocPaths accepts a single file, a directory searched recursively, or a glob +pattern, and may be repeated: + + sdk.WithPocPaths("a.yaml", "./pocs", "./extra/*.yaml") + +Paths are merged with the built-in PoCs. WithPocPathsOnly restricts the scan to +the listed paths. Pocs reports what was loaded and PocDiagnostics reports what +was skipped and why, so a malformed or deprecated PoC is visible to the caller +instead of only being printed to a console. + +# Results + +Results returns [Result] values. Each carries the complete request and response +of every step in [Exchange], as readable strings rather than protobuf byte +slices, so the whole structure can be passed straight to encoding/json. + +Use WithRequestResponse(false) to drop the raw exchanges and WithMaxStoredResults +to cap accumulation on large scans. Neither affects handlers or streams, which +always observe every finding. + +# Handlers and streams + +Results can be consumed in two ways. Handlers are callbacks registered at +construction time: + + sdk.WithResultHandler(func(r sdk.Result) { ... }) + sdk.WithFailureHandler(func(f sdk.Failure) { ... }) + +Handlers are invoked concurrently from scan workers, so they must synchronise +their own state. + +Streams are channels obtained from ResultStream, PortStream, HostStream, +WebProbeStream, ProgressStream and ScanInfoStream. A stream publishes nothing +until it is first subscribed to, so an unused stream costs nothing and cannot +stall the scan. Once subscribed it must be consumed: sends block when the +buffer fills, because silently discarding a finding is not an acceptable +failure mode for a scanner. Every stream is closed when the scan finishes, so +a range loop always terminates. + +# Errors + +Failures are reported with sentinel errors and matched using errors.Is: + + if errors.Is(err, sdk.ErrPocPathNotFound) { ... } + +# Output + +The SDK writes nothing to stdout or stderr. Use Info to obtain the scan summary +as a struct, or WithVerbose to opt into printing it. + +# Resource management + +Close stops the scan, waits for the scan goroutine to exit and releases every +background goroutine, including the out-of-band poll loop. Always defer it. + +# Concurrency + +A Scanner is safe for concurrent use. Running several scanners concurrently in +one process is not: the HTTP client, the rate limiter and the protocol probe +cache are process-global, so concurrent scanners overwrite each other's proxy, +timeout and rate-limit settings. Scan batches sequentially, or isolate them in +separate processes. +*/ +package sdk diff --git a/pkg/sdk/errors.go b/pkg/sdk/errors.go new file mode 100644 index 000000000..d2597d45e --- /dev/null +++ b/pkg/sdk/errors.go @@ -0,0 +1,50 @@ +package sdk + +import "errors" + +// Sentinel errors returned by the SDK. Callers should match them with +// errors.Is rather than comparing error strings. +var ( + // ErrNoTargets is returned when no scan target was configured. + ErrNoTargets = errors.New("afrog/sdk: no targets available") + + // ErrNoPocs is returned when no PoC is executable, either because the + // configured paths matched nothing or because every PoC was filtered out. + ErrNoPocs = errors.New("afrog/sdk: no pocs available") + + // ErrPocPathNotFound is returned when a PoC path resolves to no + // .yaml/.yml file. + ErrPocPathNotFound = errors.New("afrog/sdk: poc path not found") + + // ErrAlreadyRunning is returned when a scan is already in progress. + ErrAlreadyRunning = errors.New("afrog/sdk: scan is already running") + + // ErrAlreadyFinished is returned when a finished scanner is reused. A + // scanner is single-use; create a new one to scan again. + ErrAlreadyFinished = errors.New("afrog/sdk: scan has already finished") + + // ErrClosed is returned when the scanner has been closed. + ErrClosed = errors.New("afrog/sdk: scanner is closed") + + // ErrNotStarted is returned by Wait when the scan was never started. + ErrNotStarted = errors.New("afrog/sdk: scan has not been started") + + // ErrInvalidOptions is returned when the option combination is invalid. + ErrInvalidOptions = errors.New("afrog/sdk: invalid options") + + // ErrWebhookTokenRequired is returned when a webhook notifier is enabled + // without a token. + ErrWebhookTokenRequired = errors.New("afrog/sdk: webhook token is required") +) + +// CuratedMountError reports that the optional curated PoC source could not be +// mounted. It is never fatal: the scan proceeds without that source. +type CuratedMountError struct { + Err error +} + +func (e *CuratedMountError) Error() string { + return "afrog/sdk: curated mount failed: " + e.Err.Error() +} + +func (e *CuratedMountError) Unwrap() error { return e.Err } diff --git a/pkg/sdk/event.go b/pkg/sdk/event.go new file mode 100644 index 000000000..88bb63e09 --- /dev/null +++ b/pkg/sdk/event.go @@ -0,0 +1,388 @@ +package sdk + +import ( + "strings" + "time" + + "github.com/zan8in/afrog/v3/pkg/fingerprint" + "github.com/zan8in/afrog/v3/pkg/proto" + "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/pkg/utils" +) + +// Result is a single finding. +// +// It is the SDK's stable, JSON-serialisable view of a scan result. Raw request +// and response messages are plain strings rather than protobuf []byte fields, +// so encoding/json emits readable text instead of base64. +type Result struct { + PocID string `json:"poc_id"` + PocName string `json:"poc_name,omitempty"` + Severity string `json:"severity,omitempty"` + Author string `json:"author,omitempty"` + Description string `json:"description,omitempty"` + Reference []string `json:"reference,omitempty"` + Tags []string `json:"tags,omitempty"` + + CveID string `json:"cve_id,omitempty"` + CweID string `json:"cwe_id,omitempty"` + CvssScore float64 `json:"cvss_score,omitempty"` + CvssMetrics string `json:"cvss_metrics,omitempty"` + + // Target is the seed target; FullTarget is the URL that actually matched. + Target string `json:"target"` + FullTarget string `json:"full_target,omitempty"` + + // Extractors holds the key/value pairs produced by the PoC's extractors. + Extractors map[string]string `json:"extractors,omitempty"` + + // Fingerprints holds matched fingerprints, for fingerprint results. + Fingerprints []Fingerprint `json:"fingerprints,omitempty"` + + // Exchanges holds the request/response round trips in execution order. A + // multi-step PoC produces several. It is empty when request/response + // capture is disabled via WithRequestResponse(false). + Exchanges []Exchange `json:"exchanges,omitempty"` + + FoundAt time.Time `json:"found_at"` +} + +// Fingerprint is a matched fingerprint. +type Fingerprint struct { + ID string `json:"id"` + Name string `json:"name,omitempty"` + Severity string `json:"severity,omitempty"` + Tags string `json:"tags,omitempty"` +} + +// Exchange is one request/response round trip. +type Exchange struct { + // Request is the raw request message, including the request line, headers + // and body. + Request string `json:"request,omitempty"` + // Response is the raw response message, including the status line, headers + // and body. + Response string `json:"response,omitempty"` + + Method string `json:"method,omitempty"` + URL string `json:"url,omitempty"` + RequestHeaders map[string]string `json:"request_headers,omitempty"` + RequestBody string `json:"request_body,omitempty"` + StatusCode int `json:"status_code,omitempty"` + ResponseHeaders map[string]string `json:"response_headers,omitempty"` + ResponseBody string `json:"response_body,omitempty"` + ContentType string `json:"content_type,omitempty"` + LatencyMs int64 `json:"latency_ms,omitempty"` + + // Matched reports whether this step matched. + Matched bool `json:"matched"` + + // BodyTruncated reports that ResponseBody was cut at the MaxRespBodySize + // limit and is therefore not the complete server response. + BodyTruncated bool `json:"body_truncated,omitempty"` + + // BruteTruncated reports that a brute-force PoC stopped early because it + // reached BruteMaxRequests. + BruteTruncated bool `json:"brute_truncated,omitempty"` + // BruteRequests is the number of brute-force requests actually sent. + BruteRequests int `json:"brute_requests,omitempty"` +} + +// PortEvent reports an open port found during port pre-scanning. +type PortEvent struct { + Host string `json:"host"` + Port int `json:"port"` +} + +// HostEvent reports a live host found during host discovery. +type HostEvent struct { + Host string `json:"host"` +} + +// WebProbeEvent reports the metadata of a probed web service. +type WebProbeEvent struct { + URL string `json:"url"` + Title string `json:"title,omitempty"` + Server string `json:"server,omitempty"` + PoweredBy string `json:"powered_by,omitempty"` +} + +// Phase names reported through PhaseProgress. +const ( + PhaseHostDiscovery = "host_discovery" + PhasePortScan = "portscan" + PhaseWebProbe = "webprobe" + PhaseVuln = "vuln" +) + +// PhaseProgress reports the progress of one scan phase. +type PhaseProgress struct { + // Phase is one of the Phase* constants. + Phase string `json:"phase"` + // Status is "running", "completed" or "interrupted". + Status string `json:"status"` + Finished int64 `json:"finished"` + Total int64 `json:"total"` + Percent int `json:"percent"` +} + +// ScanInfo summarises the scan. +type ScanInfo struct { + TotalTargets int `json:"total_targets"` + TotalPocs int `json:"total_pocs"` + TotalScans int `json:"total_scans"` + Targets []string `json:"targets,omitempty"` + OOBEnabled bool `json:"oob_enabled"` + OOBStatus string `json:"oob_status,omitempty"` +} + +// Failure reports that a single PoC execution failed. A failure never aborts +// the scan; it is surfaced so that callers can observe request errors, +// expression errors and recovered panics instead of losing them silently. +type Failure struct { + Target string `json:"target"` + PocID string `json:"poc_id"` + Err error `json:"-"` +} + +func (f Failure) Error() string { + if f.Err == nil { + return "" + } + return f.Err.Error() +} + +func (f Failure) Unwrap() error { return f.Err } + +// Stats holds scan counters. Snapshots are returned by Scanner.Stats. +type Stats struct { + StartTime time.Time `json:"start_time"` + EndTime time.Time `json:"end_time"` + TotalTargets int `json:"total_targets"` + TotalPocs int `json:"total_pocs"` + TotalScans int `json:"total_scans"` + CompletedScans int64 `json:"completed_scans"` + FoundVulns int64 `json:"found_vulns"` +} + +// Duration returns the scan duration, measured to now while still running. +func (s Stats) Duration() time.Duration { + if s.EndTime.IsZero() { + return time.Since(s.StartTime) + } + return s.EndTime.Sub(s.StartTime) +} + +// DefaultRedactedHeaders lists the headers that WithRedactedHeaders masks when +// called without arguments. They are the ones that routinely carry +// credentials. +var DefaultRedactedHeaders = []string{ + "authorization", + "proxy-authorization", + "cookie", + "set-cookie", + "x-api-key", + "x-auth-token", + "x-csrf-token", +} + +// redactedValue replaces a masked header value. +const redactedValue = "[REDACTED]" + +// redactExchange masks credential-bearing headers in both the structured maps +// and the raw messages. +// +// The raw request and response are the whole point of Exchange, so redaction +// is opt-in: enabling it trades debuggability for the guarantee that results +// can be logged or persisted without leaking credentials. +func redactExchange(ex *Exchange, names map[string]struct{}) { + if len(names) == 0 { + return + } + redactHeaderMap(ex.RequestHeaders, names) + redactHeaderMap(ex.ResponseHeaders, names) + ex.Request = redactRawMessage(ex.Request, names) + ex.Response = redactRawMessage(ex.Response, names) +} + +func redactHeaderMap(headers map[string]string, names map[string]struct{}) { + for k := range headers { + if _, ok := names[strings.ToLower(k)]; ok { + headers[k] = redactedValue + } + } +} + +// redactRawMessage masks header lines in a raw HTTP message, stopping at the +// blank line that separates headers from the body. +func redactRawMessage(raw string, names map[string]struct{}) string { + if raw == "" { + return raw + } + lines := strings.Split(raw, "\n") + for i, line := range lines { + trimmed := strings.TrimRight(line, "\r") + if trimmed == "" { + break // end of headers + } + colon := strings.IndexByte(trimmed, ':') + if colon <= 0 { + continue + } + if _, ok := names[strings.ToLower(strings.TrimSpace(trimmed[:colon]))]; !ok { + continue + } + suffix := "" + if strings.HasSuffix(line, "\r") { + suffix = "\r" + } + lines[i] = trimmed[:colon] + ": " + redactedValue + suffix + } + return strings.Join(lines, "\n") +} + +// newResult converts an internal result into the SDK view. When includeRR is +// false the raw exchanges are dropped, which keeps memory bounded on large +// scans. +func newResult(r *result.Result, includeRR bool, foundAt time.Time) Result { + out := Result{ + Target: r.Target, + FullTarget: r.FullTarget, + FoundAt: foundAt, + } + if strings.TrimSpace(out.FullTarget) == "" { + out.FullTarget = out.Target + } + + if pi := r.PocInfo; pi != nil { + out.PocID = pi.Id + out.PocName = pi.Info.Name + out.Severity = pi.Info.Severity + out.Author = pi.Info.Author + out.Description = pi.Info.Description + out.Reference = append([]string(nil), pi.Info.Reference...) + out.Tags = splitAndTrim(pi.Info.Tags) + out.CveID = pi.Info.Classification.CveId + out.CweID = pi.Info.Classification.CweId + out.CvssScore = pi.Info.Classification.CvssScore + out.CvssMetrics = pi.Info.Classification.CvssMetrics + } + + if len(r.Extractor) > 0 { + out.Extractors = make(map[string]string, len(r.Extractor)) + for _, item := range r.Extractor { + key, ok := item.Key.(string) + if !ok { + continue + } + if v, ok := item.Value.(string); ok { + out.Extractors[key] = utils.Str2UTF8(v) + } + } + } + + if hits, ok := r.FingerResult.([]fingerprint.Hit); ok { + for _, h := range hits { + out.Fingerprints = append(out.Fingerprints, Fingerprint{ + ID: h.ID, + Name: h.Name, + Severity: h.Severity, + Tags: h.Tags, + }) + } + } + + if includeRR { + for _, pr := range r.AllPocResult { + if pr == nil { + continue + } + out.Exchanges = append(out.Exchanges, newExchange(pr)) + } + } + + return out +} + +// newResultRedacted is newResult with credential-bearing headers masked. +func newResultRedacted(r *result.Result, includeRR bool, foundAt time.Time, redact map[string]struct{}) Result { + out := newResult(r, includeRR, foundAt) + for i := range out.Exchanges { + redactExchange(&out.Exchanges[i], redact) + } + return out +} + +func newExchange(pr *result.PocResult) Exchange { + ex := Exchange{ + Matched: pr.IsVul, + BodyTruncated: pr.BodyTruncated, + BruteTruncated: pr.BruteTruncated, + BruteRequests: pr.BruteRequests, + } + + if req := pr.ResultRequest; req != nil { + ex.Request = utils.Str2UTF8(string(req.GetRaw())) + ex.Method = req.GetMethod() + ex.RequestHeaders = copyStringMap(req.GetHeaders()) + ex.RequestBody = utils.Str2UTF8(string(req.GetBody())) + ex.URL = protoURL(req.GetUrl()) + } + + if resp := pr.ResultResponse; resp != nil { + ex.Response = utils.Str2UTF8(string(resp.GetRaw())) + ex.StatusCode = int(resp.GetStatus()) + ex.ResponseHeaders = copyStringMap(resp.GetHeaders()) + ex.ResponseBody = utils.Str2UTF8(string(resp.GetBody())) + ex.ContentType = resp.GetContentType() + ex.LatencyMs = resp.GetLatency() + if ex.URL == "" { + ex.URL = protoURL(resp.GetUrl()) + } + } + + return ex +} + +func protoURL(u *proto.UrlType) string { + if u == nil { + return "" + } + scheme, host := u.GetScheme(), u.GetHost() + if scheme == "" && host == "" { + return "" + } + out := scheme + "://" + host + u.GetPath() + if q := u.GetQuery(); q != "" { + out += "?" + q + } + if f := u.GetFragment(); f != "" { + out += "#" + f + } + return out +} + +func copyStringMap(in map[string]string) map[string]string { + if len(in) == 0 { + return nil + } + out := make(map[string]string, len(in)) + for k, v := range in { + out[k] = v + } + return out +} + +func splitAndTrim(s string) []string { + if strings.TrimSpace(s) == "" { + return nil + } + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + if v := strings.TrimSpace(p); v != "" { + out = append(out, v) + } + } + return out +} diff --git a/pkg/sdk/event_test.go b/pkg/sdk/event_test.go new file mode 100644 index 000000000..5ec729229 --- /dev/null +++ b/pkg/sdk/event_test.go @@ -0,0 +1,84 @@ +package sdk + +import ( + "strings" + "testing" +) + +func TestRedactExchange(t *testing.T) { + names := toSet(DefaultRedactedHeaders) + + ex := Exchange{ + Request: "GET /admin HTTP/1.1\r\n" + + "Host: example.com\r\n" + + "Authorization: Bearer super-secret-token\r\n" + + "Cookie: session=abc123\r\n" + + "User-Agent: afrog\r\n" + + "\r\n" + + "Authorization: this-is-body-not-a-header\r\n", + Response: "HTTP/1.1 200 OK\r\n" + + "Set-Cookie: session=xyz789; HttpOnly\r\n" + + "Content-Type: text/html\r\n" + + "\r\n" + + "ok", + RequestHeaders: map[string]string{"authorization": "Bearer super-secret-token", "user-agent": "afrog"}, + ResponseHeaders: map[string]string{"set-cookie": "session=xyz789", "content-type": "text/html"}, + } + + redactExchange(&ex, names) + + for _, secret := range []string{"super-secret-token", "abc123", "xyz789"} { + if strings.Contains(ex.Request, secret) { + t.Errorf("request still contains %q:\n%s", secret, ex.Request) + } + if strings.Contains(ex.Response, secret) { + t.Errorf("response still contains %q:\n%s", secret, ex.Response) + } + } + if ex.RequestHeaders["authorization"] != redactedValue { + t.Errorf("request header not redacted: %q", ex.RequestHeaders["authorization"]) + } + if ex.ResponseHeaders["set-cookie"] != redactedValue { + t.Errorf("response header not redacted: %q", ex.ResponseHeaders["set-cookie"]) + } + + // Non-sensitive data must survive untouched. + if !strings.Contains(ex.Request, "User-Agent: afrog") { + t.Errorf("redaction removed a harmless header:\n%s", ex.Request) + } + if !strings.Contains(ex.Response, "ok") { + t.Errorf("redaction damaged the response body:\n%s", ex.Response) + } + // A header-looking line inside the body must not be rewritten, because + // redaction stops at the blank line that ends the headers. + if !strings.Contains(ex.Request, "this-is-body-not-a-header") { + t.Errorf("redaction leaked past the header section:\n%s", ex.Request) + } +} + +func TestRedactExchange_NoNamesIsNoOp(t *testing.T) { + original := Exchange{ + Request: "GET / HTTP/1.1\r\nAuthorization: keep-me\r\n\r\n", + RequestHeaders: map[string]string{"authorization": "keep-me"}, + } + ex := original + redactExchange(&ex, nil) + + if ex.Request != original.Request { + t.Errorf("request changed with no redaction configured:\n%s", ex.Request) + } + if ex.RequestHeaders["authorization"] != "keep-me" { + t.Errorf("header changed with no redaction configured: %q", ex.RequestHeaders["authorization"]) + } +} + +func TestStats_Duration(t *testing.T) { + var s Stats + s.StartTime = s.StartTime.Add(0) + + // A running scan measures up to now, so the duration must be non-negative + // rather than the huge negative value a zero EndTime would produce. + if d := s.Duration(); d < 0 { + t.Errorf("Duration() = %v, want >= 0 while the scan is running", d) + } +} diff --git a/pkg/sdk/lifecycle_test.go b/pkg/sdk/lifecycle_test.go new file mode 100644 index 000000000..954fad09d --- /dev/null +++ b/pkg/sdk/lifecycle_test.go @@ -0,0 +1,158 @@ +package sdk + +import ( + "context" + "errors" + "os" + "sync" + "testing" + "time" +) + +// Close must be safe at every lifecycle point, including before a scan was +// ever started and while one is still running. +func TestScanner_CloseAtEveryLifecyclePoint(t *testing.T) { + tests := []struct { + name string + before func(s *Scanner) + }{ + {name: "before start", before: func(*Scanner) {}}, + {name: "after stop without start", before: func(s *Scanner) { s.Stop() }}, + {name: "after execute", before: func(s *Scanner) { _ = s.Execute(context.Background()) }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + scanner, _ := newTestScanner(t) + tt.before(scanner) + + done := make(chan error, 1) + go func() { done <- scanner.Close() }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Close: %v", err) + } + case <-time.After(30 * time.Second): + t.Fatal("Close blocked") + } + }) + } +} + +// Close called while a scan is running must stop it and return promptly. +func TestScanner_CloseWhileRunning(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + + done := make(chan error, 1) + go func() { done <- scanner.Close() }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Close: %v", err) + } + case <-time.After(60 * time.Second): + t.Fatal("Close blocked while a scan was running") + } +} + +// Concurrent Stop and Close calls must not panic or deadlock. +func TestScanner_ConcurrentStopAndClose(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(2) + go func() { defer wg.Done(); scanner.Stop() }() + go func() { defer wg.Done(); _ = scanner.Close() }() + } + + done := make(chan struct{}) + go func() { defer close(done); wg.Wait() }() + + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("concurrent Stop/Close deadlocked") + } +} + +// A scan that finishes cleanly must report 100% progress even though the +// pre-execution task estimate rarely matches the exact number of tasks run. +func TestScanner_ProgressReaches100OnCleanFinish(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + if got := scanner.Progress(); got != 100 { + t.Fatalf("Progress() = %.2f after a clean finish, want 100", got) + } +} + +// A New that fails after the runner was already built must not leak the +// runner's background goroutines, and must put the curated environment +// variables back the way it found them. The caller gets no Scanner on that +// path, so it has no way to clean up itself. +func TestNew_ReleasesResourcesWhenConstructionFails(t *testing.T) { + t.Setenv(envCuratedDisabled, "sentinel") + + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "info.yaml", "abandon-check") + + // Warm up the lazily-initialised shared state so it is not counted below. + failedNew := func() error { + _, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + // The PoC is "info", so nothing survives the filter and + // construction fails only after the runner exists. + WithSeverity("critical"), + ) + return err + } + if err := failedNew(); !errors.Is(err, ErrNoPocs) { + t.Fatalf("New error = %v, want ErrNoPocs", err) + } + + before := waitGoroutines(t, 0) + for i := 0; i < 3; i++ { + if err := failedNew(); !errors.Is(err, ErrNoPocs) { + t.Fatalf("New error = %v, want ErrNoPocs", err) + } + } + if after := waitGoroutines(t, before); after > before+5 { + t.Fatalf("goroutine count grew from %d to %d across 3 failed constructions", before, after) + } + + if got := os.Getenv(envCuratedDisabled); got != "sentinel" { + t.Fatalf("%s = %q after a failed New, want the original %q", envCuratedDisabled, got, "sentinel") + } +} + +// Double Start must be rejected rather than launching two scans. +func TestScanner_DoubleStart(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Start(context.Background()); err != nil { + t.Fatalf("first Start: %v", err) + } + err := scanner.Start(context.Background()) + if err == nil { + t.Fatal("second Start should have failed") + } + <-scanner.Done() +} diff --git a/pkg/sdk/monitor_test.go b/pkg/sdk/monitor_test.go new file mode 100644 index 000000000..21eea4033 --- /dev/null +++ b/pkg/sdk/monitor_test.go @@ -0,0 +1,216 @@ +package sdk + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "path/filepath" + "sync" + "testing" + "time" +) + +// --- execution monitor ------------------------------------------------------ + +// The monitor's reports must reach the registered handler. +func TestScanner_ExecutionMonitorReportsThroughHandler(t *testing.T) { + var mu sync.Mutex + var lines []string + + scanner, _ := newTestScanner(t, + WithExecutionMonitor(ExecutionMonitorOptions{LogLimit: 10}), + WithMonitorHandler(func(line string) { + mu.Lock() + lines = append(lines, line) + mu.Unlock() + }), + ) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + mu.Lock() + got := len(lines) + mu.Unlock() + if got == 0 { + t.Fatal("execution monitor produced no reports") + } +} + +// The engine prints monitor output itself when no hook is installed. The SDK +// must keep the console clean whether or not a handler is registered. +func TestScanner_ExecutionMonitorStaysSilent(t *testing.T) { + tests := []struct { + name string + extra []Option + }{ + {name: "without handler", extra: nil}, + {name: "with handler", extra: []Option{WithMonitorHandler(func(string) {})}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + options := append([]Option{ + WithExecutionMonitor(ExecutionMonitorOptions{LogLimit: 10}), + }, tt.extra...) + scanner, _ := newTestScanner(t, options...) + + stdout, r, w := captureStdout(t) + execErr := scanner.Execute(context.Background()) + output := restoreStdout(t, stdout, r, w) + + if execErr != nil { + t.Fatalf("Execute: %v", execErr) + } + if output != "" { + t.Fatalf("scanner wrote to stdout:\n%s", output) + } + }) + } +} + +func TestWithExecutionMonitor_KeepsDefaultsAndValidates(t *testing.T) { + o := NewOptions() + if err := WithExecutionMonitor(ExecutionMonitorOptions{LogLimit: 5})(o); err != nil { + t.Fatalf("WithExecutionMonitor: %v", err) + } + if !o.EnableMonitor { + t.Error("WithExecutionMonitor did not enable the monitor") + } + if o.Monitor.SlowThresholdSec != DefaultMonitorSlowThresholdSec { + t.Errorf("SlowThresholdSec = %d, want %d", o.Monitor.SlowThresholdSec, DefaultMonitorSlowThresholdSec) + } + if o.Monitor.SummaryTop != DefaultMonitorSummaryTop { + t.Errorf("SummaryTop = %d, want %d", o.Monitor.SummaryTop, DefaultMonitorSummaryTop) + } + if o.Monitor.SummaryBy != MonitorSummaryByMax { + t.Errorf("SummaryBy = %q, want %q", o.Monitor.SummaryBy, MonitorSummaryByMax) + } + + if err := WithExecutionMonitor(ExecutionMonitorOptions{SummaryBy: "bogus"})(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("unknown summary key = %v, want ErrInvalidOptions", err) + } + if err := WithExecutionMonitor(ExecutionMonitorOptions{LogLimit: -1})(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("negative LogLimit = %v, want ErrInvalidOptions", err) + } +} + +// --- per-task timeout ------------------------------------------------------- + +func TestWithTaskTimeout_KeepsDefaultCapsAndValidates(t *testing.T) { + o := NewOptions() + if err := WithTaskTimeout(TaskTimeoutOptions{HardSec: 30, Smart: true})(o); err != nil { + t.Fatalf("WithTaskTimeout: %v", err) + } + caps := map[string][2]int{ + "VisibleCapSec": {o.TaskTimeout.VisibleCapSec, DefaultTaskTimeoutVisibleCapSec}, + "NetCapSec": {o.TaskTimeout.NetCapSec, DefaultTaskTimeoutNetCapSec}, + "GoCapSec": {o.TaskTimeout.GoCapSec, DefaultTaskTimeoutGoCapSec}, + } + for name, v := range caps { + if v[0] != v[1] { + t.Errorf("%s = %d, want the default %d", name, v[0], v[1]) + } + } + + if err := WithTaskTimeout(TaskTimeoutOptions{HardSec: -1})(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("negative HardSec = %v, want ErrInvalidOptions", err) + } + if err := WithTaskTimeout(TaskTimeoutOptions{NetCapSec: -5})(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("negative NetCapSec = %v, want ErrInvalidOptions", err) + } +} + +// A hard task timeout must actually cut a PoC short. Without it the scan would +// block for the full server delay. +func TestScanner_TaskHardTimeoutCutsSlowPoc(t *testing.T) { + const serverDelay = 10 * time.Second + + release := make(chan struct{}) + t.Cleanup(func() { close(release) }) + + // Only the PoC's own path is slow. The reachability probe afrog runs + // first has its own 5s ceiling, so a uniformly slow server would be cut + // there and the PoC would never execute. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/probe" { + select { + case <-time.After(serverDelay): + case <-r.Context().Done(): + return + case <-release: + } + } + _, _ = w.Write([]byte(magicToken)) + })) + t.Cleanup(srv.Close) + + dir := t.TempDir() + writePoc(t, dir, "slow.yaml", "sdk-test-slow") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + // Well above the server delay, so only the task timeout can end this. + WithTimeout(60), + WithTaskTimeout(TaskTimeoutOptions{HardSec: 1}), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + start := time.Now() + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + elapsed := time.Since(start) + + if elapsed >= serverDelay-2*time.Second { + t.Fatalf("scan took %v, so the 1s task timeout did not cut the %v request", elapsed, serverDelay) + } +} + +// The PoC path list must reach the engine before PoC loading, otherwise smart +// estimation silently produces no per-task ceiling. +func TestScanner_SmartTaskTimeoutReachesEngine(t *testing.T) { + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "smart.yaml", "sdk-test-smart") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithTaskTimeout(TaskTimeoutOptions{Smart: true, VisibleCapSec: 120}), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + if !scanner.internal.TaskSmartTimeout { + t.Error("TaskSmartTimeout did not reach the engine options") + } + if scanner.internal.TaskTimeoutVisibleCapSec != 120 { + t.Errorf("VisibleCapSec = %d, want 120", scanner.internal.TaskTimeoutVisibleCapSec) + } + + pocs := scanner.Pocs() + if len(pocs) == 0 { + t.Fatal("no pocs loaded") + } + if pocs[0].EstimatedTaskTimeoutSec <= 0 { + t.Errorf("EstimatedTaskTimeoutSec = %d, want > 0 with smart timeout on", + pocs[0].EstimatedTaskTimeoutSec) + } + if pocs[0].EstimatedTaskTimeoutSec > 120 { + t.Errorf("EstimatedTaskTimeoutSec = %d, want <= the 120s cap", pocs[0].EstimatedTaskTimeoutSec) + } + _ = filepath.Base(dir) +} diff --git a/pkg/sdk/options.go b/pkg/sdk/options.go new file mode 100644 index 000000000..f09483a53 --- /dev/null +++ b/pkg/sdk/options.go @@ -0,0 +1,1099 @@ +package sdk + +import ( + "fmt" + "strings" + "time" + + "github.com/zan8in/afrog/v3/pkg/result" +) + +// Default values for Options. These constants are the single source of truth +// for the SDK defaults, so documentation and code cannot drift apart. +const ( + DefaultRateLimit = 150 + DefaultConcurrency = 25 + DefaultRetries = 1 + DefaultTimeout = 50 // seconds + DefaultMaxHostError = 3 + DefaultMaxRespBodySize = 2 // MB + DefaultBruteMaxRequests = 5000 + DefaultStreamBuffer = 256 + DefaultOOBRateLimit = 25 + DefaultOOBConcurrency = 25 + + // DefaultOOBPollInterval and DefaultOOBHitRetention mirror the CLI + // defaults. Leaving them at zero would make the engine fall back to a + // one-second poll, hitting the out-of-band provider twice as often as the + // command line does. + DefaultOOBPollInterval = 2 // seconds + DefaultOOBHitRetention = 10 // minutes + + // Smart task-timeout estimation caps, mirroring the CLI defaults. They + // bound the timeout the engine derives from a PoC's content. + DefaultTaskTimeoutVisibleCapSec = 300 // plain HTTP PoCs + DefaultTaskTimeoutNetCapSec = 360 // tcp/udp/ssl PoCs + DefaultTaskTimeoutGoCapSec = 420 // go PoCs + + // Execution monitor defaults, mirroring the CLI. + DefaultMonitorSlowThresholdSec = 30 + DefaultMonitorSlowLogLimit = 20 + DefaultMonitorSummaryTop = 10 +) + +// DefaultCheckpointSaveInterval matches the CLI's auto-save cadence. +const DefaultCheckpointSaveInterval = 10 * time.Second + +// Execution monitor summary sort keys. +const ( + MonitorSummaryByMax = "max" + MonitorSummaryByAvg = "avg" +) + +// FingerprintFilterMode values. +const ( + FingerprintStrict = "strict" + FingerprintOpportunistic = "opportunistic" +) + +// Options is the complete scanner configuration. +// +// Build it with NewOptions and the With* functions rather than filling the +// struct by hand: the option functions validate their input. +type Options struct { + // --- targets --- + + Targets []string + TargetsFile string + + // Cyberspace sources targets from an internet-wide search engine instead + // of, or in addition to, the explicit target list. + Cyberspace CyberspaceOptions + + // TargetPreProbe probes every target's protocol and liveness in parallel + // alongside the scan. + TargetPreProbe bool + + // --- pocs --- + + // PocPaths accepts single files, directories (searched recursively) and + // glob patterns such as "dir/*.yaml". Paths are merged with the built-in + // PoCs unless PocPathsOnly is set. + PocPaths []string + // PocPathsOnly restricts the scan to PocPaths, hiding the built-in, + // curated, my and local PoC sources. + PocPathsOnly bool + + Search string + Severity string + ExcludePocs []string + ExcludePocsFile string + + // --- performance --- + + RateLimit int + Concurrency int + Retries int + Timeout int + MaxHostError int + MaxRespBodySize int + BruteMaxRequests int + ReqLimitPerTarget int + AutoReqLimit bool + Polite bool + Balanced bool + Aggressive bool + Smart bool + + // --- fingerprinting and probing --- + + DisableFingerprint bool + EnableWebProbe bool + FingerprintFilterMode string + DefaultAccept bool + + // StopOnFirstMatch stops the scan as soon as a vulnerability is found. + StopOnFirstMatch bool + + // --- per-task timeout --- + + TaskTimeout TaskTimeoutOptions + + // --- execution monitor --- + + Monitor ExecutionMonitorOptions + EnableMonitor bool + + // --- network --- + + Proxy string + Headers []string + + // --- port pre-scan --- + + PortScan PortScanOptions + EnablePortSan bool + + // --- OOB --- + + OOB OOBOptions + + // --- output --- + + // IncludeRequestResponse controls whether results carry the raw request + // and response messages. Enabled by default; disable it on large scans to + // keep memory bounded. + IncludeRequestResponse bool + // MaxStoredResults caps how many results accumulate in memory. Zero means + // unlimited. Exceeding the cap does not suppress handlers or streams. + MaxStoredResults int + // StreamBuffer is the capacity of each subscribed stream channel. + StreamBuffer int + // RedactedHeaders holds the lower-case header names masked in results. + // Empty means no redaction. + RedactedHeaders []string + // Verbose prints a scan summary to stdout before the scan starts. The SDK + // produces no console output otherwise. + Verbose bool + + // --- notifications --- + + Dingtalk bool + Wecom bool + + // --- resume --- + + Checkpoint CheckpointOptions + + // --- curated poc source --- + + Curated CuratedOptions + + // --- handlers --- + + handlers handlers +} + +// PortScanOptions configures the port pre-scan stage. +type PortScanOptions struct { + // Ports accepts "top", "full", "all", "80,443" or "1-1024". + Ports string + RateLimit int + TimeoutMs int + Retries int + SkipDiscovery bool + ChunkSize int +} + +// CyberspaceEngine values accepted by CyberspaceOptions.Engine. +const ( + // CyberspaceZoomEye is currently the only implemented engine. Its API key + // is read from the afrog configuration file (cyberspace.zoom_eyes). + CyberspaceZoomEye = "zoomeye" +) + +// CyberspaceOptions sources scan targets from an internet-wide search engine. +// +// Targets found this way are added to any explicitly configured ones, so a +// scan may be driven entirely by a search query. +type CyberspaceOptions struct { + // Engine is currently only CyberspaceZoomEye. + Engine string + // Query is the engine's own search syntax, for example `app:"tomcat"`. + Query string + // Count caps how many results are turned into targets. + Count int +} + +// CheckpointOptions makes a scan resumable, the SDK equivalent of the CLI's +// -resume. +// +// Path is read when the scanner starts, so already-finished target/PoC pairs +// are skipped, and rewritten periodically while the scan runs. Resuming only +// works when the target and PoC sets are unchanged, because progress is keyed +// by PoC id and target. +// +// This is unrelated to Scanner.Resume, which lifts a Pause. +type CheckpointOptions struct { + // Path is the checkpoint file. It is created on the first save. + Path string + // SaveInterval is how often the file is rewritten. Zero uses + // DefaultCheckpointSaveInterval. + SaveInterval time.Duration +} + +// TaskTimeoutOptions bounds how long a single target+PoC task may run. +// +// This is separate from the per-request Timeout: a PoC with many rules can +// keep a worker busy far longer than any single request. +type TaskTimeoutOptions struct { + // HardSec is a fixed ceiling in seconds. Zero disables it. + HardSec int + // Smart derives the ceiling from the PoC's content (rule count, sleeps, + // brute force, payloads). When HardSec is also set, the larger of the two + // wins, so HardSec acts as a floor rather than an override. + Smart bool + + // The Cap fields bound the smart estimate per protocol family. They have + // no effect unless Smart is set. Zero uses the defaults. + VisibleCapSec int + NetCapSec int + GoCapSec int +} + +// ExecutionMonitorOptions configures the PoC execution duration monitor, the +// SDK equivalent of the CLI's -pedm. It reports slow and stuck PoCs while the +// scan runs, and a slowest-PoC summary when it ends. +// +// Output is delivered to the handlers registered with WithMonitorHandler. The +// SDK never prints it. +type ExecutionMonitorOptions struct { + // LogLimit reports the first N started tasks. Zero disables it. + LogLimit int + // SlowThresholdSec is the number of seconds after which a task counts as + // slow. Zero disables slow reporting entirely, including the background + // monitor goroutine. + SlowThresholdSec int + // SlowLogLimit caps how many completed slow tasks are reported. Zero + // disables those reports; in-flight slow tasks are still reported. + SlowLogLimit int + // SummaryTop reports the N slowest PoCs when the scan ends. Zero disables + // the summary. + SummaryTop int + // SummaryBy is MonitorSummaryByMax or MonitorSummaryByAvg. + SummaryBy string +} + +// OOBOptions configures out-of-band detection. +type OOBOptions struct { + Enabled bool + // Adapter is one of ceyeio, dnslogcn, alphalog, xray, revsuit. + Adapter string + Key string + Domain string + ApiURL string + HttpURL string + RateLimit int + Concurrency int + FinalizeTimeout int + // PollInterval is how often the out-of-band service is polled, in + // seconds. Zero uses DefaultOOBPollInterval. + PollInterval int + // HitRetention is how long a recorded hit stays available, in minutes. + // Zero uses DefaultOOBHitRetention. + HitRetention int +} + +// CuratedOptions configures the optional curated PoC source. +type CuratedOptions struct { + Enabled string + Endpoint string + TimeoutSec int + ForceUpdate bool +} + +// handlers holds the callbacks registered through the With*Handler options. +type handlers struct { + result []func(Result) + rawResult []func(*result.Result) + failure []func(Failure) + port []func(PortEvent) + host []func(HostEvent) + webProbe []func(WebProbeEvent) + progress []func(PhaseProgress) + scanInfo []func(ScanInfo) + monitor []func(string) +} + +// NewOptions returns Options populated with the SDK defaults. +func NewOptions() *Options { + return &Options{ + RateLimit: DefaultRateLimit, + Concurrency: DefaultConcurrency, + Retries: DefaultRetries, + Timeout: DefaultTimeout, + MaxHostError: DefaultMaxHostError, + MaxRespBodySize: DefaultMaxRespBodySize, + BruteMaxRequests: DefaultBruteMaxRequests, + DefaultAccept: true, + FingerprintFilterMode: FingerprintStrict, + IncludeRequestResponse: true, + StreamBuffer: DefaultStreamBuffer, + PortScan: PortScanOptions{ + Ports: "top", + ChunkSize: 1000, + }, + TaskTimeout: TaskTimeoutOptions{ + VisibleCapSec: DefaultTaskTimeoutVisibleCapSec, + NetCapSec: DefaultTaskTimeoutNetCapSec, + GoCapSec: DefaultTaskTimeoutGoCapSec, + }, + Monitor: ExecutionMonitorOptions{ + SlowThresholdSec: DefaultMonitorSlowThresholdSec, + SlowLogLimit: DefaultMonitorSlowLogLimit, + SummaryTop: DefaultMonitorSummaryTop, + SummaryBy: MonitorSummaryByMax, + }, + OOB: OOBOptions{ + RateLimit: DefaultOOBRateLimit, + Concurrency: DefaultOOBConcurrency, + FinalizeTimeout: -1, + PollInterval: DefaultOOBPollInterval, + HitRetention: DefaultOOBHitRetention, + }, + } +} + +// validate normalises the options and reports invalid combinations. +func (o *Options) validate() error { + // A cyberspace query is a target source in its own right, so it satisfies + // the "must have targets" requirement on its own. + if len(o.Targets) == 0 && strings.TrimSpace(o.TargetsFile) == "" && strings.TrimSpace(o.Cyberspace.Query) == "" { + return ErrNoTargets + } + + if strings.TrimSpace(o.Checkpoint.Path) != "" && o.Checkpoint.SaveInterval <= 0 { + o.Checkpoint.SaveInterval = DefaultCheckpointSaveInterval + } + + limitModes := 0 + for _, on := range []bool{o.ReqLimitPerTarget > 0, o.AutoReqLimit, o.Polite, o.Balanced, o.Aggressive} { + if on { + limitModes++ + } + } + if limitModes > 1 { + return fmt.Errorf("%w: only one of ReqLimitPerTarget, AutoReqLimit, Polite, Balanced and Aggressive may be set", ErrInvalidOptions) + } + if o.ReqLimitPerTarget < 0 { + return fmt.Errorf("%w: ReqLimitPerTarget must be >= 0", ErrInvalidOptions) + } + + if o.MaxRespBodySize <= 0 { + o.MaxRespBodySize = DefaultMaxRespBodySize + } + if o.StreamBuffer <= 0 { + o.StreamBuffer = DefaultStreamBuffer + } + if o.OOB.RateLimit <= 0 { + o.OOB.RateLimit = DefaultOOBRateLimit + } + if o.OOB.Concurrency <= 0 { + o.OOB.Concurrency = DefaultOOBConcurrency + } + if o.OOB.PollInterval <= 0 { + o.OOB.PollInterval = DefaultOOBPollInterval + } + if o.OOB.HitRetention <= 0 { + o.OOB.HitRetention = DefaultOOBHitRetention + } + + switch strings.ToLower(strings.TrimSpace(o.FingerprintFilterMode)) { + case FingerprintOpportunistic: + o.FingerprintFilterMode = FingerprintOpportunistic + default: + o.FingerprintFilterMode = FingerprintStrict + } + + if o.OOB.Enabled && strings.TrimSpace(o.OOB.Adapter) == "" { + return fmt.Errorf("%w: OOB adapter must not be empty when OOB is enabled", ErrInvalidOptions) + } + + // A zero cap would silently collapse the smart estimate to a 60 second + // ceiling deep inside the engine, so fill the documented default instead. + if o.TaskTimeout.VisibleCapSec <= 0 { + o.TaskTimeout.VisibleCapSec = DefaultTaskTimeoutVisibleCapSec + } + if o.TaskTimeout.NetCapSec <= 0 { + o.TaskTimeout.NetCapSec = DefaultTaskTimeoutNetCapSec + } + if o.TaskTimeout.GoCapSec <= 0 { + o.TaskTimeout.GoCapSec = DefaultTaskTimeoutGoCapSec + } + + switch strings.ToLower(strings.TrimSpace(o.Monitor.SummaryBy)) { + case MonitorSummaryByAvg: + o.Monitor.SummaryBy = MonitorSummaryByAvg + default: + o.Monitor.SummaryBy = MonitorSummaryByMax + } + + o.applyRequestLimitPreset() + return nil +} + +// applyRequestLimitPreset derives ReqLimitPerTarget from the selected preset. +func (o *Options) applyRequestLimitPreset() { + if o.ReqLimitPerTarget != 0 { + return + } + switch { + case o.Polite: + o.ReqLimitPerTarget = 5 + case o.Balanced: + o.ReqLimitPerTarget = 15 + case o.Aggressive: + o.ReqLimitPerTarget = 50 + case o.AutoReqLimit: + rate := o.RateLimit + if rate <= 0 { + rate = DefaultRateLimit + } + limit := rate / 10 + if limit < 5 { + limit = 5 + } + if limit > 15 { + limit = 15 + } + switch con := o.Concurrency; { + case con >= 100 && limit > 8: + limit = 8 + case con >= 50 && limit > 12: + limit = 12 + } + o.ReqLimitPerTarget = limit + } +} + +// pocInputs returns every explicitly configured PoC path. +func (o *Options) pocInputs() []string { + out := make([]string, 0, len(o.PocPaths)) + for _, p := range o.PocPaths { + if v := strings.TrimSpace(p); v != "" { + out = append(out, v) + } + } + return out +} + +// Option configures a Scanner. +// +// Functional options can be applied incrementally, validate their own input, +// and let new options be added without breaking existing callers. +type Option func(*Options) error + +// WithOptions replaces the whole configuration. Later options still apply on +// top of it, which makes it a useful escape hatch for callers that keep their +// own Options value. +func WithOptions(opts *Options) Option { + return func(o *Options) error { + if opts == nil { + return fmt.Errorf("%w: options must not be nil", ErrInvalidOptions) + } + handlers := o.handlers + *o = *opts + o.handlers = handlers + return nil + } +} + +// --- targets --- + +// WithTargets appends scan targets. +func WithTargets(targets ...string) Option { + return func(o *Options) error { + o.Targets = append(o.Targets, targets...) + return nil + } +} + +// WithTargetsFile reads targets from a file, one per line. +func WithTargetsFile(path string) Option { + return func(o *Options) error { + o.TargetsFile = path + return nil + } +} + +// WithCyberspace sources targets from an internet-wide search engine, so a +// scan can be driven by a query instead of an explicit target list. +// +// The engine's API key must be present in the afrog configuration file. +func WithCyberspace(cfg CyberspaceOptions) Option { + return func(o *Options) error { + engine := strings.ToLower(strings.TrimSpace(cfg.Engine)) + if engine == "" { + return fmt.Errorf("%w: cyberspace engine must not be empty", ErrInvalidOptions) + } + if engine != CyberspaceZoomEye { + return fmt.Errorf("%w: unsupported cyberspace engine %q, only %q is implemented", + ErrInvalidOptions, cfg.Engine, CyberspaceZoomEye) + } + if strings.TrimSpace(cfg.Query) == "" { + return fmt.Errorf("%w: cyberspace query must not be empty", ErrInvalidOptions) + } + if cfg.Count < 0 { + return fmt.Errorf("%w: cyberspace count must be >= 0, got %d", ErrInvalidOptions, cfg.Count) + } + cfg.Engine = engine + o.Cyberspace = cfg + return nil + } +} + +// WithTargetPreProbe probes each target's protocol and liveness in parallel +// with the scan, blacklisting hosts that exceed MaxHostError. +// +// This is the CLI's -mt flag. Despite the flag's name it does not watch the +// targets file for changes. +func WithTargetPreProbe() Option { + return func(o *Options) error { o.TargetPreProbe = true; return nil } +} + +// --- resume --- + +// WithCheckpoint makes the scan resumable by recording finished target/PoC +// pairs to a file and skipping them on a later run. +// +// See [CheckpointOptions] for the constraints. This is unrelated to +// Scanner.Resume, which lifts a Pause. +func WithCheckpoint(cfg CheckpointOptions) Option { + return func(o *Options) error { + if strings.TrimSpace(cfg.Path) == "" { + return fmt.Errorf("%w: checkpoint path must not be empty", ErrInvalidOptions) + } + if cfg.SaveInterval < 0 { + return fmt.Errorf("%w: checkpoint save interval must be >= 0, got %v", ErrInvalidOptions, cfg.SaveInterval) + } + if cfg.SaveInterval == 0 { + cfg.SaveInterval = DefaultCheckpointSaveInterval + } + o.Checkpoint = cfg + return nil + } +} + +// --- pocs --- + +// WithPocPaths appends PoC inputs. Each entry may be a single file, a +// directory searched recursively, or a glob pattern such as "dir/*.yaml". +func WithPocPaths(paths ...string) Option { + return func(o *Options) error { + o.PocPaths = append(o.PocPaths, paths...) + return nil + } +} + +// WithPocPathsOnly restricts the scan to the explicitly configured PoCs, +// hiding the built-in and user-directory sources. +func WithPocPathsOnly() Option { + return func(o *Options) error { + o.PocPathsOnly = true + return nil + } +} + +// WithSearch filters PoCs by keyword, for example "tomcat,phpinfo". +func WithSearch(keyword string) Option { + return func(o *Options) error { + o.Search = keyword + return nil + } +} + +// WithSeverity filters PoCs by severity, for example "high,critical". +func WithSeverity(severity string) Option { + return func(o *Options) error { + o.Severity = severity + return nil + } +} + +// WithExcludePocs excludes PoCs by id. +func WithExcludePocs(pocs ...string) Option { + return func(o *Options) error { + o.ExcludePocs = append(o.ExcludePocs, pocs...) + return nil + } +} + +// WithExcludePocsFile excludes the PoCs listed in a file. +func WithExcludePocsFile(path string) Option { + return func(o *Options) error { + o.ExcludePocsFile = path + return nil + } +} + +// --- performance --- + +// WithConcurrency sets the number of concurrent scan workers. +func WithConcurrency(n int) Option { + return func(o *Options) error { + if n <= 0 { + return fmt.Errorf("%w: concurrency must be > 0, got %d", ErrInvalidOptions, n) + } + o.Concurrency = n + return nil + } +} + +// WithRateLimit sets the maximum number of requests per second. +func WithRateLimit(n int) Option { + return func(o *Options) error { + if n <= 0 { + return fmt.Errorf("%w: rate limit must be > 0, got %d", ErrInvalidOptions, n) + } + o.RateLimit = n + return nil + } +} + +// WithTimeout sets the per-request timeout in seconds. +func WithTimeout(seconds int) Option { + return func(o *Options) error { + if seconds <= 0 { + return fmt.Errorf("%w: timeout must be > 0, got %d", ErrInvalidOptions, seconds) + } + o.Timeout = seconds + return nil + } +} + +// WithRetries sets how many times a failed request is retried. +func WithRetries(n int) Option { + return func(o *Options) error { + if n < 0 { + return fmt.Errorf("%w: retries must be >= 0, got %d", ErrInvalidOptions, n) + } + o.Retries = n + return nil + } +} + +// WithMaxHostError sets how many errors a host may produce before it is +// skipped. +func WithMaxHostError(n int) Option { + return func(o *Options) error { + if n < 0 { + return fmt.Errorf("%w: max host error must be >= 0, got %d", ErrInvalidOptions, n) + } + o.MaxHostError = n + return nil + } +} + +// WithMaxRespBodySize sets the response body read limit in megabytes. +// Responses beyond the limit are truncated and Exchange.BodyTruncated is set. +func WithMaxRespBodySize(mb int) Option { + return func(o *Options) error { + if mb <= 0 { + return fmt.Errorf("%w: max response body size must be > 0, got %d", ErrInvalidOptions, mb) + } + o.MaxRespBodySize = mb + return nil + } +} + +// WithRequestLimitPerTarget caps concurrent requests per target. +func WithRequestLimitPerTarget(n int) Option { + return func(o *Options) error { + if n < 0 { + return fmt.Errorf("%w: request limit per target must be >= 0, got %d", ErrInvalidOptions, n) + } + o.ReqLimitPerTarget = n + return nil + } +} + +// WithPolite applies the conservative per-target request preset. +func WithPolite() Option { + return func(o *Options) error { o.Polite = true; return nil } +} + +// WithBalanced applies the balanced per-target request preset. +func WithBalanced() Option { + return func(o *Options) error { o.Balanced = true; return nil } +} + +// WithAggressive applies the aggressive per-target request preset. +func WithAggressive() Option { + return func(o *Options) error { o.Aggressive = true; return nil } +} + +// WithAutoRequestLimit derives the per-target request limit from the rate +// limit and concurrency. +func WithAutoRequestLimit() Option { + return func(o *Options) error { o.AutoReqLimit = true; return nil } +} + +// WithSmartConcurrency adjusts concurrency based on the target count. +func WithSmartConcurrency() Option { + return func(o *Options) error { o.Smart = true; return nil } +} + +// WithStopOnFirstMatch stops the scan as soon as a vulnerability is found. +func WithStopOnFirstMatch() Option { + return func(o *Options) error { o.StopOnFirstMatch = true; return nil } +} + +// --- per-task timeout --- + +// WithTaskTimeout bounds how long a single target+PoC task may run, so that +// one pathological PoC cannot hold a worker indefinitely. +// +// Zero-valued cap fields keep the current defaults. +func WithTaskTimeout(cfg TaskTimeoutOptions) Option { + return func(o *Options) error { + if cfg.HardSec < 0 { + return fmt.Errorf("%w: task hard timeout must be >= 0, got %d", ErrInvalidOptions, cfg.HardSec) + } + for name, v := range map[string]int{ + "VisibleCapSec": cfg.VisibleCapSec, + "NetCapSec": cfg.NetCapSec, + "GoCapSec": cfg.GoCapSec, + } { + if v < 0 { + return fmt.Errorf("%w: task timeout %s must be >= 0, got %d", ErrInvalidOptions, name, v) + } + } + if cfg.VisibleCapSec <= 0 { + cfg.VisibleCapSec = o.TaskTimeout.VisibleCapSec + } + if cfg.NetCapSec <= 0 { + cfg.NetCapSec = o.TaskTimeout.NetCapSec + } + if cfg.GoCapSec <= 0 { + cfg.GoCapSec = o.TaskTimeout.GoCapSec + } + o.TaskTimeout = cfg + return nil + } +} + +// --- execution monitor --- + +// WithExecutionMonitor enables the PoC execution duration monitor, which +// surfaces slow and stuck PoCs during the scan and a slowest-PoC summary at +// the end. +// +// Register at least one WithMonitorHandler to receive the reports; the SDK +// never writes them to the console. Zero-valued fields keep the defaults. +func WithExecutionMonitor(cfg ExecutionMonitorOptions) Option { + return func(o *Options) error { + for name, v := range map[string]int{ + "LogLimit": cfg.LogLimit, + "SlowThresholdSec": cfg.SlowThresholdSec, + "SlowLogLimit": cfg.SlowLogLimit, + "SummaryTop": cfg.SummaryTop, + } { + if v < 0 { + return fmt.Errorf("%w: monitor %s must be >= 0, got %d", ErrInvalidOptions, name, v) + } + } + if strings.TrimSpace(cfg.SummaryBy) == "" { + cfg.SummaryBy = o.Monitor.SummaryBy + } + switch strings.ToLower(strings.TrimSpace(cfg.SummaryBy)) { + case MonitorSummaryByMax, MonitorSummaryByAvg: + default: + return fmt.Errorf("%w: unknown monitor summary key %q", ErrInvalidOptions, cfg.SummaryBy) + } + if cfg.SlowThresholdSec == 0 { + cfg.SlowThresholdSec = o.Monitor.SlowThresholdSec + } + if cfg.SlowLogLimit == 0 { + cfg.SlowLogLimit = o.Monitor.SlowLogLimit + } + if cfg.SummaryTop == 0 { + cfg.SummaryTop = o.Monitor.SummaryTop + } + o.Monitor = cfg + o.EnableMonitor = true + return nil + } +} + +// --- fingerprinting --- + +// WithFingerprintDisabled skips the fingerprinting stage. +func WithFingerprintDisabled() Option { + return func(o *Options) error { o.DisableFingerprint = true; return nil } +} + +// WithFingerprintFilterMode selects how fingerprint gating filters PoCs. +// Valid values are FingerprintStrict and FingerprintOpportunistic. +func WithFingerprintFilterMode(mode string) Option { + return func(o *Options) error { + switch strings.ToLower(strings.TrimSpace(mode)) { + case FingerprintStrict, FingerprintOpportunistic: + o.FingerprintFilterMode = strings.ToLower(strings.TrimSpace(mode)) + return nil + default: + return fmt.Errorf("%w: unknown fingerprint filter mode %q", ErrInvalidOptions, mode) + } + } +} + +// WithWebProbe enables the web probing stage. +func WithWebProbe() Option { + return func(o *Options) error { o.EnableWebProbe = true; return nil } +} + +// --- network --- + +// WithProxy sets an HTTP or SOCKS5 proxy. +func WithProxy(proxy string) Option { + return func(o *Options) error { + o.Proxy = proxy + return nil + } +} + +// WithHeaders appends custom request headers in "Name: value" form. +func WithHeaders(headers ...string) Option { + return func(o *Options) error { + for _, h := range headers { + if !strings.Contains(h, ":") { + return fmt.Errorf("%w: header %q must be in \"Name: value\" form", ErrInvalidOptions, h) + } + } + o.Headers = append(o.Headers, headers...) + return nil + } +} + +// --- port pre-scan --- + +// WithPortScan runs a port pre-scan before the PoC scan. Discovered open ports +// are appended to the target set as host:port. +func WithPortScan(cfg PortScanOptions) Option { + return func(o *Options) error { + if cfg.Ports == "" { + cfg.Ports = o.PortScan.Ports + } + if cfg.ChunkSize <= 0 { + cfg.ChunkSize = o.PortScan.ChunkSize + } + o.EnablePortSan = true + o.PortScan = cfg + return nil + } +} + +// --- OOB --- + +// WithOOB enables out-of-band detection. +func WithOOB(cfg OOBOptions) Option { + return func(o *Options) error { + if strings.TrimSpace(cfg.Adapter) == "" { + return fmt.Errorf("%w: OOB adapter must not be empty", ErrInvalidOptions) + } + if cfg.RateLimit <= 0 { + cfg.RateLimit = o.OOB.RateLimit + } + if cfg.Concurrency <= 0 { + cfg.Concurrency = o.OOB.Concurrency + } + if cfg.FinalizeTimeout == 0 { + cfg.FinalizeTimeout = o.OOB.FinalizeTimeout + } + if cfg.PollInterval <= 0 { + cfg.PollInterval = o.OOB.PollInterval + } + if cfg.HitRetention <= 0 { + cfg.HitRetention = o.OOB.HitRetention + } + cfg.Enabled = true + o.OOB = cfg + return nil + } +} + +// --- curated --- + +// WithCurated configures the optional curated PoC source. +func WithCurated(cfg CuratedOptions) Option { + return func(o *Options) error { + o.Curated = cfg + return nil + } +} + +// --- notifications --- + +// WithDingtalk enables DingTalk webhook notifications. The token must be +// present in the afrog configuration file. +func WithDingtalk() Option { + return func(o *Options) error { o.Dingtalk = true; return nil } +} + +// WithWecom enables WeCom webhook notifications. The token must be present in +// the afrog configuration file. +func WithWecom() Option { + return func(o *Options) error { o.Wecom = true; return nil } +} + +// --- output --- + +// WithRequestResponse controls whether results carry raw request and response +// messages. Enabled by default. +func WithRequestResponse(include bool) Option { + return func(o *Options) error { + o.IncludeRequestResponse = include + return nil + } +} + +// WithMaxStoredResults caps how many results accumulate in memory. Zero means +// unlimited. Handlers and streams still receive every result. +func WithMaxStoredResults(n int) Option { + return func(o *Options) error { + if n < 0 { + return fmt.Errorf("%w: max stored results must be >= 0, got %d", ErrInvalidOptions, n) + } + o.MaxStoredResults = n + return nil + } +} + +// WithStreamBuffer sets the capacity of each subscribed stream channel. +func WithStreamBuffer(n int) Option { + return func(o *Options) error { + if n <= 0 { + return fmt.Errorf("%w: stream buffer must be > 0, got %d", ErrInvalidOptions, n) + } + o.StreamBuffer = n + return nil + } +} + +// WithVerbose prints a scan summary to stdout before the scan starts. +func WithVerbose() Option { + return func(o *Options) error { o.Verbose = true; return nil } +} + +// WithRedactedHeaders masks the given headers in the raw request and response +// of every [Exchange], as well as in the structured header maps. +// +// Calling it with no arguments uses [DefaultRedactedHeaders], which covers the +// usual credential-bearing headers such as Authorization and Cookie. +// +// Redaction is opt-in because the raw messages are the point of Exchange, and +// masking them by default would quietly weaken the SDK's main debugging aid. +// Enable it whenever results are logged, persisted or returned over an API, +// especially when scanning authenticated targets or using WithHeaders to send +// credentials. +func WithRedactedHeaders(names ...string) Option { + return func(o *Options) error { + if len(names) == 0 { + names = DefaultRedactedHeaders + } + for _, n := range names { + if v := strings.ToLower(strings.TrimSpace(n)); v != "" { + o.RedactedHeaders = append(o.RedactedHeaders, v) + } + } + return nil + } +} + +// --- handlers --- + +// WithResultHandler registers a result callback. It may be used several times. +// +// Handlers are invoked concurrently from scan workers, so implementations must +// synchronise their own shared state. +func WithResultHandler(fn func(Result)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.result = append(o.handlers.result, fn) + } + return nil + } +} + +// WithRawResultHandler registers a callback that receives the engine's own +// result type. +// +// This is an escape hatch for advanced integrations that need fields the +// stable [Result] view does not expose, such as persisting the full internal +// structure. Unlike Result, the shape of result.Result is an internal detail +// and may change between releases. Prefer WithResultHandler. +// +// The value passed to the handler is a snapshot owned by the caller. +func WithRawResultHandler(fn func(*result.Result)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.rawResult = append(o.handlers.rawResult, fn) + } + return nil + } +} + +// WithFailureHandler registers a callback for PoC execution failures such as +// request errors, expression errors and recovered panics. +func WithFailureHandler(fn func(Failure)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.failure = append(o.handlers.failure, fn) + } + return nil + } +} + +// WithPortHandler registers a callback for open ports found during the port +// pre-scan. +func WithPortHandler(fn func(PortEvent)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.port = append(o.handlers.port, fn) + } + return nil + } +} + +// WithHostHandler registers a callback for hosts found during discovery. +func WithHostHandler(fn func(HostEvent)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.host = append(o.handlers.host, fn) + } + return nil + } +} + +// WithWebProbeHandler registers a callback for probed web services. +func WithWebProbeHandler(fn func(WebProbeEvent)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.webProbe = append(o.handlers.webProbe, fn) + } + return nil + } +} + +// WithProgressHandler registers a callback for phase progress updates. +func WithProgressHandler(fn func(PhaseProgress)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.progress = append(o.handlers.progress, fn) + } + return nil + } +} + +// WithScanInfoHandler registers a callback for scan summary updates. +func WithScanInfoHandler(fn func(ScanInfo)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.scanInfo = append(o.handlers.scanInfo, fn) + } + return nil + } +} + +// WithMonitorHandler registers a callback for execution monitor reports. Each +// call receives one preformatted line. +// +// Without a handler the monitor still runs but its output goes nowhere, which +// is why WithExecutionMonitor is only useful together with this option. +func WithMonitorHandler(fn func(line string)) Option { + return func(o *Options) error { + if fn != nil { + o.handlers.monitor = append(o.handlers.monitor, fn) + } + return nil + } +} diff --git a/pkg/sdk/options_test.go b/pkg/sdk/options_test.go new file mode 100644 index 000000000..3f26069ae --- /dev/null +++ b/pkg/sdk/options_test.go @@ -0,0 +1,217 @@ +package sdk + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" +) + +func TestNewOptions_Defaults(t *testing.T) { + o := NewOptions() + + tests := []struct { + name string + got any + want any + }{ + {"RateLimit", o.RateLimit, DefaultRateLimit}, + {"Concurrency", o.Concurrency, DefaultConcurrency}, + {"Retries", o.Retries, DefaultRetries}, + {"Timeout", o.Timeout, DefaultTimeout}, + {"MaxHostError", o.MaxHostError, DefaultMaxHostError}, + {"MaxRespBodySize", o.MaxRespBodySize, DefaultMaxRespBodySize}, + {"BruteMaxRequests", o.BruteMaxRequests, DefaultBruteMaxRequests}, + {"StreamBuffer", o.StreamBuffer, DefaultStreamBuffer}, + {"FingerprintFilterMode", o.FingerprintFilterMode, FingerprintStrict}, + {"IncludeRequestResponse", o.IncludeRequestResponse, true}, + {"Verbose", o.Verbose, false}, + {"MaxStoredResults", o.MaxStoredResults, 0}, + {"OOB.RateLimit", o.OOB.RateLimit, DefaultOOBRateLimit}, + {"OOB.Concurrency", o.OOB.Concurrency, DefaultOOBConcurrency}, + {"OOB.PollInterval", o.OOB.PollInterval, DefaultOOBPollInterval}, + {"OOB.HitRetention", o.OOB.HitRetention, DefaultOOBHitRetention}, + } + for _, tt := range tests { + if tt.got != tt.want { + t.Errorf("%s = %v, want %v", tt.name, tt.got, tt.want) + } + } +} + +// The OOB poll cadence must match the command line's. Leaving PollInterval at +// zero makes the engine fall back to one second, which polls the out-of-band +// provider twice as often as afrog's own CLI does. +func TestWithOOB_KeepsPollDefaultsWhenUnset(t *testing.T) { + o := NewOptions() + if err := WithOOB(OOBOptions{Adapter: "ceyeio"})(o); err != nil { + t.Fatalf("WithOOB: %v", err) + } + if o.OOB.PollInterval != DefaultOOBPollInterval { + t.Errorf("PollInterval = %d, want %d", o.OOB.PollInterval, DefaultOOBPollInterval) + } + if o.OOB.HitRetention != DefaultOOBHitRetention { + t.Errorf("HitRetention = %d, want %d", o.OOB.HitRetention, DefaultOOBHitRetention) + } + + // An explicit value must survive. + o2 := NewOptions() + if err := WithOOB(OOBOptions{Adapter: "ceyeio", PollInterval: 7, HitRetention: 3})(o2); err != nil { + t.Fatalf("WithOOB: %v", err) + } + if o2.OOB.PollInterval != 7 || o2.OOB.HitRetention != 3 { + t.Errorf("explicit values lost: PollInterval=%d HitRetention=%d", o2.OOB.PollInterval, o2.OOB.HitRetention) + } +} + +func TestOptions_ValidationErrors(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "a.yaml"), []byte("id: a\ninfo:\n name: a\n severity: info\nrules:\n r0:\n request:\n method: GET\n path: /\n expression: response.status == 200\nexpression: r0()\n"), 0o644); err != nil { + t.Fatalf("write poc: %v", err) + } + + tests := []struct { + name string + option Option + }{ + {"zero concurrency", WithConcurrency(0)}, + {"negative concurrency", WithConcurrency(-1)}, + {"zero rate limit", WithRateLimit(0)}, + {"zero timeout", WithTimeout(0)}, + {"negative retries", WithRetries(-1)}, + {"negative max host error", WithMaxHostError(-1)}, + {"zero response body size", WithMaxRespBodySize(0)}, + {"negative max stored results", WithMaxStoredResults(-1)}, + {"zero stream buffer", WithStreamBuffer(0)}, + {"empty oob adapter", WithOOB(OOBOptions{})}, + {"malformed header", WithHeaders("no-colon")}, + {"unknown fingerprint mode", WithFingerprintFilterMode("bogus")}, + {"nil options", WithOptions(nil)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := New(context.Background(), + WithTargets("http://127.0.0.1:1"), + WithPocPaths(dir), + tt.option, + ) + if !errors.Is(err, ErrInvalidOptions) { + t.Fatalf("error = %v, want ErrInvalidOptions", err) + } + }) + } +} + +func TestOptions_MutuallyExclusiveRequestLimits(t *testing.T) { + o := NewOptions() + o.Targets = []string{"http://127.0.0.1:1"} + o.Polite = true + o.Aggressive = true + + if err := o.validate(); !errors.Is(err, ErrInvalidOptions) { + t.Fatalf("error = %v, want ErrInvalidOptions", err) + } +} + +func TestOptions_RequestLimitPresets(t *testing.T) { + tests := []struct { + name string + apply func(*Options) + want int + }{ + {"polite", func(o *Options) { o.Polite = true }, 5}, + {"balanced", func(o *Options) { o.Balanced = true }, 15}, + {"aggressive", func(o *Options) { o.Aggressive = true }, 50}, + {"auto derives from rate limit", func(o *Options) { o.AutoReqLimit = true; o.RateLimit = 100 }, 10}, + {"auto clamps low", func(o *Options) { o.AutoReqLimit = true; o.RateLimit = 10 }, 5}, + {"auto clamps high", func(o *Options) { o.AutoReqLimit = true; o.RateLimit = 1000 }, 15}, + {"auto reduced by high concurrency", func(o *Options) { o.AutoReqLimit = true; o.RateLimit = 1000; o.Concurrency = 100 }, 8}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + o := NewOptions() + o.Targets = []string{"http://127.0.0.1:1"} + tt.apply(o) + + if err := o.validate(); err != nil { + t.Fatalf("validate: %v", err) + } + if o.ReqLimitPerTarget != tt.want { + t.Errorf("ReqLimitPerTarget = %d, want %d", o.ReqLimitPerTarget, tt.want) + } + }) + } +} + +func TestOptions_NormalisesInvalidValues(t *testing.T) { + o := NewOptions() + o.Targets = []string{"http://127.0.0.1:1"} + o.MaxRespBodySize = 0 + o.StreamBuffer = -1 + o.FingerprintFilterMode = "NONSENSE" + o.OOB.RateLimit = 0 + o.OOB.Concurrency = -5 + + if err := o.validate(); err != nil { + t.Fatalf("validate: %v", err) + } + + if o.MaxRespBodySize != DefaultMaxRespBodySize { + t.Errorf("MaxRespBodySize = %d, want %d", o.MaxRespBodySize, DefaultMaxRespBodySize) + } + if o.StreamBuffer != DefaultStreamBuffer { + t.Errorf("StreamBuffer = %d, want %d", o.StreamBuffer, DefaultStreamBuffer) + } + if o.FingerprintFilterMode != FingerprintStrict { + t.Errorf("FingerprintFilterMode = %q, want %q", o.FingerprintFilterMode, FingerprintStrict) + } + if o.OOB.RateLimit != DefaultOOBRateLimit { + t.Errorf("OOB.RateLimit = %d, want %d", o.OOB.RateLimit, DefaultOOBRateLimit) + } + if o.OOB.Concurrency != DefaultOOBConcurrency { + t.Errorf("OOB.Concurrency = %d, want %d", o.OOB.Concurrency, DefaultOOBConcurrency) + } +} + +func TestWithOptions_PreservesHandlers(t *testing.T) { + o := NewOptions() + if err := WithResultHandler(func(Result) {})(o); err != nil { + t.Fatalf("WithResultHandler: %v", err) + } + + replacement := NewOptions() + replacement.Concurrency = 99 + if err := WithOptions(replacement)(o); err != nil { + t.Fatalf("WithOptions: %v", err) + } + + if o.Concurrency != 99 { + t.Errorf("Concurrency = %d, want 99", o.Concurrency) + } + if len(o.handlers.result) != 1 { + t.Errorf("handlers were dropped by WithOptions: got %d, want 1", len(o.handlers.result)) + } +} + +func TestWithPortScan_KeepsDefaultsForZeroFields(t *testing.T) { + o := NewOptions() + if err := WithPortScan(PortScanOptions{TimeoutMs: 500})(o); err != nil { + t.Fatalf("WithPortScan: %v", err) + } + + if !o.EnablePortSan { + t.Error("WithPortScan should enable the port pre-scan") + } + if o.PortScan.Ports != "top" { + t.Errorf("Ports = %q, want the default \"top\"", o.PortScan.Ports) + } + if o.PortScan.ChunkSize != 1000 { + t.Errorf("ChunkSize = %d, want the default 1000", o.PortScan.ChunkSize) + } + if o.PortScan.TimeoutMs != 500 { + t.Errorf("TimeoutMs = %d, want 500", o.PortScan.TimeoutMs) + } +} diff --git a/pkg/sdk/scanner.go b/pkg/sdk/scanner.go new file mode 100644 index 000000000..8716af31b --- /dev/null +++ b/pkg/sdk/scanner.go @@ -0,0 +1,1345 @@ +package sdk + +import ( + "context" + "fmt" + "os" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/zan8in/oobadapter/pkg/oobadapter" + + "github.com/zan8in/afrog/v3/pkg/config" + "github.com/zan8in/afrog/v3/pkg/curated/service" + "github.com/zan8in/afrog/v3/pkg/fingerprint" + "github.com/zan8in/afrog/v3/pkg/poc" + "github.com/zan8in/afrog/v3/pkg/protocols/http/retryhttpclient" + "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/pkg/runner" + "github.com/zan8in/afrog/v3/pkg/targets" + "github.com/zan8in/afrog/v3/pkg/utils" +) + +// Scan lifecycle states. +const ( + stateIdle int32 = iota + stateRunning + stateFinished +) + +// Scanner runs a vulnerability scan. +// +// A Scanner is single-use: once a scan finishes the instance cannot be +// restarted. All methods are safe for concurrent use. +type Scanner struct { + opts *Options + runner *runner.Runner + // internal is the engine-level configuration derived from opts. + internal *config.Options + + ctx context.Context + cancel context.CancelFunc + + state atomic.Int32 + stopping atomic.Bool + closed atomic.Bool + + runDone chan struct{} + runDoneOnce sync.Once + + runErrMu sync.Mutex + runErr error + + resultsMu sync.Mutex + results []Result + + portsMu sync.Mutex + openPorts map[string]map[int]struct{} + + statsMu sync.RWMutex + stats Stats + // completed and found are hot counters kept outside statsMu. + completed atomic.Int64 + found atomic.Int64 + + phasesMu sync.Mutex + phases map[string]PhaseProgress + lastVuln atomic.Int32 + + resultEvents *emitter[Result] + portEvents *emitter[PortEvent] + hostEvents *emitter[HostEvent] + webProbeEvents *emitter[WebProbeEvent] + progressEvents *emitter[PhaseProgress] + scanInfoEvents *emitter[ScanInfo] + + pocs []poc.Poc + pocDiags []config.PocLoadError + curatedErr error + includeRR bool + maxStored int + redact map[string]struct{} + oobProbedMu sync.Mutex + + // curatedEnv restores the process environment that mountCurated changed, + // so a second scanner in the same process is not affected by the first. + curatedEnv []envEntry +} + +// envEntry remembers an environment variable's prior state. +type envEntry struct { + key string + value string + set bool +} + +// New creates a Scanner from the given options. +// +// The parent context bounds the scanner's lifetime: cancelling it stops any +// running scan and releases every background goroutine. +func New(ctx context.Context, options ...Option) (*Scanner, error) { + if ctx == nil { + ctx = context.Background() + } + + opts := NewOptions() + for _, apply := range options { + if apply == nil { + continue + } + if err := apply(opts); err != nil { + return nil, err + } + } + if err := opts.validate(); err != nil { + return nil, err + } + + built, err := buildInternalOptions(opts) + if err != nil { + return nil, err + } + internal, curatedErr := built.options, built.curatedErr + + scanCtx, cancel := context.WithCancel(ctx) + s := &Scanner{ + opts: opts, + internal: internal, + ctx: scanCtx, + cancel: cancel, + runDone: make(chan struct{}), + openPorts: make(map[string]map[int]struct{}), + phases: make(map[string]PhaseProgress), + curatedErr: curatedErr, + includeRR: opts.IncludeRequestResponse, + maxStored: opts.MaxStoredResults, + redact: toSet(opts.RedactedHeaders), + curatedEnv: built.envRestore, + resultEvents: newEmitter(opts.StreamBuffer, opts.handlers.result), + portEvents: newEmitter(opts.StreamBuffer, opts.handlers.port), + hostEvents: newEmitter(opts.StreamBuffer, opts.handlers.host), + webProbeEvents: newEmitter(opts.StreamBuffer, opts.handlers.webProbe), + progressEvents: newEmitter(opts.StreamBuffer, opts.handlers.progress), + scanInfoEvents: newEmitter(opts.StreamBuffer, opts.handlers.scanInfo), + } + s.stats.StartTime = time.Now() + s.lastVuln.Store(-1) + + s.installEngineHooks() + + r, err := buildRunner(internal) + if err != nil { + s.abandon(cancel) + return nil, err + } + s.runner = r + s.installRunnerHooks() + + if err := s.loadPocs(); err != nil { + s.abandon(cancel) + return nil, err + } + + return s, nil +} + +// abandon releases everything a partially built scanner already owns. Without +// it a failed New would leak the runner's background goroutines and leave the +// curated environment variables pointing at this scanner's PoC directory. +func (s *Scanner) abandon(cancel context.CancelFunc) { + cancel() + s.closed.Store(true) + s.runner.Release() + s.restoreEnv() +} + +// installEngineHooks wires the config-level callbacks. They must be installed +// before the runner is built because the engine captures them. +func (s *Scanner) installEngineHooks() { + s.internal.OnPortScanResult = func(host string, port int) { + s.recordOpenPort(host, port) + s.portEvents.emit(s.ctx, PortEvent{Host: host, Port: port}) + } + + s.internal.OnHostDiscovered = func(host string) { + if host = strings.TrimSpace(host); host == "" { + return + } + s.hostEvents.emit(s.ctx, HostEvent{Host: host}) + } + + s.internal.OnPhaseProgress = func(phase, status string, finished, total int64, percent int) { + phase = strings.ToLower(strings.TrimSpace(phase)) + if phase == "" { + return + } + pp := PhaseProgress{ + Phase: phase, + Status: strings.ToLower(strings.TrimSpace(status)), + Finished: finished, + Total: total, + Percent: clampPercent(percent), + } + + s.phasesMu.Lock() + s.phases[phase] = pp + s.phasesMu.Unlock() + + s.progressEvents.emit(s.ctx, pp) + } + + // Always non-nil: the engine falls back to printing execution monitor + // lines itself when this hook is missing, which would break the SDK's + // promise to write nothing to the console. + s.internal.OnPedmLog = func(line string) { + for _, fn := range s.opts.handlers.monitor { + fn(line) + } + } + + s.internal.OnScanInfoUpdate = func(info config.ScanInfoUpdate) { + s.statsMu.Lock() + s.stats.TotalTargets = info.TotalTargets + s.stats.TotalPocs = info.TotalPocs + s.stats.TotalScans = info.TotalScans + s.statsMu.Unlock() + + s.scanInfoEvents.emit(s.ctx, ScanInfo{ + TotalTargets: info.TotalTargets, + Targets: append([]string(nil), info.Targets...), + TotalPocs: info.TotalPocs, + TotalScans: info.TotalScans, + OOBEnabled: info.OOBEnabled, + OOBStatus: info.OOBStatus, + }) + } +} + +// installRunnerHooks wires the runner-level callbacks. +func (s *Scanner) installRunnerHooks() { + s.runner.OnWebProbe = func(meta runner.WebMeta) { + s.webProbeEvents.emit(s.ctx, WebProbeEvent{ + URL: strings.TrimSpace(meta.URL), + Title: strings.TrimSpace(meta.Title), + Server: strings.TrimSpace(meta.Server), + PoweredBy: strings.TrimSpace(meta.PoweredBy), + }) + } + + s.runner.OnFailure = func(target, pocID string, err error) { + if err == nil { + return + } + f := Failure{Target: target, PocID: pocID, Err: err} + for _, fn := range s.opts.handlers.failure { + fn(f) + } + } + + s.runner.OnResult = s.handleResult + s.runner.OnFingerprint = s.handleFingerprint +} + +// handleResult receives every executed task, whether it matched or not. +func (s *Scanner) handleResult(r *result.Result) { + if r == nil || !r.SkipCount { + s.completed.Add(1) + // The checkpoint's "done tasks" counter lives on the engine options + // and is pre-seeded when resuming, so it must be incremented rather + // than overwritten. + atomic.AddUint32(&s.internal.CurrentCount, 1) + } + s.emitVulnProgress() + + if r == nil || !r.IsVul { + return + } + + // Snapshot at the SDK boundary: the engine owns r and its protobuf + // payloads, so handing that pointer to callers would tie their data to the + // engine's object lifetime. + snap := r.Snapshot() + if !s.includeRR { + snap.AllPocResult = nil + } + + s.found.Add(1) + out := newResultRedacted(snap, s.includeRR, time.Now(), s.redact) + + s.resultsMu.Lock() + if s.maxStored <= 0 || len(s.results) < s.maxStored { + s.results = append(s.results, out) + } + s.resultsMu.Unlock() + + for _, fn := range s.opts.handlers.rawResult { + fn(snap) + } + s.resultEvents.emit(s.ctx, out) + + if s.opts.StopOnFirstMatch { + s.Stop() + } +} + +// handleFingerprint turns fingerprint hits into results. +func (s *Scanner) handleFingerprint(targetKey string, hits []fingerprint.Hit) { + for _, hit := range hits { + severity := strings.TrimSpace(hit.Severity) + if severity == "" { + severity = "info" + } + name := strings.TrimSpace(hit.Name) + if name == "" { + name = strings.TrimSpace(hit.ID) + } + + rst := s.runner.FingerprintResult(targetKey, hit.ID) + if rst == nil { + rst = &result.Result{ + Target: targetKey, + FullTarget: targetKey, + PocInfo: &poc.Poc{Id: hit.ID}, + } + } else if rst.PocInfo == nil { + rst.PocInfo = &poc.Poc{Id: hit.ID} + } else { + rst.PocInfo.Id = hit.ID + } + + rst.IsVul = true + rst.PocInfo.Info.Name = name + rst.PocInfo.Info.Severity = severity + rst.PocInfo.Info.Tags = hit.Tags + rst.FingerResult = []fingerprint.Hit{hit} + if strings.TrimSpace(rst.FullTarget) == "" { + rst.FullTarget = rst.Target + } + + s.handleResult(rst) + } +} + +// emitVulnProgress publishes vulnerability-stage progress, throttled to whole +// percentage points so that a large scan does not flood subscribers. +func (s *Scanner) emitVulnProgress() { + if s.internal.OnPhaseProgress == nil { + return + } + completed, total, percent := s.vulnProgress() + if int32(percent) == s.lastVuln.Load() { + return + } + s.lastVuln.Store(int32(percent)) + + status := "running" + if total == 0 || (completed >= total && percent >= 100) { + status = "completed" + } + s.internal.OnPhaseProgress(PhaseVuln, status, completed, total, percent) +} + +// vulnProgress reports how far the vulnerability stage has advanced. A scan +// with no estimated tasks counts as complete rather than as 0%. +func (s *Scanner) vulnProgress() (completed, total int64, percent int) { + total = int64(s.totalScans()) + completed = s.completed.Load() + percent = 100 + if total > 0 { + percent = clampPercent(int(completed * 100 / total)) + } + return completed, total, percent +} + +// loadPocs resolves the PoC set once and reuses it for the scan statistics. +func (s *Scanner) loadPocs() error { + pocs, diags := s.internal.CreatePocListWithDiagnostics() + s.pocs = pocs + s.pocDiags = diags + + fingerprintPocs, rest := s.internal.FingerprintPoCs(pocs) + + seeds := make([]string, 0, s.internal.Targets.Len()) + for _, t := range s.internal.Targets.List() { + v, ok := t.(string) + if !ok { + continue + } + if v = strings.TrimSpace(v); v != "" { + seeds = append(seeds, v) + } + } + netTargets := targets.BuildTargetIndex(seeds).NetTargets() + + taskCount := 0 + if !s.internal.DisableFingerprint && len(fingerprintPocs) > 0 { + taskCount += len(fingerprintPocs) * len(seeds) + } + for _, p := range rest { + if isNetOnlyPoc(p) { + taskCount += len(netTargets) + } else { + taskCount += len(seeds) + } + } + + totalPocs := len(rest) + if !s.internal.DisableFingerprint && len(fingerprintPocs) > 0 { + totalPocs += len(fingerprintPocs) + } + + s.statsMu.Lock() + s.stats.TotalTargets = len(seeds) + s.stats.TotalPocs = totalPocs + s.stats.TotalScans = taskCount + s.statsMu.Unlock() + + if totalPocs == 0 { + return ErrNoPocs + } + return nil +} + +// isNetOnlyPoc reports whether a PoC only targets raw network protocols and so +// must not be scheduled against plain HTTP targets. +func isNetOnlyPoc(p poc.Poc) bool { + var hasHTTP, hasNet, hasGo bool + for _, rule := range p.Rules { + switch strings.ToLower(strings.TrimSpace(rule.Value.Request.Type)) { + case "", poc.HTTP_Type, poc.HTTPS_Type: + hasHTTP = true + case poc.TCP_Type, poc.UDP_Type, poc.SSL_Type: + hasNet = true + case poc.GO_Type: + hasGo = true + default: + hasHTTP = true + } + } + if hasGo { + return false + } + return hasNet && !hasHTTP +} + +// ---------------------------------------------------------------- lifecycle + +// Execute runs the scan synchronously and returns when it finishes. +// +// Cancelling ctx stops the scan; Execute then returns the context error. +func (s *Scanner) Execute(ctx context.Context) error { + if err := s.Start(ctx); err != nil { + return err + } + return s.Wait(ctx) +} + +// Start runs the scan asynchronously and returns immediately. Use Wait or Done +// to observe completion. +// +// Start returns ErrAlreadyRunning if a scan is in progress and +// ErrAlreadyFinished if this scanner has already been used. +func (s *Scanner) Start(ctx context.Context) error { + if s.closed.Load() { + return ErrClosed + } + if !s.state.CompareAndSwap(stateIdle, stateRunning) { + if s.state.Load() == stateRunning { + return ErrAlreadyRunning + } + return ErrAlreadyFinished + } + + // Bridge the caller's context to the scan. The watcher exits when the scan + // finishes, so it cannot outlive the scanner. + if ctx != nil && ctx.Done() != nil { + go func() { + select { + case <-ctx.Done(): + s.Stop() + case <-s.runDone: + } + }() + } + + go func() { + // Deferred in reverse order: streams close first so that a subscriber + // blocked in range is released, then runDone unblocks Wait, then the + // state flips. All three must happen even if the scan panics, + // otherwise a subscriber or a Wait would hang forever. + defer s.state.Store(stateFinished) + defer s.runDoneOnce.Do(func() { close(s.runDone) }) + defer s.closeStreams() + + defer func() { + if r := recover(); r != nil { + s.setRunErr(fmt.Errorf("afrog/sdk: scan panicked: %v", r)) + } + }() + + s.printSummary() + s.setRunErr(s.run()) + }() + return nil +} + +func (s *Scanner) setRunErr(err error) { + s.runErrMu.Lock() + s.runErr = err + s.runErrMu.Unlock() +} + +// run executes the scan and finalises the statistics. +func (s *Scanner) run() error { + stopSaving := s.startCheckpointSaver() + + // Run rather than Execute: it additionally starts the target pre-probe + // when TargetPreProbe is set. Without a pre-probe configured the two are + // identical. + _ = s.runner.Run() + + stopSaving() + s.saveCheckpoint() + + s.emitFinalProgress() + + s.statsMu.Lock() + s.stats.EndTime = time.Now() + s.statsMu.Unlock() + + return s.ctx.Err() +} + +// startCheckpointSaver rewrites the checkpoint file on a fixed interval and +// returns a function that stops it and waits for the goroutine to exit. +// +// The engine only reads the checkpoint; persisting it is the embedder's job, +// which is why the SDK has to run this loop itself. +func (s *Scanner) startCheckpointSaver() func() { + if s.opts.Checkpoint.Path == "" || s.runner == nil || s.runner.ScanProgress == nil { + return func() {} + } + + stop := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + ticker := time.NewTicker(s.opts.Checkpoint.SaveInterval) + defer ticker.Stop() + for { + select { + case <-stop: + return + case <-s.ctx.Done(): + return + case <-ticker.C: + s.saveCheckpoint() + } + } + }() + + var once sync.Once + return func() { + once.Do(func() { close(stop) }) + <-done + } +} + +// saveCheckpoint writes the current progress so an interrupted scan can be +// resumed. Failures are reported through the failure handlers rather than +// aborting the scan: losing a checkpoint is not worth losing scan results. +func (s *Scanner) saveCheckpoint() { + if s.opts.Checkpoint.Path == "" || s.runner == nil || s.runner.ScanProgress == nil { + return + } + // The task total comes from the scanner's own guarded copy rather than + // engine options.Count, which the engine writes unsynchronised while + // scheduling. + err := s.runner.ScanProgress.AtomicSave( + s.opts.Checkpoint.Path, + atomic.LoadUint32(&s.internal.CurrentCount), + uint32(s.totalScans()), + ) + if err == nil { + return + } + for _, fn := range s.opts.handlers.failure { + fn(Failure{Err: fmt.Errorf("afrog/sdk: checkpoint save failed: %w", err)}) + } +} + +func (s *Scanner) emitFinalProgress() { + if s.internal.OnPhaseProgress == nil { + return + } + completed, total, percent := s.vulnProgress() + + status := "completed" + if s.ctx.Err() != nil || (total > 0 && completed < total) { + status = "interrupted" + } + if status == "completed" { + percent = 100 + if total > 0 { + completed = total + } + } + s.internal.OnPhaseProgress(PhaseVuln, status, completed, total, percent) +} + +// Wait blocks until the scan finishes and returns its error. +// +// Cancelling ctx makes Wait return ctx.Err() without stopping the scan; call +// Stop for that. +func (s *Scanner) Wait(ctx context.Context) error { + if s.state.Load() == stateIdle { + return ErrNotStarted + } + if ctx == nil { + ctx = context.Background() + } + select { + case <-s.runDone: + return s.Err() + case <-ctx.Done(): + return ctx.Err() + } +} + +// Done returns a channel closed when the scan finishes. +func (s *Scanner) Done() <-chan struct{} { return s.runDone } + +// Err returns the scan error, or nil while the scan is still running. +func (s *Scanner) Err() error { + s.runErrMu.Lock() + defer s.runErrMu.Unlock() + return s.runErr +} + +// Stop asks the scan to stop and returns immediately. It is safe to call +// concurrently and more than once. Use Wait or Close to block until the scan +// has actually stopped. +func (s *Scanner) Stop() { + s.stopping.Store(true) + s.cancel() + if s.runner != nil { + s.runner.Stop() + } +} + +// Close stops the scan, waits for the scan goroutine to exit and releases +// every resource held by the scanner. +// +// Close is idempotent and safe to defer immediately after New. +func (s *Scanner) Close() error { + if s.closed.Swap(true) { + return nil + } + s.Stop() + + // Only wait when a scan was actually started, otherwise runDone never + // closes and Close would block forever. + if s.state.Load() != stateIdle { + <-s.runDone + } + + s.closeStreams() + + if s.runner != nil { + s.runner.Release() + } + s.restoreEnv() + + s.resultsMu.Lock() + s.results = nil + s.resultsMu.Unlock() + + return nil +} + +// restoreEnv undoes the process-wide environment changes made when the curated +// PoC source was mounted. +func (s *Scanner) restoreEnv() { + for _, e := range s.curatedEnv { + if e.set { + _ = os.Setenv(e.key, e.value) + } else { + _ = os.Unsetenv(e.key) + } + } + s.curatedEnv = nil +} + +func toSet(values []string) map[string]struct{} { + if len(values) == 0 { + return nil + } + out := make(map[string]struct{}, len(values)) + for _, v := range values { + out[v] = struct{}{} + } + return out +} + +func (s *Scanner) closeStreams() { + s.resultEvents.close() + s.portEvents.close() + s.hostEvents.close() + s.webProbeEvents.close() + s.progressEvents.close() + s.scanInfoEvents.close() +} + +// Pause suspends task scheduling. +func (s *Scanner) Pause() { + if s.runner != nil { + s.runner.Pause() + } +} + +// Resume resumes a paused scan. +func (s *Scanner) Resume() { + if s.runner != nil { + s.runner.Resume() + } +} + +// IsPaused reports whether the scan is paused. +func (s *Scanner) IsPaused() bool { + return s.runner != nil && s.runner.IsPaused() +} + +// IsStopping reports whether Stop has been called. +func (s *Scanner) IsStopping() bool { return s.stopping.Load() } + +// IsRunning reports whether a scan is currently in progress. +func (s *Scanner) IsRunning() bool { return s.state.Load() == stateRunning } + +// ------------------------------------------------------------------ streams + +// ResultStream returns a channel of findings, closed when the scan finishes. +// +// Nothing is published until the first call, so an unused stream costs +// nothing. Once subscribed the channel must be consumed: sends block when the +// buffer fills so that findings are never silently dropped. +func (s *Scanner) ResultStream() <-chan Result { return s.resultEvents.subscribe() } + +// PortStream returns a channel of open ports found during the port pre-scan. +// The same subscription rules as ResultStream apply. +func (s *Scanner) PortStream() <-chan PortEvent { return s.portEvents.subscribe() } + +// HostStream returns a channel of hosts found during discovery. +func (s *Scanner) HostStream() <-chan HostEvent { return s.hostEvents.subscribe() } + +// WebProbeStream returns a channel of probed web services. +func (s *Scanner) WebProbeStream() <-chan WebProbeEvent { return s.webProbeEvents.subscribe() } + +// ProgressStream returns a channel of phase progress updates. +func (s *Scanner) ProgressStream() <-chan PhaseProgress { return s.progressEvents.subscribe() } + +// ScanInfoStream returns a channel of scan summary updates. +func (s *Scanner) ScanInfoStream() <-chan ScanInfo { return s.scanInfoEvents.subscribe() } + +// ------------------------------------------------------------------ results + +// Results returns a snapshot of the findings collected so far. +func (s *Scanner) Results() []Result { + s.resultsMu.Lock() + defer s.resultsMu.Unlock() + out := make([]Result, len(s.results)) + copy(out, s.results) + return out +} + +// ResultCount returns the number of findings, including any beyond +// MaxStoredResults. +func (s *Scanner) ResultCount() int { return int(s.found.Load()) } + +// HasResults reports whether any vulnerability was found. +func (s *Scanner) HasResults() bool { return s.found.Load() > 0 } + +// OpenPorts returns the open ports discovered by the port pre-scan. +func (s *Scanner) OpenPorts() map[string][]int { + s.portsMu.Lock() + defer s.portsMu.Unlock() + + out := make(map[string][]int, len(s.openPorts)) + for host, ports := range s.openPorts { + if len(ports) == 0 { + continue + } + list := make([]int, 0, len(ports)) + for p := range ports { + list = append(list, p) + } + out[host] = list + } + return out +} + +func (s *Scanner) recordOpenPort(host string, port int) { + s.portsMu.Lock() + defer s.portsMu.Unlock() + ports, ok := s.openPorts[host] + if !ok { + ports = make(map[int]struct{}) + s.openPorts[host] = ports + } + ports[port] = struct{}{} +} + +// -------------------------------------------------------------- diagnostics + +// Stats returns a snapshot of the scan counters. +func (s *Scanner) Stats() Stats { + s.statsMu.RLock() + out := s.stats + s.statsMu.RUnlock() + out.CompletedScans = s.completed.Load() + out.FoundVulns = s.found.Load() + return out +} + +func (s *Scanner) totalScans() int { + s.statsMu.RLock() + defer s.statsMu.RUnlock() + return s.stats.TotalScans +} + +// Progress returns overall scan progress in the range [0, 100], weighting the +// port pre-scan and web probe stages when they are enabled. +func (s *Scanner) Progress() float64 { + // TotalScans is an estimate made before execution, and the engine may run + // slightly fewer tasks than estimated (deduplication, protocol-specific + // skips). Reporting the raw ratio would leave a completed scan stuck just + // below 100, so a scan that finished without being interrupted is 100% + // by definition. + if s.state.Load() == stateFinished && !s.stopping.Load() && s.Err() == nil { + return 100 + } + + vuln := 100.0 + if total := s.totalScans(); total > 0 { + vuln = float64(s.completed.Load()) / float64(total) * 100 + } + + weightPort, weightWeb := 0.0, 0.0 + portStage, webStage := 100.0, 100.0 + + if s.opts.EnablePortSan { + weightPort = 0.2 + discovery := s.phasePercent(PhaseHostDiscovery) + portscan := s.phasePercent(PhasePortScan) + if discovery == 0 && portscan == 0 { + portStage = 0 + } else { + portStage = float64(discovery+portscan) / 2 + } + } + if s.opts.EnableWebProbe { + weightWeb = 0.2 + webStage = float64(s.phasePercent(PhaseWebProbe)) + } + + weightVuln := 1.0 - weightPort - weightWeb + if weightVuln < 0 { + weightVuln = 0 + } + + return clampFloat(weightPort*portStage + weightWeb*webStage + weightVuln*clampFloat(vuln)) +} + +func (s *Scanner) phasePercent(name string) int { + s.phasesMu.Lock() + defer s.phasesMu.Unlock() + pp, ok := s.phases[name] + if !ok { + return 0 + } + return clampPercent(pp.Percent) +} + +// Pocs returns the PoCs loaded for this scan. +func (s *Scanner) Pocs() []poc.Poc { + out := make([]poc.Poc, len(s.pocs)) + copy(out, s.pocs) + return out +} + +// PocCount returns how many PoCs were loaded. +func (s *Scanner) PocCount() int { return len(s.pocs) } + +// PocDiagnostics returns the PoCs skipped during loading and the reason for +// each. These used to be printed to the console and were invisible to callers. +func (s *Scanner) PocDiagnostics() []config.PocLoadError { + out := make([]config.PocLoadError, len(s.pocDiags)) + copy(out, s.pocDiags) + return out +} + +// CuratedError reports why the optional curated PoC source failed to mount. +// It is never fatal; nil means the source mounted or was disabled. +func (s *Scanner) CuratedError() error { return s.curatedErr } + +// Info returns a summary of the scan without writing anything to the console. +// +// When OOB is enabled this performs a live connectivity probe against the OOB +// service. +func (s *Scanner) Info() ScanInfo { + stats := s.Stats() + + list := make([]string, 0, stats.TotalTargets) + for _, t := range s.internal.Targets.List() { + if v, ok := t.(string); ok { + list = append(list, v) + } + } + + info := ScanInfo{ + TotalTargets: stats.TotalTargets, + TotalPocs: stats.TotalPocs, + TotalScans: stats.TotalScans, + Targets: list, + OOBStatus: "disabled", + } + if s.opts.OOB.Enabled { + info.OOBEnabled, info.OOBStatus = s.OOBStatus() + } + return info +} + +// OOBStatus reports whether out-of-band detection is usable, together with a +// human-readable description. It performs a live probe against the configured +// OOB service. +func (s *Scanner) OOBStatus() (bool, string) { + if !s.opts.OOB.Enabled { + return false, "disabled" + } + + adapter := strings.ToLower(strings.TrimSpace(s.opts.OOB.Adapter)) + if adapter == "" { + return false, "not configured" + } + if s.internal.OOBKey == "" && s.internal.OOBDomain == "" { + return false, fmt.Sprintf("%s (incomplete configuration)", adapter) + } + + // Serialise probes: concurrent callers would otherwise open redundant + // connections to the OOB service. + s.oobProbedMu.Lock() + defer s.oobProbedMu.Unlock() + + client, err := oobadapter.NewOOBAdapter(s.internal.OOB, &oobadapter.ConnectorParams{ + Key: s.internal.OOBKey, + Domain: s.internal.OOBDomain, + HTTPUrl: s.internal.OOBHttpUrl, + ApiUrl: s.internal.OOBApiUrl, + }) + if err != nil { + return false, fmt.Sprintf("%s (initialisation failed: %v)", adapter, err) + } + if !client.IsVaild() { + return false, fmt.Sprintf("%s (unreachable)", adapter) + } + return true, fmt.Sprintf("%s (ok)", adapter) +} + +// printSummary writes the scan summary to stdout, but only when the caller +// opted in with WithVerbose. The SDK is otherwise completely silent. +func (s *Scanner) printSummary() { + if !s.opts.Verbose { + return + } + info := s.Info() + + fmt.Printf("\n========== afrog scan ==========\n") + fmt.Printf("targets : %d\n", info.TotalTargets) + fmt.Printf("pocs : %d\n", info.TotalPocs) + fmt.Printf("tasks : %d\n", info.TotalScans) + + const preview = 3 + if len(info.Targets) <= 5 { + fmt.Printf("scope : %s\n", strings.Join(info.Targets, ", ")) + } else { + fmt.Printf("scope : %s ... (+%d)\n", strings.Join(info.Targets[:preview], ", "), len(info.Targets)-preview) + } + fmt.Printf("oob : %s\n", info.OOBStatus) + fmt.Printf("================================\n") +} + +// ------------------------------------------------------------------ helpers + +func clampPercent(v int) int { + if v < 0 { + return 0 + } + if v > 100 { + return 100 + } + return v +} + +func clampFloat(v float64) float64 { + if v < 0 { + return 0 + } + if v > 100 { + return 100 + } + return v +} + +// -------------------------------------------------- internal option plumbing + +// builtConfig is the outcome of translating SDK options into engine options. +// +// curatedErr is kept separate from the function's error result because the two +// have different severities: a curated mount failure degrades the PoC set but +// must not prevent the scan from running. +type builtConfig struct { + options *config.Options + curatedErr error + // envRestore undoes the process-wide environment changes made while + // mounting the curated PoC source. + envRestore []envEntry +} + +// buildInternalOptions converts the SDK options into the engine configuration. +func buildInternalOptions(opts *Options) (*builtConfig, error) { + internal := &config.Options{ + TargetsFile: opts.TargetsFile, + PocPaths: opts.PocPaths, + PocPathsOnly: opts.PocPathsOnly, + Search: opts.Search, + Severity: opts.Severity, + ExcludePocs: opts.ExcludePocs, + ExcludePocsFile: opts.ExcludePocsFile, + RateLimit: opts.RateLimit, + Concurrency: opts.Concurrency, + Retries: opts.Retries, + Timeout: opts.Timeout, + MaxHostError: opts.MaxHostError, + MaxRespBodySize: opts.MaxRespBodySize, + BruteMaxRequests: opts.BruteMaxRequests, + ReqLimitPerTarget: opts.ReqLimitPerTarget, + AutoReqLimit: opts.AutoReqLimit, + Polite: opts.Polite, + Balanced: opts.Balanced, + Aggressive: opts.Aggressive, + Smart: opts.Smart, + DefaultAccept: opts.DefaultAccept, + DisableFingerprint: opts.DisableFingerprint, + EnableWebProbe: opts.EnableWebProbe, + FingerprintFilterMode: opts.FingerprintFilterMode, + VulnerabilityScannerBreakpoint: opts.StopOnFirstMatch, + Proxy: opts.Proxy, + Header: opts.Headers, + Dingtalk: opts.Dingtalk, + Wecom: opts.Wecom, + + PortScan: opts.EnablePortSan, + PSPorts: opts.PortScan.Ports, + PSRateLimit: opts.PortScan.RateLimit, + PSTimeout: opts.PortScan.TimeoutMs, + PSRetries: opts.PortScan.Retries, + PSSkipDiscovery: opts.PortScan.SkipDiscovery, + PSS4Chunk: opts.PortScan.ChunkSize, + + OOBRateLimit: opts.OOB.RateLimit, + OOBConcurrency: opts.OOB.Concurrency, + OOBFinalizeTimeout: opts.OOB.FinalizeTimeout, + OOBPollInterval: opts.OOB.PollInterval, + OOBHitRetention: opts.OOB.HitRetention, + + TaskHardTimeoutSec: opts.TaskTimeout.HardSec, + TaskSmartTimeout: opts.TaskTimeout.Smart, + TaskTimeoutVisibleCapSec: opts.TaskTimeout.VisibleCapSec, + TaskTimeoutNetCapSec: opts.TaskTimeout.NetCapSec, + TaskTimeoutGoCapSec: opts.TaskTimeout.GoCapSec, + + Cyberspace: opts.Cyberspace.Engine, + Query: opts.Cyberspace.Query, + QueryCount: opts.Cyberspace.Count, + + MonitorTargets: opts.TargetPreProbe, + Resume: opts.Checkpoint.Path, + + PocExecutionDurationMonitor: opts.EnableMonitor, + PedmLogLimit: opts.Monitor.LogLimit, + PedmSlowThresholdSec: opts.Monitor.SlowThresholdSec, + PedmSlowLogLimit: opts.Monitor.SlowLogLimit, + PedmSummaryTop: opts.Monitor.SummaryTop, + PedmSummaryBy: opts.Monitor.SummaryBy, + + CuratedEnabled: opts.Curated.Enabled, + CuratedEndpoint: opts.Curated.Endpoint, + CuratedTimeout: opts.Curated.TimeoutSec, + CuratedForceUpdate: opts.Curated.ForceUpdate, + } + + internal.Target = append(internal.Target, opts.Targets...) + + // SDK mode: silent, no update checks, no report files. + internal.SDKMode = true + internal.Silent = true + internal.DisableUpdateCheck = true + internal.DisableOutputHtml = true + internal.Json = "" + internal.JsonAll = "" + internal.Output = "" + + // Read the user configuration without touching the filesystem. A library + // constructor must not create directories, write files, or rewrite an + // existing user config. + cfg, err := config.LoadConfigReadOnly("") + if err != nil { + cfg = &config.Config{} + } + internal.Config = cfg + + if v := strings.TrimSpace(opts.Curated.Enabled); v != "" { + cfg.Curated.Enabled = v + } + if v := strings.TrimSpace(opts.Curated.Endpoint); v != "" { + cfg.Curated.Endpoint = v + } + if opts.Curated.TimeoutSec > 0 { + cfg.Curated.TimeoutSec = opts.Curated.TimeoutSec + } + + if err := applyOOB(internal, opts, cfg); err != nil { + return nil, err + } + if err := validateWebhooks(internal); err != nil { + return nil, err + } + if inputs := opts.pocInputs(); len(inputs) > 0 { + if err := config.ValidatePocInputs(inputs); err != nil { + return nil, fmt.Errorf("%w: %v", ErrPocPathNotFound, err) + } + } + + var envRestore []envEntry + curatedErr := mountCurated(internal, &envRestore) + return &builtConfig{options: internal, curatedErr: curatedErr, envRestore: envRestore}, nil +} + +// applyOOB resolves the OOB configuration, falling back to the user config +// file when the caller did not supply credentials. +func applyOOB(internal *config.Options, opts *Options, cfg *config.Config) error { + if !opts.OOB.Enabled { + return nil + } + + internal.EnableOOB = true + internal.OOB = opts.OOB.Adapter + internal.OOBKey = opts.OOB.Key + internal.OOBDomain = opts.OOB.Domain + internal.OOBApiUrl = opts.OOB.ApiURL + internal.OOBHttpUrl = opts.OOB.HttpURL + + noCredentials := internal.OOBKey == "" && internal.OOBDomain == "" && + internal.OOBApiUrl == "" && internal.OOBHttpUrl == "" + if noCredentials { + reverse := cfg.Reverse + switch strings.ToLower(internal.OOB) { + case "ceyeio": + internal.OOBKey = reverse.Ceye.ApiKey + internal.OOBDomain = reverse.Ceye.Domain + case "dnslogcn": + internal.OOBDomain = reverse.Dnslogcn.Domain + case "alphalog": + internal.OOBDomain = reverse.Alphalog.Domain + internal.OOBApiUrl = reverse.Alphalog.ApiUrl + case "xray": + internal.OOBKey = reverse.Xray.XToken + internal.OOBDomain = reverse.Xray.Domain + internal.OOBApiUrl = reverse.Xray.ApiUrl + case "revsuit": + internal.OOBKey = reverse.Revsuit.Token + internal.OOBDomain = reverse.Revsuit.DnsDomain + internal.OOBApiUrl = reverse.Revsuit.ApiUrl + internal.OOBHttpUrl = reverse.Revsuit.HttpUrl + } + } + + if strings.EqualFold(internal.OOB, "dnslogcn") && internal.OOBDomain == "" { + internal.OOBDomain = "dnslog.cn" + } + return nil +} + +func validateWebhooks(internal *config.Options) error { + if internal.Config == nil { + return nil + } + empty := func(tokens []string) bool { + for _, t := range tokens { + if strings.TrimSpace(t) != "" { + return false + } + } + return true + } + if internal.Dingtalk && empty(internal.Config.Webhook.Dingtalk.Tokens) { + return fmt.Errorf("%w: dingtalk", ErrWebhookTokenRequired) + } + if internal.Wecom && empty(internal.Config.Webhook.Wecom.Tokens) { + return fmt.Errorf("%w: wecom", ErrWebhookTokenRequired) + } + return nil +} + +// Environment variables through which the PoC repository learns about the +// curated source. +const ( + envCuratedDisabled = "AFROG_CURATED_DISABLED" + envCuratedDir = "AFROG_POCS_CURATED_DIR" +) + +// setEnv changes an environment variable and records how to undo it. +func setEnv(restore *[]envEntry, key, value string) { + prev, ok := os.LookupEnv(key) + *restore = append(*restore, envEntry{key: key, value: prev, set: ok}) + if value == "" { + _ = os.Unsetenv(key) + return + } + _ = os.Setenv(key, value) +} + +// mountCurated mounts the optional curated PoC source. Failures are reported +// but never abort the scan. +// +// It communicates with the PoC repository through process-wide environment +// variables, so it records the previous values for Close to restore. Without +// that, a later scanner in the same process would inherit this one's curated +// directory and silently scan with the wrong PoC set. +func mountCurated(internal *config.Options, restore *[]envEntry) error { + cur := internal.Config.Curated + enabled := strings.ToLower(strings.TrimSpace(cur.Enabled)) + endpoint := strings.TrimSpace(cur.Endpoint) + + if enabled == "off" || enabled == "false" || enabled == "0" || endpoint == "" { + setEnv(restore, envCuratedDisabled, "1") + setEnv(restore, envCuratedDir, "") + return nil + } + setEnv(restore, envCuratedDisabled, "") + + svc := service.New(service.Config{ + Endpoint: endpoint, + Channel: strings.TrimSpace(cur.Channel), + LicenseKey: strings.TrimSpace(cur.LicenseKey), + NoUpdate: cur.AutoUpdate != nil && !*cur.AutoUpdate && !internal.CuratedForceUpdate, + ForceUpdate: internal.CuratedForceUpdate, + ClientVersion: config.Version, + }) + + ctx := context.Background() + var cancel context.CancelFunc + if cur.TimeoutSec > 0 { + ctx, cancel = context.WithTimeout(ctx, time.Duration(cur.TimeoutSec)*time.Second) + } else { + ctx, cancel = context.WithCancel(ctx) + } + defer cancel() + + dir, err := svc.Mount(ctx) + if err != nil { + return &CuratedMountError{Err: err} + } + if strings.TrimSpace(dir) != "" { + setEnv(restore, envCuratedDir, dir) + } + return nil +} + +// buildRunner initialises the HTTP client, resolves the target set and builds +// the scan engine. +func buildRunner(internal *config.Options) (*runner.Runner, error) { + if err := retryhttpclient.Init(&retryhttpclient.Options{ + Proxy: internal.Proxy, + Timeout: internal.Timeout, + Retries: internal.Retries, + MaxRespBodySize: internal.MaxRespBodySize, + ReqLimitPerTarget: internal.ReqLimitPerTarget, + DefaultAccept: internal.DefaultAccept, + }); err != nil { + return nil, err + } + + seen := make(map[string]struct{}) + appendTargets := func(raws []string) { + for _, raw := range raws { + t := strings.TrimSpace(raw) + if t == "" { + continue + } + if _, dup := seen[t]; dup { + continue + } + seen[t] = struct{}{} + internal.Targets.Append(t) + } + } + + appendTargets(internal.Target) + if internal.TargetsFile != "" { + lines, err := utils.ReadFileLineByLine(internal.TargetsFile) + if err != nil { + return nil, err + } + appendTargets(lines) + } + + // A cyberspace query resolves to targets inside NewRunner, so the target + // set can legitimately still be empty at this point. + fromCyberspace := strings.TrimSpace(internal.Cyberspace) != "" && strings.TrimSpace(internal.Query) != "" + if internal.Targets.Len() == 0 && !fromCyberspace { + return nil, ErrNoTargets + } + + for _, p := range internal.PocPaths { + if v := strings.TrimSpace(p); v != "" { + internal.PocsDirectory.Set(v) + } + } + + // NewRunner would otherwise re-append these to Targets. + internal.Target = nil + + r, err := runner.NewRunner(internal) + if err != nil { + return nil, err + } + // A search that matched nothing leaves the runner with no work to do. The + // CLI tolerates that; a library caller is better served by an error. + if internal.Targets.Len() == 0 { + r.Release() + return nil, ErrNoTargets + } + return r, nil +} diff --git a/pkg/sdk/sdk_test.go b/pkg/sdk/sdk_test.go new file mode 100644 index 000000000..7ee7eac4f --- /dev/null +++ b/pkg/sdk/sdk_test.go @@ -0,0 +1,720 @@ +package sdk + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/zan8in/afrog/v3/pkg/result" +) + +const magicToken = "AFROG_SDK_TEST_TOKEN" + +// newTestServer returns a server whose response matches the test PoC. +func newTestServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Server", "afrog-test/1.0") + w.Header().Set("X-Custom", "custom-value") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("afrog test" + magicToken + "")) + })) + t.Cleanup(srv.Close) + return srv +} + +// writePoc writes a PoC that matches newTestServer's response. +func writePoc(t *testing.T, dir, name, id string) string { + t.Helper() + path := filepath.Join(dir, name) + body := "id: " + id + ` +info: + name: ` + id + ` + author: afrog-test + severity: info + description: sdk test poc +rules: + r0: + request: + method: GET + path: /probe?q=1 + expression: response.status == 200 && response.body.bcontains(b"` + magicToken + `") +expression: r0() +` + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + t.Fatalf("write poc %s: %v", path, err) + } + return path +} + +// newTestScanner wires a scanner to a local server and a temporary PoC dir. +func newTestScanner(t *testing.T, extra ...Option) (*Scanner, *httptest.Server) { + t.Helper() + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "match.yaml", "sdk-test-match") + + options := append([]Option{ + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithTimeout(10), + }, extra...) + + scanner, err := New(context.Background(), options...) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + return scanner, srv +} + +// --- PoC input --------------------------------------------------------------- + +func TestScanner_PocInputForms(t *testing.T) { + srv := newTestServer(t) + dir := t.TempDir() + single := writePoc(t, dir, "one.yaml", "one") + writePoc(t, dir, "two.yaml", "two") + + tests := []struct { + name string + paths []string + want int + }{ + {name: "single file", paths: []string{single}, want: 1}, + {name: "directory", paths: []string{dir}, want: 2}, + {name: "glob", paths: []string{filepath.Join(dir, "*.yaml")}, want: 2}, + {name: "mixed inputs are merged and de-duplicated", paths: []string{single, dir}, want: 2}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(tt.paths...), + WithPocPathsOnly(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + defer scanner.Close() + + if got := scanner.PocCount(); got != tt.want { + t.Errorf("PocCount() = %d, want %d", got, tt.want) + } + }) + } +} + +func TestScanner_RejectsUnresolvablePocPath(t *testing.T) { + srv := newTestServer(t) + + _, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(filepath.Join(t.TempDir(), "missing")), + ) + if !errors.Is(err, ErrPocPathNotFound) { + t.Fatalf("error = %v, want ErrPocPathNotFound", err) + } +} + +func TestScanner_RequiresTargets(t *testing.T) { + dir := t.TempDir() + writePoc(t, dir, "a.yaml", "a") + + _, err := New(context.Background(), WithPocPaths(dir)) + if !errors.Is(err, ErrNoTargets) { + t.Fatalf("error = %v, want ErrNoTargets", err) + } +} + +func TestScanner_PocDiagnosticsExposeSkippedPocs(t *testing.T) { + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "good.yaml", "good") + + broken := filepath.Join(dir, "broken.yaml") + if err := os.WriteFile(broken, []byte("id: broken\ninfo: [not a mapping\n"), 0o644); err != nil { + t.Fatalf("write broken poc: %v", err) + } + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + defer scanner.Close() + + if len(scanner.PocDiagnostics()) == 0 { + t.Fatal("expected a diagnostic for the malformed PoC") + } + if scanner.PocCount() != 1 { + t.Errorf("PocCount() = %d, want 1", scanner.PocCount()) + } +} + +// --- full data output -------------------------------------------------------- + +func TestScanner_ResultsCarryFullRequestAndResponse(t *testing.T) { + scanner, srv := newTestScanner(t) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + results := scanner.Results() + if len(results) != 1 { + t.Fatalf("got %d results, want 1", len(results)) + } + + r := results[0] + if r.PocID != "sdk-test-match" { + t.Errorf("PocID = %q, want sdk-test-match", r.PocID) + } + if !strings.HasPrefix(r.FullTarget, srv.URL) { + t.Errorf("FullTarget = %q, want prefix %q", r.FullTarget, srv.URL) + } + if len(r.Exchanges) == 0 { + t.Fatal("expected at least one exchange") + } + + ex := r.Exchanges[0] + if !strings.Contains(ex.Request, "GET /probe?q=1") { + t.Errorf("raw request missing request line:\n%s", ex.Request) + } + if !strings.Contains(ex.Response, magicToken) { + t.Errorf("raw response missing body token:\n%s", ex.Response) + } + if ex.StatusCode != http.StatusOK { + t.Errorf("StatusCode = %d, want 200", ex.StatusCode) + } + if ex.Method != http.MethodGet { + t.Errorf("Method = %q, want GET", ex.Method) + } + if ex.ResponseHeaders["server"] != "afrog-test/1.0" { + t.Errorf("response header server = %q, want afrog-test/1.0", ex.ResponseHeaders["server"]) + } + if ex.BodyTruncated { + t.Error("BodyTruncated should be false for a small response") + } +} + +// Redaction has to survive the whole path from the engine to Results(), not +// just the Exchange helper: the point of the option is that a caller can log +// or persist the results without leaking credentials. +func TestScanner_RedactsCredentialsInScanResults(t *testing.T) { + const secret = "super-secret-credential" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Set-Cookie", "session="+secret) + w.Header().Set("X-Custom", "keep-me") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("" + magicToken + "")) + })) + t.Cleanup(srv.Close) + + dir := t.TempDir() + writePoc(t, dir, "redact.yaml", "sdk-test-redact") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithTimeout(10), + WithHeaders("Authorization: Bearer "+secret), + WithRedactedHeaders(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + results := scanner.Results() + if len(results) == 0 { + t.Fatal("no result to inspect") + } + + // Serialising is how these leak in practice, so assert on the encoded form. + blob, err := json.Marshal(results) + if err != nil { + t.Fatalf("marshal results: %v", err) + } + encoded := string(blob) + + if strings.Contains(encoded, secret) { + t.Errorf("serialised results still carry the credential:\n%s", encoded) + } + if !strings.Contains(encoded, redactedValue) { + t.Errorf("nothing was redacted:\n%s", encoded) + } + if !strings.Contains(encoded, "keep-me") { + t.Errorf("redaction removed a harmless header:\n%s", encoded) + } +} + +func TestScanner_ResultsAreJSONSerialisable(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + data, err := json.Marshal(scanner.Results()) + if err != nil { + t.Fatalf("json.Marshal: %v", err) + } + // Raw exchanges must survive as readable text, not base64 blobs. + if !strings.Contains(string(data), magicToken) { + t.Errorf("serialised results should contain the raw response body:\n%s", data) + } + + var round []Result + if err := json.Unmarshal(data, &round); err != nil { + t.Fatalf("json.Unmarshal: %v", err) + } + if len(round) != 1 { + t.Fatalf("round-tripped %d results, want 1", len(round)) + } +} + +func TestScanner_RequestResponseCanBeDisabled(t *testing.T) { + scanner, _ := newTestScanner(t, WithRequestResponse(false)) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + results := scanner.Results() + if len(results) != 1 { + t.Fatalf("got %d results, want 1", len(results)) + } + if len(results[0].Exchanges) != 0 { + t.Errorf("Exchanges = %d, want 0 when capture is disabled", len(results[0].Exchanges)) + } +} + +func TestScanner_MaxStoredResultsDoesNotSuppressHandlers(t *testing.T) { + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "a.yaml", "poc-a") + writePoc(t, dir, "b.yaml", "poc-b") + + var handled atomic.Int64 + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithMaxStoredResults(1), + WithResultHandler(func(Result) { handled.Add(1) }), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + defer scanner.Close() + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + if got := len(scanner.Results()); got != 1 { + t.Errorf("stored %d results, want 1", got) + } + if got := handled.Load(); got != 2 { + t.Errorf("handler fired %d times, want 2", got) + } + if got := scanner.ResultCount(); got != 2 { + t.Errorf("ResultCount() = %d, want 2", got) + } +} + +func TestScanner_RawResultHandlerReceivesEngineType(t *testing.T) { + var got atomic.Int64 + scanner, _ := newTestScanner(t, WithRawResultHandler(func(r *result.Result) { + if r != nil && len(r.AllPocResult) > 0 { + got.Add(1) + } + })) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + if got.Load() != 1 { + t.Fatalf("raw handler fired %d times with exchanges, want 1", got.Load()) + } +} + +// --- lifecycle --------------------------------------------------------------- + +func TestScanner_ExecuteIsSingleUse(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("first Execute: %v", err) + } + if err := scanner.Execute(context.Background()); !errors.Is(err, ErrAlreadyFinished) { + t.Fatalf("second Execute = %v, want ErrAlreadyFinished", err) + } +} + +func TestScanner_WaitBeforeStart(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Wait(context.Background()); !errors.Is(err, ErrNotStarted) { + t.Fatalf("Wait = %v, want ErrNotStarted", err) + } +} + +func TestScanner_DoneClosesAfterScan(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + + select { + case <-scanner.Done(): + case <-time.After(60 * time.Second): + t.Fatal("scan did not finish within the timeout") + } + if err := scanner.Err(); err != nil { + t.Errorf("Err() = %v, want nil", err) + } +} + +func TestScanner_ContextCancellationStopsScan(t *testing.T) { + scanner, _ := newTestScanner(t) + + ctx, cancel := context.WithCancel(context.Background()) + if err := scanner.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + cancel() + + select { + case <-scanner.Done(): + case <-time.After(60 * time.Second): + t.Fatal("cancelling the context did not stop the scan") + } + if !scanner.IsStopping() { + t.Error("IsStopping() should be true after the context was cancelled") + } +} + +func TestScanner_StopIsIndependentOfStopOnFirstMatch(t *testing.T) { + // StopOnFirstMatch used to share a field with the stop flag, so enabling it + // made IsStopping report true before the scan even began. + scanner, _ := newTestScanner(t, WithStopOnFirstMatch()) + + if scanner.IsStopping() { + t.Fatal("IsStopping() must be false before Stop is called") + } + scanner.Stop() + if !scanner.IsStopping() { + t.Fatal("IsStopping() must be true after Stop") + } +} + +func TestScanner_CloseIsIdempotent(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + if err := scanner.Close(); err != nil { + t.Fatalf("first Close: %v", err) + } + if err := scanner.Close(); err != nil { + t.Fatalf("second Close: %v", err) + } + if err := scanner.Start(context.Background()); !errors.Is(err, ErrClosed) { + t.Fatalf("Start after Close = %v, want ErrClosed", err) + } +} + +func TestScanner_CloseWithoutStartDoesNotBlock(t *testing.T) { + scanner, _ := newTestScanner(t) + + done := make(chan struct{}) + go func() { + defer close(done) + _ = scanner.Close() + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("Close blocked when no scan had been started") + } +} + +func TestScanner_MultipleResultHandlers(t *testing.T) { + var first, second atomic.Int64 + + scanner, _ := newTestScanner(t, + WithResultHandler(func(Result) { first.Add(1) }), + WithResultHandler(func(Result) { second.Add(1) }), + ) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + if first.Load() != 1 || second.Load() != 1 { + t.Fatalf("handlers fired %d and %d times, want 1 each", first.Load(), second.Load()) + } +} + +// --- streams ----------------------------------------------------------------- + +func TestScanner_ResultStreamDeliversEveryFinding(t *testing.T) { + scanner, _ := newTestScanner(t) + + results := scanner.ResultStream() + if err := scanner.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + + var ( + wg sync.WaitGroup + count int + ) + wg.Add(1) + go func() { + defer wg.Done() + for range results { + count++ + } + }() + + if err := scanner.Wait(context.Background()); err != nil { + t.Fatalf("Wait: %v", err) + } + wg.Wait() + + if count != 1 { + t.Fatalf("streamed %d results, want 1", count) + } +} + +// An unsubscribed stream must never stall the scan, which is the failure mode +// a naive blocking implementation would introduce. +func TestScanner_UnsubscribedStreamDoesNotBlockScan(t *testing.T) { + scanner, _ := newTestScanner(t) + + done := make(chan error, 1) + go func() { done <- scanner.Execute(context.Background()) }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Execute: %v", err) + } + case <-time.After(60 * time.Second): + t.Fatal("scan stalled even though no stream was subscribed") + } +} + +func TestScanner_StreamClosesSoRangeTerminates(t *testing.T) { + scanner, _ := newTestScanner(t) + + results := scanner.ResultStream() + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + drained := make(chan struct{}) + go func() { + defer close(drained) + for range results { + } + }() + + select { + case <-drained: + case <-time.After(10 * time.Second): + t.Fatal("result stream was not closed when the scan finished") + } +} + +func TestScanner_SubscribingAfterScanYieldsClosedStream(t *testing.T) { + scanner, _ := newTestScanner(t) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + drained := make(chan struct{}) + go func() { + defer close(drained) + for range scanner.ResultStream() { + } + }() + + select { + case <-drained: + case <-time.After(10 * time.Second): + t.Fatal("subscribing after the scan should yield an already-closed stream") + } +} + +// --- resource management ----------------------------------------------------- + +// The out-of-band poll loop used to outlive every completed scan, so a +// long-running host process accumulated one goroutine per scan. +func TestScanner_DoesNotLeakGoroutines(t *testing.T) { + // Warm up shared lazily-initialised state so it is not counted as a leak. + warm, _ := newTestScanner(t) + if err := warm.Execute(context.Background()); err != nil { + t.Fatalf("warm-up Execute: %v", err) + } + _ = warm.Close() + + before := waitGoroutines(t, 0) + + for i := 0; i < 3; i++ { + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "leak.yaml", "leak-check") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + if err := scanner.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + } + + after := waitGoroutines(t, before) + // A small tolerance covers runtime and net/http background workers. + if after > before+5 { + t.Fatalf("goroutine count grew from %d to %d across 3 scans", before, after) + } +} + +// waitGoroutines settles the scheduler and returns the goroutine count, +// returning early once it drops to target or below. +func waitGoroutines(t *testing.T, target int) int { + t.Helper() + var n int + for i := 0; i < 40; i++ { + runtime.GC() + time.Sleep(50 * time.Millisecond) + n = runtime.NumGoroutine() + if target > 0 && n <= target { + return n + } + } + return n +} + +// --- output hygiene ---------------------------------------------------------- + +func TestScanner_IsSilentByDefault(t *testing.T) { + scanner, _ := newTestScanner(t) + + stdout, r, w := captureStdout(t) + execErr := scanner.Execute(context.Background()) + output := restoreStdout(t, stdout, r, w) + + if execErr != nil { + t.Fatalf("Execute: %v", execErr) + } + if output != "" { + t.Errorf("SDK wrote to stdout by default:\n%s", output) + } +} + +func TestScanner_VerbosePrintsSummary(t *testing.T) { + scanner, _ := newTestScanner(t, WithVerbose()) + + stdout, r, w := captureStdout(t) + execErr := scanner.Execute(context.Background()) + output := restoreStdout(t, stdout, r, w) + + if execErr != nil { + t.Fatalf("Execute: %v", execErr) + } + if !strings.Contains(output, "afrog scan") { + t.Errorf("WithVerbose should print a summary, got:\n%s", output) + } +} + +func captureStdout(t *testing.T) (*os.File, *os.File, *os.File) { + t.Helper() + stdout := os.Stdout + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe: %v", err) + } + os.Stdout = w + return stdout, r, w +} + +func restoreStdout(t *testing.T, stdout, r, w *os.File) string { + t.Helper() + os.Stdout = stdout + _ = w.Close() + + var sb strings.Builder + done := make(chan struct{}) + go func() { + defer close(done) + buf := make([]byte, 4096) + for { + n, err := r.Read(buf) + if n > 0 { + sb.Write(buf[:n]) + } + if err != nil { + return + } + } + }() + <-done + _ = r.Close() + return sb.String() +} + +func TestScanner_InfoReportsCountsWithoutPrinting(t *testing.T) { + scanner, _ := newTestScanner(t) + + info := scanner.Info() + if info.TotalTargets != 1 { + t.Errorf("TotalTargets = %d, want 1", info.TotalTargets) + } + if info.TotalPocs != 1 { + t.Errorf("TotalPocs = %d, want 1", info.TotalPocs) + } + if info.OOBEnabled { + t.Error("OOBEnabled should be false when OOB is not configured") + } +} diff --git a/pkg/sdk/stream.go b/pkg/sdk/stream.go new file mode 100644 index 000000000..5bd579f95 --- /dev/null +++ b/pkg/sdk/stream.go @@ -0,0 +1,99 @@ +package sdk + +import ( + "context" + "sync" +) + +// stream is a lazily created, single-close broadcast channel. +// +// Two properties matter here: +// +// - Nothing is allocated and nothing is sent until a caller subscribes, so an +// unsubscribed stream can never block the scan or leak a goroutine. +// - send holds a read lock and close holds the write lock, so close waits for +// in-flight sends instead of racing them into a send-on-closed panic. +// +// Sends are blocking by design. Silently discarding a vulnerability is not an +// acceptable failure mode for a scanner, so a subscriber that stops consuming +// applies backpressure to the scan rather than losing findings. The ctx guard +// keeps that from becoming a permanent stall: cancelling the scanner releases +// any blocked send. +type stream[T any] struct { + buf int + + mu sync.RWMutex + ch chan T + closed bool +} + +func newStream[T any](buf int) *stream[T] { + if buf <= 0 { + buf = DefaultStreamBuffer + } + return &stream[T]{buf: buf} +} + +// subscribe returns the receive side of the stream, creating it on first use. +// Subscribing after the stream closed yields an already-closed channel so that +// a range loop terminates immediately instead of blocking forever. +func (s *stream[T]) subscribe() <-chan T { + s.mu.Lock() + defer s.mu.Unlock() + if s.ch == nil { + s.ch = make(chan T, s.buf) + if s.closed { + close(s.ch) + } + } + return s.ch +} + +// send delivers a value to the subscriber, if any. +func (s *stream[T]) send(ctx context.Context, v T) { + s.mu.RLock() + defer s.mu.RUnlock() + if s.ch == nil || s.closed { + return + } + select { + case s.ch <- v: + case <-ctx.Done(): + } +} + +// close closes the stream. It is safe to call more than once. +func (s *stream[T]) close() { + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return + } + s.closed = true + if s.ch != nil { + close(s.ch) + } +} + +// emitter fans one event out to the registered handlers and the subscription +// stream. Handlers run before the stream send: a handler is the reliable +// delivery path and must not be delayed by a slow (or absent) stream consumer. +type emitter[T any] struct { + stream *stream[T] + handlers []func(T) +} + +func newEmitter[T any](buf int, handlers []func(T)) *emitter[T] { + return &emitter[T]{stream: newStream[T](buf), handlers: handlers} +} + +func (e *emitter[T]) emit(ctx context.Context, v T) { + for _, fn := range e.handlers { + fn(v) + } + e.stream.send(ctx, v) +} + +func (e *emitter[T]) subscribe() <-chan T { return e.stream.subscribe() } + +func (e *emitter[T]) close() { e.stream.close() } diff --git a/pkg/sdk/targetsource_test.go b/pkg/sdk/targetsource_test.go new file mode 100644 index 000000000..db46750ed --- /dev/null +++ b/pkg/sdk/targetsource_test.go @@ -0,0 +1,167 @@ +package sdk + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + "time" +) + +// --- checkpoint / resume ----------------------------------------------------- + +// The engine only reads a checkpoint; writing it is the embedder's job. If the +// SDK does not persist one, WithCheckpoint would look configured but resume +// nothing. +func TestScanner_CheckpointIsWrittenAndSkipsFinishedWork(t *testing.T) { + srv := newTestServer(t) + pocDir := t.TempDir() + writePoc(t, pocDir, "match.yaml", "sdk-test-checkpoint") + path := filepath.Join(t.TempDir(), "scan.afg") + + scan := func(checkpoint string) *Scanner { + t.Helper() + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(pocDir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithTimeout(10), + WithCheckpoint(CheckpointOptions{Path: checkpoint}), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + return scanner + } + + first := scan(path) + if got := first.ResultCount(); got != 1 { + t.Fatalf("first scan found %d results, want 1", got) + } + + info, err := os.Stat(path) + if err != nil { + t.Fatalf("checkpoint was not written: %v", err) + } + if info.Size() == 0 { + t.Fatal("checkpoint file is empty") + } + + // The same scan again must skip the work recorded in the checkpoint. + second := scan(path) + if got := second.ResultCount(); got != 0 { + t.Errorf("second scan found %d results, want 0 because the checkpoint marked the PoC done", got) + } + if got := second.Stats().CompletedScans; got != 0 { + t.Errorf("second scan ran %d tasks, want 0", got) + } + + // Proves the skip above came from the checkpoint rather than from the + // target or PoC going stale between runs. + fresh := scan(filepath.Join(t.TempDir(), "fresh.afg")) + if got := fresh.ResultCount(); got != 1 { + t.Errorf("scan with a fresh checkpoint found %d results, want 1", got) + } +} + +// A checkpoint path pointing at a file that does not exist yet is the normal +// first run, not an error. +func TestScanner_CheckpointAcceptsMissingFile(t *testing.T) { + scanner, _ := newTestScanner(t, WithCheckpoint(CheckpointOptions{ + Path: filepath.Join(t.TempDir(), "does-not-exist.afg"), + SaveInterval: 50 * time.Millisecond, + })) + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + if got := scanner.ResultCount(); got != 1 { + t.Errorf("found %d results, want 1", got) + } +} + +func TestWithCheckpoint_Validation(t *testing.T) { + if err := WithCheckpoint(CheckpointOptions{})(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("empty path = %v, want ErrInvalidOptions", err) + } + if err := WithCheckpoint(CheckpointOptions{Path: "x.afg", SaveInterval: -time.Second})(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("negative interval = %v, want ErrInvalidOptions", err) + } + + o := NewOptions() + if err := WithCheckpoint(CheckpointOptions{Path: "x.afg"})(o); err != nil { + t.Fatalf("WithCheckpoint: %v", err) + } + if o.Checkpoint.SaveInterval != DefaultCheckpointSaveInterval { + t.Errorf("SaveInterval = %v, want the default %v", o.Checkpoint.SaveInterval, DefaultCheckpointSaveInterval) + } +} + +// --- cyberspace -------------------------------------------------------------- + +func TestWithCyberspace_Validation(t *testing.T) { + tests := []struct { + name string + cfg CyberspaceOptions + }{ + {"empty engine", CyberspaceOptions{Query: "app:tomcat"}}, + {"unsupported engine", CyberspaceOptions{Engine: "fofa", Query: "app:tomcat"}}, + {"empty query", CyberspaceOptions{Engine: CyberspaceZoomEye}}, + {"negative count", CyberspaceOptions{Engine: CyberspaceZoomEye, Query: "x", Count: -1}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := WithCyberspace(tt.cfg)(NewOptions()); !errors.Is(err, ErrInvalidOptions) { + t.Errorf("got %v, want ErrInvalidOptions", err) + } + }) + } +} + +// A search query is a target source, so it must satisfy the "targets are +// required" rule on its own. +func TestOptions_CyberspaceQuerySatisfiesTargetRequirement(t *testing.T) { + o := NewOptions() + if err := o.validate(); !errors.Is(err, ErrNoTargets) { + t.Fatalf("bare options = %v, want ErrNoTargets", err) + } + + o2 := NewOptions() + if err := WithCyberspace(CyberspaceOptions{Engine: CyberspaceZoomEye, Query: `app:"tomcat"`})(o2); err != nil { + t.Fatalf("WithCyberspace: %v", err) + } + if err := o2.validate(); err != nil { + t.Fatalf("validate with a cyberspace query = %v, want nil", err) + } +} + +// --- target pre-probe -------------------------------------------------------- + +// The pre-probe only starts from Runner.Run. The SDK used to call Execute +// directly, so the option would have been silently inert. +func TestScanner_TargetPreProbeReachesEngineAndStaysSilent(t *testing.T) { + scanner, _ := newTestScanner(t, WithTargetPreProbe()) + + if !scanner.internal.MonitorTargets { + t.Fatal("TargetPreProbe did not reach the engine options") + } + + stdout, r, w := captureStdout(t) + execErr := scanner.Execute(context.Background()) + output := restoreStdout(t, stdout, r, w) + + if execErr != nil { + t.Fatalf("Execute: %v", execErr) + } + if output != "" { + t.Fatalf("pre-probe wrote to stdout:\n%s", output) + } + if got := scanner.ResultCount(); got != 1 { + t.Errorf("found %d results, want 1", got) + } +} diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go index 71f6f6a29..a8159c614 100644 --- a/pkg/utils/utils.go +++ b/pkg/utils/utils.go @@ -3,7 +3,6 @@ package utils import ( "encoding/hex" "io" - "log" "net/http" "net/url" "os" @@ -180,15 +179,17 @@ func GetNumberText(number int) string { return num } -// 16进制解码 +// HexDecode 解码十六进制字符串，输入非法时返回 nil。 +// +// 这里刻意不使用 log.Fatal：作为库函数，它不能因为一个畸形的 PoC 输入 +// 就终止宿主进程。 func HexDecode(s string) []byte { - dst := make([]byte, hex.DecodedLen(len(s))) //申请一个切片, 指明大小. 必须使用hex.DecodedLen - n, err := hex.Decode(dst, []byte(s)) //进制转换, src->dst + dst := make([]byte, hex.DecodedLen(len(s))) + n, err := hex.Decode(dst, []byte(s)) if err != nil { - log.Fatal(err) return nil } - return dst[:n] //返回0:n的数据. + return dst[:n] } // 字符串转为16进制 diff --git a/pkg/web/handlers.go b/pkg/web/handlers.go index 1928d944b..e0004c384 100644 --- a/pkg/web/handlers.go +++ b/pkg/web/handlers.go @@ -180,7 +180,7 @@ func serverInfoHandler(w http.ResponseWriter, r *http.Request) { activeTaskCount := 0 m.mu.Lock() for _, t := range m.tasks { - if t.Status == TaskRunning || t.Status == TaskPaused || t.Status == TaskStarting { + if isActive(t.Status()) { activeTaskCount++ } } @@ -206,7 +206,7 @@ func instancesListHandler(w http.ResponseWriter, r *http.Request) { active := make([]string, 0, 16) m.mu.Lock() for id, t := range m.tasks { - if t.Status == TaskRunning || t.Status == TaskPaused || t.Status == TaskStarting { + if isActive(t.Status()) { active = append(active, id) } } @@ -262,37 +262,26 @@ func instanceForceStopHandler(w http.ResponseWriter, r *http.Request) { return } m := getTaskManager() - found := false m.mu.Lock() - for id, t := range m.tasks { - if id == req.TaskID { - // 仅允许关闭属于该实例的活跃任务 - if t.Status == TaskRunning || t.Status == TaskPaused || t.Status == TaskStarting { - found = true - } - break - } - } + t := m.tasks[req.TaskID] m.mu.Unlock() - if !found { + + // 仅允许关闭属于该实例的活跃任务 + if t == nil || !isActive(t.Status()) { w.WriteHeader(http.StatusBadRequest) _ = json.NewEncoder(w).Encode(APIResponse{Success: false, Message: "taskId 不属于该实例或非活跃任务"}) return } - m.mu.Lock() - t := m.tasks[req.TaskID] - m.mu.Unlock() - if t != nil { - t.Scanner.Stop() - t.Status = TaskCancelled - publish(t, ScanEvent{Type: "status", Data: map[string]string{"status": string(TaskCancelled)}}) - finalizeTask(m, t, TaskCancelled) - if t.Scanner.IsStopping() { - gologger.Debug().Str("taskId", req.TaskID).Str("instanceId", instanceID).Msg("force-stop succeeded: task cancelled and server shutting down") - } else { - gologger.Debug().Str("taskId", req.TaskID).Str("instanceId", instanceID).Msg("force-stop uncertain: cancel flag not set") - } + t.Scanner.Stop() + t.setStatus(TaskCancelled) + publish(t, ScanEvent{Type: "status", Data: map[string]string{"status": string(TaskCancelled)}}) + finalizeTask(m, t, TaskCancelled) + if t.Scanner.IsStopping() { + gologger.Debug().Str("taskId", req.TaskID).Str("instanceId", instanceID).Msg("force-stop succeeded: task cancelled and server shutting down") + } else { + gologger.Debug().Str("taskId", req.TaskID).Str("instanceId", instanceID).Msg("force-stop uncertain: cancel flag not set") } + go func() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() diff --git a/pkg/web/scans.go b/pkg/web/scans.go index de2860a5c..92e92374f 100644 --- a/pkg/web/scans.go +++ b/pkg/web/scans.go @@ -2,6 +2,7 @@ package web import ( "bufio" + "context" "encoding/json" "fmt" "net/http" @@ -10,13 +11,14 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/gorilla/mux" - afrog "github.com/zan8in/afrog/v3" "github.com/zan8in/afrog/v3/pkg/db/sqlite" "github.com/zan8in/afrog/v3/pkg/pocsrepo" "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/pkg/sdk" "github.com/zan8in/gologger" ) @@ -31,22 +33,60 @@ const ( TaskCancelled TaskStatus = "cancelled" ) +// isActive reports whether a task still occupies a slot in the task manager. +func isActive(s TaskStatus) bool { + return s == TaskRunning || s == TaskPaused || s == TaskStarting +} + type ScanEvent struct { Type string `json:"type"` Data interface{} `json:"data"` } +// Task tracks one scan. Its mutable fields are read and written from the HTTP +// handler goroutines and from the event-drain goroutine at the same time, so +// they are reached only through the accessors below. type Task struct { ID string Name string CreatedAt time.Time - Status TaskStatus - Scanner *afrog.SDKScanner + Scanner *sdk.Scanner SeverityStats map[string]int Subscribers map[chan ScanEvent]struct{} - mu sync.Mutex - startTime time.Time - lastProgress time.Time + + mu sync.Mutex + status TaskStatus + startTime time.Time + + // finalized makes finalizeTask run exactly once. Both the stop handler and + // the drain goroutine reach it when a scan is cancelled, and running it + // twice would decrement the manager's running count twice and let the + // queue admit more scans than maxRunning allows. + finalized atomic.Bool +} + +func (t *Task) Status() TaskStatus { + t.mu.Lock() + defer t.mu.Unlock() + return t.status +} + +func (t *Task) setStatus(s TaskStatus) { + t.mu.Lock() + t.status = s + t.mu.Unlock() +} + +func (t *Task) started() time.Time { + t.mu.Lock() + defer t.mu.Unlock() + return t.startTime +} + +func (t *Task) setStarted(at time.Time) { + t.mu.Lock() + t.startTime = at + t.mu.Unlock() } type TaskManager struct { @@ -141,23 +181,25 @@ func startTask(m *TaskManager, t *Task) { m.running++ m.mu.Unlock() - t.Status = TaskRunning - t.startTime = time.Now() + t.setStatus(TaskRunning) + t.setStarted(time.Now()) gologger.Debug().Msgf("start scan running: taskId=%s capacity available", t.ID) publish(t, ScanEvent{Type: "status", Data: map[string]string{"status": "running"}}) + // Subscribe before starting the scan so that no event is missed. + resultCh := t.Scanner.ResultStream() + portCh := t.Scanner.PortStream() + hostCh := t.Scanner.HostStream() + webProbeCh := t.Scanner.WebProbeStream() + phaseCh := t.Scanner.ProgressStream() + scanInfoCh := t.Scanner.ScanInfoStream() + go func() { ticker := time.NewTicker(1 * time.Second) defer ticker.Stop() - resultCh := t.Scanner.ResultChan - portCh := t.Scanner.PortChan - hostCh := t.Scanner.HostChan - webProbeCh := t.Scanner.WebProbeChan - phaseCh := t.Scanner.PhaseProgressChan - scanInfoCh := t.Scanner.ScanInfoChan for { if resultCh == nil && portCh == nil && hostCh == nil && webProbeCh == nil && phaseCh == nil && scanInfoCh == nil { - if t.Status != TaskCancelled { + if t.Status() != TaskCancelled { finalizeTask(m, t, TaskCompleted) } return @@ -168,20 +210,19 @@ func startTask(m *TaskManager, t *Task) { resultCh = nil continue } - sev := strings.ToLower(r.PocInfo.Info.Severity) + sev := strings.ToLower(r.Severity) if t.SeverityStats == nil { t.SeverityStats = make(map[string]int) } t.SeverityStats[sev]++ - _ = persistHit(t.ID, r) publish(t, ScanEvent{Type: "result", Data: map[string]interface{}{ "target": r.FullTarget, - "severity": r.PocInfo.Info.Severity, + "severity": r.Severity, "poc": map[string]string{ - "id": r.PocInfo.Id, - "name": r.PocInfo.Info.Name, + "id": r.PocID, + "name": r.PocName, }, - "message": fmt.Sprintf("命中 %s", r.PocInfo.Info.Severity), + "message": fmt.Sprintf("命中 %s", r.Severity), "ts": time.Now().UnixMilli(), }}) case pr, ok := <-portCh: @@ -247,34 +288,37 @@ func startTask(m *TaskManager, t *Task) { "ts": time.Now().UnixMilli(), }}) case <-ticker.C: - st := t.Scanner.GetStats() - prog := t.Scanner.GetProgress() + st := t.Scanner.Stats() + prog := t.Scanner.Progress() publish(t, ScanEvent{Type: "progress", Data: map[string]interface{}{ "percent": int(prog + 0.5), "finished": int(st.CompletedScans), "total": st.TotalScans, - "rate": calcRate(t.startTime, st.CompletedScans), - "elapsedMs": time.Since(t.startTime).Milliseconds(), + "rate": calcRate(t.started(), st.CompletedScans), + "elapsedMs": time.Since(t.started()).Milliseconds(), }}) } } }() - _ = t.Scanner.RunAsync() + _ = t.Scanner.Start(context.Background()) } func finalizeTask(m *TaskManager, t *Task, status TaskStatus) { - t.Status = status + if t.finalized.Swap(true) { + return + } + t.setStatus(status) if t.Scanner != nil { - st := t.Scanner.GetStats() - prog := t.Scanner.GetProgress() + st := t.Scanner.Stats() + prog := t.Scanner.Progress() publish(t, ScanEvent{Type: "progress", Data: map[string]interface{}{ "percent": int(prog + 0.5), "finished": int(st.CompletedScans), "total": st.TotalScans, - "rate": calcRate(t.startTime, st.CompletedScans), - "elapsedMs": time.Since(t.startTime).Milliseconds(), + "rate": calcRate(t.started(), st.CompletedScans), + "elapsedMs": time.Since(t.started()).Milliseconds(), }}) - oobEnabled, oobStatus := t.Scanner.GetOOBStatus() + oobEnabled, oobStatus := t.Scanner.OOBStatus() publish(t, ScanEvent{Type: "scan_info", Data: map[string]interface{}{ "total_targets": st.TotalTargets, "total_pocs": st.TotalPocs, @@ -285,20 +329,25 @@ func finalizeTask(m *TaskManager, t *Task, status TaskStatus) { }}) } publish(t, ScanEvent{Type: "status", Data: map[string]string{"status": string(status)}}) + + // Release the scanner's background goroutines. Without this a long-running + // server accumulates one engine and one OOB poller per finished task. + if t.Scanner != nil { + _ = t.Scanner.Close() + } + m.mu.Lock() if m.running > 0 { m.running-- } - var nextID string + var next *Task if len(m.queue) > 0 { - nextID = m.queue[0] + next = m.tasks[m.queue[0]] m.queue = m.queue[1:] } m.mu.Unlock() - if nextID != "" { - if nt, ok := m.tasks[nextID]; ok { - startTask(m, nt) - } + if next != nil { + startTask(m, next) } } @@ -398,56 +447,68 @@ func scansCreateHandler(w http.ResponseWriter, r *http.Request) { } } - sdkOpts := afrog.NewSDKOptions() - sdkOpts.Targets = targets - sdkOpts.PocFile = pocPath - sdkOpts.AppendPoc = appendPocs - if useIDs { - sdkOpts.Search = "" - sdkOpts.Severity = "" - } else { - sdkOpts.Search = strings.TrimSpace(req.Search) - sdkOpts.Severity = strings.TrimSpace(req.Severity) + taskID := nextTaskID(getTaskManager()) + + sdkOpts := []sdk.Option{ + sdk.WithTargets(targets...), + sdk.WithPocPaths(pocPath), + sdk.WithPocPathsOnly(), + // Persist the engine-level result so that the stored request and + // response keep the exact shape the reports and UI expect. + sdk.WithRawResultHandler(func(r *result.Result) { + _ = persistHit(taskID, r) + }), + } + if len(appendPocs) > 0 { + sdkOpts = append(sdkOpts, sdk.WithPocPaths(appendPocs...)) + } + if !useIDs { + sdkOpts = append(sdkOpts, + sdk.WithSearch(strings.TrimSpace(req.Search)), + sdk.WithSeverity(strings.TrimSpace(req.Severity)), + ) } if req.Concurrency > 0 { - sdkOpts.Concurrency = req.Concurrency + sdkOpts = append(sdkOpts, sdk.WithConcurrency(req.Concurrency)) } if req.RateLimit > 0 { - sdkOpts.RateLimit = req.RateLimit + sdkOpts = append(sdkOpts, sdk.WithRateLimit(req.RateLimit)) } if req.Timeout > 0 { - sdkOpts.Timeout = req.Timeout + sdkOpts = append(sdkOpts, sdk.WithTimeout(req.Timeout)) } - if req.Retries >= 0 { - sdkOpts.Retries = req.Retries + if req.Retries > 0 { + sdkOpts = append(sdkOpts, sdk.WithRetries(req.Retries)) } if req.MaxHostError > 0 { - sdkOpts.MaxHostError = req.MaxHostError + sdkOpts = append(sdkOpts, sdk.WithMaxHostError(req.MaxHostError)) } - if strings.TrimSpace(req.Proxy) != "" { - sdkOpts.Proxy = strings.TrimSpace(req.Proxy) + if v := strings.TrimSpace(req.Proxy); v != "" { + sdkOpts = append(sdkOpts, sdk.WithProxy(v)) } if req.Smart { - sdkOpts.Smart = true + sdkOpts = append(sdkOpts, sdk.WithSmartConcurrency()) } - - sdkOpts.EnableOOB = req.EnableOOB - sdkOpts.OOB = strings.TrimSpace(req.OOB) - sdkOpts.OOBKey = strings.TrimSpace(req.OOBKey) - sdkOpts.OOBDomain = strings.TrimSpace(req.OOBDomain) - sdkOpts.OOBApiUrl = strings.TrimSpace(req.OOBApiUrl) - sdkOpts.OOBHttpUrl = strings.TrimSpace(req.OOBHttpUrl) - sdkOpts.PortScan = req.PortScan || req.PortScanCompat - sdkOpts.PSSkipDiscovery = req.SkipHostDisc - if v := strings.TrimSpace(req.Ports); v != "" { - sdkOpts.PSPorts = v + if req.EnableOOB { + sdkOpts = append(sdkOpts, sdk.WithOOB(sdk.OOBOptions{ + Adapter: strings.TrimSpace(req.OOB), + Key: strings.TrimSpace(req.OOBKey), + Domain: strings.TrimSpace(req.OOBDomain), + ApiURL: strings.TrimSpace(req.OOBApiUrl), + HttpURL: strings.TrimSpace(req.OOBHttpUrl), + })) + } + if req.PortScan || req.PortScanCompat { + sdkOpts = append(sdkOpts, sdk.WithPortScan(sdk.PortScanOptions{ + Ports: strings.TrimSpace(req.Ports), + SkipDiscovery: req.SkipHostDisc, + })) } if req.WebProbe || req.WebFingerprint { - sdkOpts.EnableWebProbe = true + sdkOpts = append(sdkOpts, sdk.WithWebProbe()) } - sdkOpts.EnableStream = true - scanner, err := afrog.NewSDKScanner(sdkOpts) + scanner, err := sdk.New(context.Background(), sdkOpts...) if err != nil { w.WriteHeader(http.StatusBadRequest) gologger.Debug().Str("path", r.URL.Path).Str("error", err.Error()).Msg("start scan failed: create scanner error") @@ -456,8 +517,8 @@ func scansCreateHandler(w http.ResponseWriter, r *http.Request) { } m := getTaskManager() - id := nextTaskID(m) - t := &Task{ID: id, Name: strings.TrimSpace(req.TaskName), Status: TaskStarting, Scanner: scanner, CreatedAt: time.Now()} + id := taskID + t := &Task{ID: id, Name: strings.TrimSpace(req.TaskName), status: TaskStarting, Scanner: scanner, CreatedAt: time.Now()} m.mu.Lock() m.tasks[id] = t m.mu.Unlock() @@ -551,8 +612,8 @@ func scansCreateHandler(w http.ResponseWriter, r *http.Request) { startTask(m, t) // 获取扫描初始化信息 - stats := scanner.GetStats() - oobEnabled, oobStatus := scanner.GetOOBStatus() + stats := scanner.Stats() + oobEnabled, oobStatus := scanner.OOBStatus() // 获取扫描目标（截取前5个用于展示，与CLI保持一致） displayTargets := []string{} @@ -631,8 +692,9 @@ func scanEventsHandler(w http.ResponseWriter, r *http.Request) { fl.Flush() } - writeEvent(ScanEvent{Type: "status", Data: map[string]string{"status": string(t.Status)}}) - if t.Status == TaskCompleted || t.Status == TaskFailed || t.Status == TaskCancelled { + current := t.Status() + writeEvent(ScanEvent{Type: "status", Data: map[string]string{"status": string(current)}}) + if current == TaskCompleted || current == TaskFailed || current == TaskCancelled { return } @@ -683,15 +745,15 @@ func scanStatusHandler(w http.ResponseWriter, r *http.Request) { _ = json.NewEncoder(w).Encode(APIResponse{Success: false, Message: "任务不存在"}) return } - st := t.Scanner.GetStats() + st := t.Scanner.Stats() resp := ScanStatusData{ - Status: string(t.Status), + Status: string(t.Status()), Progress: ScanProgressData{ - Percent: int(t.Scanner.GetProgress() + 0.5), + Percent: int(t.Scanner.Progress() + 0.5), Finished: int(st.CompletedScans), Total: st.TotalScans, - Rate: calcRate(t.startTime, st.CompletedScans), - ElapsedMs: time.Since(t.startTime).Milliseconds(), + Rate: calcRate(t.started(), st.CompletedScans), + ElapsedMs: time.Since(t.started()).Milliseconds(), }, TaskID: taskID, InstanceID: serverInstanceID, @@ -730,7 +792,7 @@ func scanPauseHandler(w http.ResponseWriter, r *http.Request) { return } t.Scanner.Pause() - t.Status = TaskPaused + t.setStatus(TaskPaused) if t.Scanner.IsPaused() { gologger.Debug().Str("taskId", taskID).Msg("pause succeeded: engine gated") } else { @@ -767,7 +829,7 @@ func scanResumeHandler(w http.ResponseWriter, r *http.Request) { return } t.Scanner.Resume() - t.Status = TaskRunning + t.setStatus(TaskRunning) if !t.Scanner.IsPaused() { gologger.Debug().Str("taskId", taskID).Msg("resume succeeded: engine released") } else { @@ -808,17 +870,15 @@ func scanStopHandler(w http.ResponseWriter, r *http.Request) { } else { gologger.Debug().Str("taskId", taskID).Msg("stop uncertain: cancel flag not set") } - t.Status = TaskCancelled + t.setStatus(TaskCancelled) finalizeTask(m, t, TaskCancelled) _ = json.NewEncoder(w).Encode(APIResponse{Success: true, Message: "stopped", Data: map[string]bool{"stopped": true}}) } -func calcRate(start time.Time, completed int32) int { +func calcRate(start time.Time, completed int64) int { secs := time.Since(start).Seconds() if secs <= 0 { return 0 } return int(float64(completed) / secs) } - -// 使用 SDKScanner 自带的统计，已在 scanStatus/progress 中读取 diff --git a/pkg/web/scans_test.go b/pkg/web/scans_test.go new file mode 100644 index 000000000..7e2b809c4 --- /dev/null +++ b/pkg/web/scans_test.go @@ -0,0 +1,97 @@ +package web + +import ( + "sync" + "testing" + "time" +) + +// Cancelling a scan reaches finalizeTask from two directions at once: the stop +// handler calls it directly, and the drain goroutine calls it when the scanner +// closes its streams. Releasing the manager slot twice would let the queue +// admit more concurrent scans than maxRunning allows. +func TestFinalizeTask_ReleasesTheSlotExactlyOnce(t *testing.T) { + m := newTaskManager() + m.running = 3 + + task := &Task{ID: "t1", status: TaskRunning} + + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + finalizeTask(m, task, TaskCancelled) + }() + } + wg.Wait() + + m.mu.Lock() + running := m.running + m.mu.Unlock() + + if running != 2 { + t.Fatalf("running = %d after 16 concurrent finalize calls, want 2", running) + } + if got := task.Status(); got != TaskCancelled { + t.Fatalf("Status() = %q, want %q", got, TaskCancelled) + } +} + +// A second finalize for the same task must not pull another task off the +// queue: that would start it while the first promotion is still running. +func TestFinalizeTask_DuplicateCallDoesNotDrainTheQueue(t *testing.T) { + m := newTaskManager() + m.running = 2 + + queued := &Task{ID: "queued", status: TaskStarting} + m.tasks[queued.ID] = queued + m.queue = []string{queued.ID} + + done := &Task{ID: "done", status: TaskRunning} + done.finalized.Store(true) // stands in for an earlier finalize + finalizeTask(m, done, TaskCompleted) + + m.mu.Lock() + queueLen := len(m.queue) + m.mu.Unlock() + if queueLen != 1 { + t.Fatalf("queue length = %d after a duplicate finalize, want 1", queueLen) + } +} + +// The task fields are read and written from HTTP handlers and the drain +// goroutine at the same time. This test is meaningful under -race. +func TestTask_ConcurrentFieldAccessIsRaceFree(t *testing.T) { + task := &Task{ID: "t1", status: TaskStarting} + sub := addSubscriber(task) + go func() { + for range sub { + } + }() + + statuses := []TaskStatus{TaskRunning, TaskPaused, TaskCancelled, TaskCompleted} + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + i := i + wg.Add(4) + go func() { defer wg.Done(); task.setStatus(statuses[i%len(statuses)]) }() + go func() { defer wg.Done(); _ = isActive(task.Status()) }() + go func() { defer wg.Done(); task.setStarted(time.Now()); _ = task.started() }() + go func() { + defer wg.Done() + publish(task, ScanEvent{Type: "status", Data: map[string]string{"status": "running"}}) + }() + } + + waited := make(chan struct{}) + go func() { defer close(waited); wg.Wait() }() + + select { + case <-waited: + case <-time.After(30 * time.Second): + t.Fatal("concurrent task field access deadlocked") + } + removeSubscriber(task, sub) +} diff --git a/pkg/web/server_routes copy.go b/pkg/web/server_routes copy.go deleted file mode 100644 index bcfb5f16a..000000000 --- a/pkg/web/server_routes copy.go +++ /dev/null @@ -1,197 +0,0 @@ -package web - -// import ( -// "fmt" -// "io" -// "io/fs" -// "log" -// "net/http" -// "path/filepath" -// "strings" -// ) - -// func setupHandler() (http.Handler, error) { -// mux := http.NewServeMux() - -// // API 路由组 - 使用子路由器确保精确匹配 -// apiMux := http.NewServeMux() -// apiMux.HandleFunc("/login", loginRateLimitMiddleware(loginHandler)) -// apiMux.HandleFunc("/logout", jwtAuthMiddleware(logoutHandler)) -// apiMux.HandleFunc("/vulns", jwtAuthMiddleware(vulnsHandler)) -// apiMux.HandleFunc("/reports", jwtAuthMiddleware(reportsHandler)) -// apiMux.HandleFunc("/reports/detail/", jwtAuthMiddleware(reportsDetailHandler)) // 修改路径避免冲突 -// apiMux.HandleFunc("/reports/poc/", jwtAuthMiddleware(pocDetailHandler)) -// apiMux.HandleFunc("/pocs/stats", jwtAuthMiddleware(pocsStatsHandler)) -// apiMux.HandleFunc("/health", healthCheckHandler) - -// // 将 API 路由挂载到 /api/ 下，并应用 API 专用中间件 -// mux.Handle("/api/", http.StripPrefix("/api", apiMiddleware(apiMux))) - -// // 静态文件和 SPA 处理 -// buildRoot, err := fs.Sub(GetWebpathFS(), "webpath") -// if err != nil { -// return nil, fmt.Errorf("无法加载静态文件: %v", err) -// } - -// // 使用优化的 SPA Handler -// spaHandler := &spaHandler{ -// staticFS: buildRoot, -// indexPath: GetWebpathIndexPath(), -// } -// mux.Handle("/", spaHandler) - -// return secureHeadersMiddleware(mux), nil -// // return secureHeadersMiddleware(loggingMiddleware(mux)), nil -// } - -// // API 专用中间件 -// func apiMiddleware(next http.Handler) http.Handler { -// return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { -// // 确保 API 响应始终为 JSON -// w.Header().Set("Content-Type", "application/json") -// w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate") -// w.Header().Set("Pragma", "no-cache") - -// next.ServeHTTP(w, r) -// }) -// } - -// // 健康检查处理器 -// func healthCheckHandler(w http.ResponseWriter, r *http.Request) { -// w.Header().Set("Content-Type", "application/json") -// w.WriteHeader(http.StatusOK) -// w.Write([]byte(`{"status":"ok","service":"afrog-web"}`)) -// } - -// // 优化的 SPA Handler -// type spaHandler struct { -// staticFS fs.FS -// indexPath string -// } - -// func (h *spaHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { -// // 如果是 API 请求，直接返回 404（不应该到达这里） -// if strings.HasPrefix(r.URL.Path, "/api/") { -// http.NotFound(w, r) -// return -// } - -// // 清理路径 -// path := filepath.Clean(r.URL.Path) -// if path == "/" { -// path = "index.html" -// } else { -// path = strings.TrimPrefix(path, "/") -// } - -// // 尝试打开文件 -// file, err := h.staticFS.Open(path) -// if err != nil { -// // 文件不存在，检查是否为前端路由 -// if h.isFrontendRoute(r.URL.Path) { -// h.serveIndex(w, r) -// return -// } -// http.NotFound(w, r) -// return -// } -// defer file.Close() - -// // 检查是否为目录 -// stat, err := file.Stat() -// if err != nil { -// http.Error(w, "Unable to stat file", http.StatusInternalServerError) -// return -// } - -// // 如果是目录，返回 index.html 让前端路由处理 -// if stat.IsDir() { -// h.serveIndex(w, r) -// return -// } - -// // 设置缓存策略 -// h.setCacheHeaders(w, path) - -// // 确保文件实现了 io.ReadSeeker 接口 -// readSeeker, ok := file.(io.ReadSeeker) -// if !ok { -// http.Error(w, "File does not support seeking", http.StatusInternalServerError) -// return -// } - -// // 使用标准的文件服务器处理 -// http.ServeContent(w, r, path, stat.ModTime(), readSeeker) -// } - -// // 判断是否为前端路由 -// func (h *spaHandler) isFrontendRoute(path string) bool { -// // SvelteKit 的前端路由路径 -// frontendRoutes := []string{"/login", "/reports", "/pocs", "/docs"} -// for _, route := range frontendRoutes { -// if strings.HasPrefix(path, route) { -// return true -// } -// } -// return false -// } - -// // 设置差异化缓存策略 -// func (h *spaHandler) setCacheHeaders(w http.ResponseWriter, path string) { -// ext := filepath.Ext(path) -// switch ext { -// case ".html": -// // HTML 文件不缓存，确保路由更新 -// w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate") -// w.Header().Set("Pragma", "no-cache") -// w.Header().Set("Expires", "0") -// case ".js", ".css": -// // JS/CSS 文件长期缓存（SvelteKit 会生成带 hash 的文件名） -// w.Header().Set("Cache-Control", "public, max-age=31536000, immutable") -// case ".png", ".jpg", ".jpeg", ".gif", ".svg", ".ico", ".webp": -// // 图片资源长期缓存 -// w.Header().Set("Cache-Control", "public, max-age=31536000") -// case ".woff", ".woff2", ".ttf", ".eot": -// // 字体文件长期缓存 -// w.Header().Set("Cache-Control", "public, max-age=31536000") -// default: -// // 其他文件短期缓存 -// w.Header().Set("Cache-Control", "public, max-age=3600") -// } -// } - -// // serveIndex 提供 index.html 文件 -// func (h *spaHandler) serveIndex(w http.ResponseWriter, r *http.Request) { -// indexFile, err := h.staticFS.Open(h.indexPath) -// if err != nil { -// http.Error(w, "Index file not found", http.StatusNotFound) -// return -// } -// defer indexFile.Close() - -// stat, err := indexFile.Stat() -// if err != nil { -// http.Error(w, "Unable to stat index file", http.StatusInternalServerError) -// return -// } - -// readSeeker, ok := indexFile.(io.ReadSeeker) -// if !ok { -// http.Error(w, "Index file does not support seeking", http.StatusInternalServerError) -// return -// } - -// // HTML 不缓存 -// w.Header().Set("Content-Type", "text/html; charset=utf-8") -// w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate") -// w.Header().Set("Pragma", "no-cache") -// w.Header().Set("Expires", "0") -// http.ServeContent(w, r, h.indexPath, stat.ModTime(), readSeeker) -// } - -// func loggingMiddleware(next http.Handler) http.Handler { -// return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { -// log.Printf("Request: %s %s", r.Method, r.URL.Path) -// next.ServeHTTP(w, r) -// }) -// } From 52bc104eccddaec360b7dc0495bc834c7eba1c1a Mon Sep 17 00:00:00 2001 From: zhizhuo Date: Tue, 11 Aug 2026 20:55:02 +0800 Subject: [PATCH 2/5] =?UTF-8?q?=E8=A1=A5=E5=85=85=E5=BC=95=E6=93=8E?= =?UTF-8?q?=E5=B1=82=E4=BF=AE=E5=A4=8D=E7=9A=84=E6=B5=8B=E8=AF=95=E8=A6=86?= =?UTF-8?q?=E7=9B=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此前这几项修复通过了全量回归，但没有任何测试断言它们自身的正确性： 回归只能说明没有破坏既有功能，不能说明修复真的生效。 每个测试都通过临时回退对应修复验证过会失败，确认不是空测： - OOB 轮询协程回收：回退后 Stop 永远等不到协程退出，测试超时 - ticker/oobMgr 原子指针：并发 setTicker/waitTick/stopTicker/Stop 在 -race 下的专项测试 - ticker 零间隔：回退后 rate=2e9 触发 "non-positive interval for NewTicker" panic - OnFailure：回退后连接被拒的目标不产生任何失败上报 - BodyTruncated：补上真实截断为 true 的场景，此前只测了 false - IPv6 地址拼接：回退后报 "too many colons in address" - HexDecode：回退后 log.Fatal 直接终止测试进程 pkg/utils 此前没有任何测试文件，现已覆盖。 --- pkg/portscan/ipv6_test.go | 96 +++++++++++++++++++ pkg/runner/lifecycle_test.go | 181 +++++++++++++++++++++++++++++++++++ pkg/sdk/enginefix_test.go | 159 ++++++++++++++++++++++++++++++ pkg/utils/utils_test.go | 58 +++++++++++ 4 files changed, 494 insertions(+) create mode 100644 pkg/portscan/ipv6_test.go create mode 100644 pkg/runner/lifecycle_test.go create mode 100644 pkg/sdk/enginefix_test.go create mode 100644 pkg/utils/utils_test.go diff --git a/pkg/portscan/ipv6_test.go b/pkg/portscan/ipv6_test.go new file mode 100644 index 000000000..58a1b4ca8 --- /dev/null +++ b/pkg/portscan/ipv6_test.go @@ -0,0 +1,96 @@ +package portscan + +import ( + "net" + "strconv" + "testing" + "time" +) + +// checkPortOpen used to build its dial address with fmt.Sprintf("%s:%d"), +// which produces "::1:8080" for an IPv6 literal. net.Dial cannot parse that, +// so every IPv6 target failed regardless of whether the port was open. +func TestCheckPortOpen_IPv6Literal(t *testing.T) { + ln, err := net.Listen("tcp", "[::1]:0") + if err != nil { + t.Skipf("IPv6 loopback unavailable in this environment: %v", err) + } + defer ln.Close() + + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + _ = c.Close() + } + }() + + _, portStr, err := net.SplitHostPort(ln.Addr().String()) + if err != nil { + t.Fatalf("SplitHostPort(%q): %v", ln.Addr(), err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + t.Fatalf("Atoi(%q): %v", portStr, err) + } + + s := &Scanner{options: &Options{Timeout: 3 * time.Second, Retries: 0}} + + conn, err := s.checkPortOpen("::1", port) + if err != nil { + t.Fatalf("checkPortOpen on an open IPv6 port failed: %v", err) + } + _ = conn.Close() +} + +// The IPv4 path must be byte-for-byte what it was before the change. +func TestCheckPortOpen_IPv4StillWorks(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer ln.Close() + + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + _ = c.Close() + } + }() + + _, portStr, _ := net.SplitHostPort(ln.Addr().String()) + port, _ := strconv.Atoi(portStr) + + s := &Scanner{options: &Options{Timeout: 3 * time.Second, Retries: 0}} + + conn, err := s.checkPortOpen("127.0.0.1", port) + if err != nil { + t.Fatalf("checkPortOpen on an open IPv4 port failed: %v", err) + } + _ = conn.Close() +} + +// net.JoinHostPort is what makes the IPv6 case work; pin the difference so a +// future refactor back to Sprintf is caught here rather than in the field. +func TestDialAddressBracketsIPv6(t *testing.T) { + tests := []struct { + host string + port int + want string + }{ + {"127.0.0.1", 80, "127.0.0.1:80"}, + {"example.com", 8443, "example.com:8443"}, + {"::1", 80, "[::1]:80"}, + {"fe80::1", 22, "[fe80::1]:22"}, + } + for _, tt := range tests { + if got := net.JoinHostPort(tt.host, strconv.Itoa(tt.port)); got != tt.want { + t.Errorf("JoinHostPort(%q, %d) = %q, want %q", tt.host, tt.port, got, tt.want) + } + } +} diff --git a/pkg/runner/lifecycle_test.go b/pkg/runner/lifecycle_test.go new file mode 100644 index 000000000..ed5e62605 --- /dev/null +++ b/pkg/runner/lifecycle_test.go @@ -0,0 +1,181 @@ +package runner + +import ( + "context" + "runtime" + "sync" + "testing" + "time" + + "github.com/zan8in/afrog/v3/pkg/config" +) + +// settleGoroutines returns the goroutine count once it stops shrinking, so a +// just-signalled goroutine is not counted as still running. +func settleGoroutines(t *testing.T, target int) int { + t.Helper() + n := runtime.NumGoroutine() + for i := 0; i < 40; i++ { + runtime.GC() + time.Sleep(25 * time.Millisecond) + n = runtime.NumGoroutine() + if target > 0 && n <= target { + return n + } + } + return n +} + +// The poll loop used to outlive every completed scan, because the runner +// context is only cancelled by an explicit Stop. Stop must terminate it and +// wait for the goroutine to be gone, not merely signal it. +func TestOOBManager_StopTerminatesPollLoop(t *testing.T) { + before := settleGoroutines(t, 0) + + for i := 0; i < 5; i++ { + m := NewOOBManager(context.Background(), nil, 10*time.Millisecond, time.Minute) + m.Stop() + + // Stop waits for the loop to exit, so the goroutine must already be + // gone rather than merely scheduled to stop. + select { + case <-m.done: + default: + t.Fatal("Stop returned before the poll loop exited") + } + } + + if after := settleGoroutines(t, before); after > before+2 { + t.Fatalf("goroutine count grew from %d to %d across 5 manager lifecycles", before, after) + } +} + +// Stop must be safe to call repeatedly and from several goroutines at once. +func TestOOBManager_StopIsIdempotentAndConcurrencySafe(t *testing.T) { + m := NewOOBManager(context.Background(), nil, 10*time.Millisecond, time.Minute) + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { defer wg.Done(); m.Stop() }() + } + + done := make(chan struct{}) + go func() { defer close(done); wg.Wait() }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("concurrent Stop deadlocked") + } +} + +// A nil manager must be inert so callers need no nil checks. +func TestOOBManager_StopOnNilIsSafe(t *testing.T) { + var m *OOBManager + m.Stop() +} + +// Installing a manager has to stop the previous one, otherwise a second scan +// in the same process leaks the first scan's poller. +func TestEngine_SetOOBManagerStopsThePrevious(t *testing.T) { + e := &Engine{} + + first := NewOOBManager(context.Background(), nil, 10*time.Millisecond, time.Minute) + e.setOOBManager(first) + if e.OOBMgr() != first { + t.Fatal("setOOBManager did not install the manager") + } + + second := NewOOBManager(context.Background(), nil, 10*time.Millisecond, time.Minute) + e.setOOBManager(second) + + select { + case <-first.done: + case <-time.After(5 * time.Second): + t.Fatal("the replaced manager's poll loop is still running") + } + if e.OOBMgr() != second { + t.Fatal("setOOBManager did not install the replacement") + } + + e.stopOOBManager() + select { + case <-second.done: + case <-time.After(5 * time.Second): + t.Fatal("stopOOBManager left the poll loop running") + } + if e.OOBMgr() != nil { + t.Fatal("stopOOBManager did not clear the reference") + } +} + +// stopOOBManager runs from Stop, from Release and from Execute's defer, so it +// must tolerate being called with nothing installed and more than once. +func TestEngine_StopOOBManagerIsSafeWhenIdle(t *testing.T) { + var nilEngine *Engine + nilEngine.stopOOBManager() + if nilEngine.OOBMgr() != nil { + t.Fatal("a nil engine should report no manager") + } + + e := &Engine{} + e.stopOOBManager() + e.stopOOBManager() +} + +// The ticker and the OOB manager are written while scheduling and read from +// Stop on another goroutine. Before they became atomic this raced; the test is +// meaningful under -race. +func TestEngine_ConcurrentTickerAndStopIsRaceFree(t *testing.T) { + e := NewEngine(&config.Options{}) + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(3) + go func() { + defer wg.Done() + e.setTicker(time.NewTicker(time.Millisecond)) + }() + go func() { + defer wg.Done() + e.waitTick() + }() + go func() { + defer wg.Done() + e.stopTicker() + }() + } + + // Stop closes e.quit, which releases any waitTick blocked on the ticker. + wg.Add(1) + go func() { defer wg.Done(); e.Stop() }() + + done := make(chan struct{}) + go func() { defer close(done); wg.Wait() }() + + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("concurrent ticker access deadlocked") + } + + // Stop must leave nothing installed for a later Release to trip over. + e.stopTicker() +} + +// Stop must also release the OOB poller: it is the path an interrupted scan +// takes, and the poller would otherwise survive the scan. +func TestEngine_StopReleasesTheOOBPoller(t *testing.T) { + e := NewEngine(&config.Options{}) + m := NewOOBManager(context.Background(), nil, 10*time.Millisecond, time.Minute) + e.setOOBManager(m) + + e.Stop() + + select { + case <-m.done: + case <-time.After(5 * time.Second): + t.Fatal("Stop left the OOB poll loop running") + } +} diff --git a/pkg/sdk/enginefix_test.go b/pkg/sdk/enginefix_test.go new file mode 100644 index 000000000..ecb2981e8 --- /dev/null +++ b/pkg/sdk/enginefix_test.go @@ -0,0 +1,159 @@ +package sdk + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync" + "testing" +) + +// time.Second/rate truncates to zero once the rate limit reaches one billion, +// and time.NewTicker panics on a non-positive interval. The panic surfaced as +// a crashed scan goroutine rather than an error. +func TestScanner_ExtremeRateLimitDoesNotPanic(t *testing.T) { + for _, rate := range []int{1_000_000_000, 2_000_000_000} { + t.Run(strconv.Itoa(rate), func(t *testing.T) { + scanner, _ := newTestScanner(t, WithRateLimit(rate)) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute with rate limit %d: %v", rate, err) + } + // A panicking scan goroutine is recovered into the run error, so a + // clean Execute plus a real finding proves the scan actually ran. + if got := scanner.ResultCount(); got != 1 { + t.Fatalf("found %d results at rate limit %d, want 1", got, rate) + } + }) + } +} + +// PoC execution failures used to be swallowed: a request error or a broken +// expression produced no result and no signal, so callers could not tell a +// clean "not vulnerable" from a PoC that never ran. +func TestScanner_FailureHandlerObservesRequestErrors(t *testing.T) { + // Bind then release a port so connections to it are refused outright. + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + addr := ln.Addr().String() + _ = ln.Close() + + dir := t.TempDir() + writePoc(t, dir, "fail.yaml", "sdk-test-failure") + + var mu sync.Mutex + var failures []Failure + + scanner, err := New(context.Background(), + WithTargets("http://"+addr), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithTimeout(3), + WithRetries(0), + WithFailureHandler(func(f Failure) { + mu.Lock() + failures = append(failures, f) + mu.Unlock() + }), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + mu.Lock() + defer mu.Unlock() + if len(failures) == 0 { + t.Fatal("no failure reported for a target that refuses every connection") + } + f := failures[0] + if f.Err == nil { + t.Error("Failure.Err is nil") + } + if f.PocID != "sdk-test-failure" { + t.Errorf("Failure.PocID = %q, want sdk-test-failure", f.PocID) + } + if strings.TrimSpace(f.Error()) == "" { + t.Error("Failure.Error() is empty") + } +} + +// A response larger than MaxRespBodySize is cut short, and callers have to be +// able to tell that the body they received is not the whole server response. +func TestScanner_BodyTruncatedIsReportedOnLargeResponse(t *testing.T) { + const maxMB = 1 + // Comfortably past the 1 MB ceiling so the read stops mid-body. + payload := strings.Repeat("A", 3*1024*1024) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + // The token has to land before the cut so the PoC still matches. + _, _ = w.Write([]byte(magicToken)) + _, _ = w.Write([]byte(payload)) + })) + t.Cleanup(srv.Close) + + dir := t.TempDir() + writePoc(t, dir, "big.yaml", "sdk-test-truncate") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithTimeout(20), + WithMaxRespBodySize(maxMB), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + results := scanner.Results() + if len(results) == 0 { + t.Fatal("no result to inspect") + } + if len(results[0].Exchanges) == 0 { + t.Fatal("result carries no exchange") + } + + ex := results[0].Exchanges[0] + if !ex.BodyTruncated { + t.Errorf("BodyTruncated = false for a %d byte body read with a %d MB limit", + len(payload), maxMB) + } + if len(ex.ResponseBody) >= len(payload) { + t.Errorf("response body is %d bytes, so it was not actually truncated", len(ex.ResponseBody)) + } +} + +// The counterpart: a small response must not be flagged, otherwise the flag +// would be useless noise. +func TestScanner_BodyTruncatedIsFalseForSmallResponse(t *testing.T) { + scanner, _ := newTestScanner(t, WithMaxRespBodySize(2)) + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + results := scanner.Results() + if len(results) == 0 || len(results[0].Exchanges) == 0 { + t.Fatal("no exchange to inspect") + } + if results[0].Exchanges[0].BodyTruncated { + t.Error("BodyTruncated = true for a response well under the limit") + } +} diff --git a/pkg/utils/utils_test.go b/pkg/utils/utils_test.go new file mode 100644 index 000000000..63893ec3f --- /dev/null +++ b/pkg/utils/utils_test.go @@ -0,0 +1,58 @@ +package utils + +import ( + "bytes" + "testing" +) + +// HexDecode used to call log.Fatal on malformed input, which terminates the +// whole process. A library function must not kill its host just because one +// PoC carried a bad hex string, so it now returns nil instead. +func TestHexDecode(t *testing.T) { + tests := []struct { + name string + input string + want []byte + }{ + {"empty", "", []byte{}}, + {"lowercase", "48656c6c6f", []byte("Hello")}, + {"uppercase", "48656C6C6F", []byte("Hello")}, + {"zero byte", "00", []byte{0x00}}, + + // Every case below terminated the process before the fix. + {"odd length", "abc", nil}, + {"non-hex character", "zz", nil}, + {"trailing garbage", "48656c6c6fzz", nil}, + {"whitespace", "48 65", nil}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := HexDecode(tt.input) + if tt.want == nil { + if got != nil { + t.Fatalf("HexDecode(%q) = %v, want nil", tt.input, got) + } + return + } + if !bytes.Equal(got, tt.want) { + t.Fatalf("HexDecode(%q) = %v, want %v", tt.input, got, tt.want) + } + }) + } +} + +// A malformed string reaching HexDecode must leave the caller running so the +// scan can continue with the remaining PoCs. +func TestHexDecode_MalformedInputDoesNotAbort(t *testing.T) { + for _, s := range []string{"nothex", "1", "%%%%"} { + if got := HexDecode(s); got != nil { + t.Errorf("HexDecode(%q) = %v, want nil", s, got) + } + } + // Reaching this line at all is the assertion: the old implementation + // exited the test binary before it. + if got := HexDecode("41"); !bytes.Equal(got, []byte("A")) { + t.Fatalf("HexDecode still works after malformed input: got %v", got) + } +} From 6cc00397176fe869599efd478704886b48b16398 Mon Sep 17 00:00:00 2001 From: zhizhuo Date: Tue, 11 Aug 2026 22:38:22 +0800 Subject: [PATCH 3/5] =?UTF-8?q?=E6=81=A2=E5=A4=8D=E6=97=A7=20SDK=20?= =?UTF-8?q?=E5=85=BC=E5=AE=B9=E5=B1=82=EF=BC=8C=E4=B8=A4=E5=A5=97=20API=20?= =?UTF-8?q?=E5=85=B1=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 上一次提交把根目录的 afrog.go 删掉了，使用 afrog.NewSDKScanner 的 存量代码会直接编译不过。这次把旧接口原样恢复为根包的兼容门面， 内部委托给 pkg/sdk，因此只有一份实现，缺陷修复与新能力对两套 API 同时生效。 兼容性 - SDKOptions 的 51 个原始字段、默认值、语义保持不变 - NewSDKOptions / NewSDKScanner 与 19 个方法签名逐一对齐 - OnResult / OnPort / OnWebProbe 仍是构造后赋值的公开字段， 处理器在调用时读取而非构造时捕获 - 六个通道仍由 EnableStream 控制创建，并保留"满了就丢"的 非阻塞语义：改成阻塞会让不读通道的旧代码直接卡死 - SetProxy / SetRateLimit / SetConcurrency 在 Run 前生效， 底层扫描器按需重建 - 配置错误仍在 NewSDKScanner 阶段返回，而不是推迟到 Run 旧行为的唯一修正：同时指定 PocFile 与 AppendPoc 时， 旧版静默丢弃 AppendPoc，现在两者都会加载。 新能力以可选字段开放给旧写法：PocPaths/PocPathsOnly、ResumeFile、 TaskHardTimeoutSec/TaskSmartTimeout、Cyberspace/Query/QueryCount、 MonitorTargets、OOBPollInterval/OOBHitRetention、MaxStoredResults、 RedactedHeaders、OnFailure、Silent。留空即与以前完全一致。 另有 Scanner() 返回底层 *sdk.Scanner，便于逐步迁移。 验证 - 新增 15 个兼容测试，覆盖基础扫描、回调、流式、异步、setter、 统计、Close 幂等、构造期报错、通道不读不卡死 - 从 git 历史取出改动前的 7 个真实示例，未经修改全部编译通过 - 全量 15 个包在 -race 下通过 - 中英文文档补充两套 API 的对照与迁移说明 --- afrog.go | 906 ++++++++++++++++++ afrog_test.go | 448 +++++++++ docs/SDK_Usage_Guide_English.md | 36 + ...7\345\215\227_\344\270\255\346\226\207.md" | 36 + 4 files changed, 1426 insertions(+) create mode 100644 afrog.go create mode 100644 afrog_test.go diff --git a/afrog.go b/afrog.go new file mode 100644 index 000000000..f8525b505 --- /dev/null +++ b/afrog.go @@ -0,0 +1,906 @@ +// Package afrog is the backward-compatible facade over the afrog scanner. +// +// It preserves the original SDK surface — [NewSDKOptions], [NewSDKScanner] and +// the [SDKScanner] methods — so existing integrations keep compiling and +// behaving as before. Everything here delegates to [pkg/sdk], which is the +// current, context-aware API; new code should prefer that package directly. +// +// The facade exists so the two can coexist: one implementation, two entry +// points. Fixes and features added to pkg/sdk reach callers of this package +// automatically. +package afrog + +import ( + "context" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/zan8in/afrog/v3/pkg/protocols/http/retryhttpclient" + "github.com/zan8in/afrog/v3/pkg/result" + "github.com/zan8in/afrog/v3/pkg/sdk" +) + +// OOBAdapter reports whether an out-of-band service is reachable. +type OOBAdapter interface { + IsVaild() bool +} + +// PortScanResult is an open port found during the port pre-scan. +type PortScanResult struct { + Host string + Port int +} + +// HostDiscoveryResult is a live host found during discovery. +type HostDiscoveryResult struct { + Host string +} + +// WebProbeResult is the metadata of a probed web service. +type WebProbeResult struct { + URL string + Title string + Server string + PoweredBy string +} + +// PhaseProgress is the progress of one scan phase. +type PhaseProgress struct { + Phase string + Status string + Finished int64 + Total int64 + Percent int +} + +// ScanInfoUpdate summarises the scan as the engine resolves it. +type ScanInfoUpdate struct { + TotalTargets int + Targets []string + TotalPocs int + TotalScans int + OOBEnabled bool + OOBStatus string +} + +// ScanStats holds the scan counters. +type ScanStats struct { + StartTime time.Time + EndTime time.Time + TotalTargets int + TotalPocs int + TotalScans int + CompletedScans int32 + FoundVulns int32 +} + +// SDKOptions is the scanner configuration. +// +// The fields up to and including the curated section are the original ones and +// keep their original meaning. The block at the end exposes capabilities added +// after this facade was introduced; leaving them zero reproduces the previous +// behaviour exactly. +type SDKOptions struct { + Targets []string + TargetsFile string + + // PocFile restricts the scan to the given file or directory, hiding the + // built-in PoCs. AppendPoc adds to the built-in set instead. + // + // Specifying both used to silently drop AppendPoc; both are now honoured. + PocFile string + AppendPoc []string + + Search string + Severity string + ExcludePocs []string + ExcludePocsFile string + + RateLimit int + ReqLimitPerTarget int + AutoReqLimit bool + Polite bool + Balanced bool + Aggressive bool + Concurrency int + Retries int + Timeout int + MaxHostError int + Smart bool + + DisableFingerprint bool + EnableWebProbe bool + FingerprintFilterMode string + + MaxRespBodySize int + BruteMaxRequests int + DefaultAccept bool + + // VulnerabilityScannerBreakpoint stops the scan on the first finding. + VulnerabilityScannerBreakpoint bool + + PortScan bool + PSPorts string + PSRateLimit int + PSTimeout int + PSRetries int + PSSkipDiscovery bool + PSS4Chunk int + + Proxy string + Headers []string + + EnableOOB bool + OOB string + OOBKey string + OOBDomain string + OOBApiUrl string + OOBHttpUrl string + OOBRateLimit int + OOBConcurrency int + OOBFinalizeTimeout int + + // EnableStream allocates the SDKScanner channels. Without it the channels + // stay nil and only the callbacks fire. + EnableStream bool + + Dingtalk bool + Wecom bool + + CuratedEnabled string + CuratedEndpoint string + CuratedTimeout int + CuratedForceUpdate bool + + // --- added after the original SDK; zero values keep the old behaviour --- + + // Silent suppresses the scan summary Run prints before starting. The + // original SDK always printed it. + Silent bool + + // PocPaths accepts files, directories and glob patterns such as + // "dir/*.yaml". Entries are appended to the built-in PoCs unless + // PocPathsOnly is set. + PocPaths []string + PocPathsOnly bool + + // IncludeRequestResponse keeps the raw request and response on every + // result. It defaults to true; set DisableRequestResponse to opt out on + // large scans. + DisableRequestResponse bool + // MaxStoredResults caps how many results accumulate in memory. Zero means + // unlimited. + MaxStoredResults int + // RedactedHeaders masks these headers in stored results. Use + // [sdk.DefaultRedactedHeaders] for the usual credential-bearing set. + RedactedHeaders []string + + // OOBPollInterval is the out-of-band poll interval in seconds and + // OOBHitRetention how long a hit is kept, in minutes. Zero uses the same + // defaults as the command line. + OOBPollInterval int + OOBHitRetention int + + // TaskHardTimeoutSec caps a single target+PoC task, in seconds. + // TaskSmartTimeout derives that cap from the PoC's content instead; when + // both are set the larger wins. + TaskHardTimeoutSec int + TaskSmartTimeout bool + + // ResumeFile makes the scan resumable: finished target/PoC pairs are + // recorded there and skipped on a later run. + ResumeFile string + + // Cyberspace sources targets from a search engine instead of, or in + // addition to, Targets. Only "zoomeye" is implemented. + Cyberspace string + Query string + QueryCount int + + // MonitorTargets probes each target's protocol and liveness in parallel + // with the scan. + MonitorTargets bool + + // OnFailure reports PoC executions that failed. These used to be + // discarded, leaving no way to tell a clean scan from a broken one. + OnFailure func(target string, pocID string, err error) +} + +// NewSDKOptions returns the default configuration. +func NewSDKOptions() *SDKOptions { + return &SDKOptions{ + RateLimit: 150, + Concurrency: 25, + Retries: 1, + Timeout: 50, + MaxHostError: 3, + MaxRespBodySize: 2, + BruteMaxRequests: 5000, + DefaultAccept: true, + FingerprintFilterMode: "strict", + PSPorts: "top", + PSS4Chunk: 1000, + OOBRateLimit: 25, + OOBConcurrency: 25, + OOBFinalizeTimeout: -1, + } +} + +// SDKScanner runs a vulnerability scan. +// +// Assign the On* callbacks and read the channels exactly as before. The +// channels are only allocated when SDKOptions.EnableStream is set, and a send +// to a full channel is dropped rather than blocking the scan, which is what +// the original implementation did. +type SDKScanner struct { + // OnResult receives every finding, as the engine's own result type. + OnResult func(*result.Result) + // OnPort receives every open port found by the port pre-scan. + OnPort func(host string, port int) + // OnWebProbe receives every probed web service. + OnWebProbe func(r WebProbeResult) + + ResultChan chan *result.Result + PortChan chan PortScanResult + HostChan chan HostDiscoveryResult + WebProbeChan chan WebProbeResult + PhaseProgressChan chan PhaseProgress + ScanInfoChan chan ScanInfoUpdate + + opts *SDKOptions + + // scanner is rebuilt when one of the Set* methods changes the + // configuration after construction, because the underlying scanner takes + // its options once, at creation. + mu sync.Mutex + scanner *sdk.Scanner + dirty bool + + resultsMu sync.Mutex + results []*result.Result + + stats ScanStats + + ctx context.Context + cancel context.CancelFunc + + closeChansOnce sync.Once + runDoneOnce sync.Once + runDone chan struct{} + started atomic.Bool + closed atomic.Bool +} + +// NewSDKScanner creates a scanner from the given options. A nil opts uses the +// defaults. +func NewSDKScanner(opts *SDKOptions) (*SDKScanner, error) { + if opts == nil { + opts = NewSDKOptions() + } + + ctx, cancel := context.WithCancel(context.Background()) + s := &SDKScanner{ + opts: opts, + ctx: ctx, + cancel: cancel, + runDone: make(chan struct{}), + } + + if opts.EnableStream { + s.ResultChan = make(chan *result.Result, 100) + s.PortChan = make(chan PortScanResult, 100) + s.HostChan = make(chan HostDiscoveryResult, 256) + s.WebProbeChan = make(chan WebProbeResult, 100) + s.PhaseProgressChan = make(chan PhaseProgress, 64) + s.ScanInfoChan = make(chan ScanInfoUpdate, 16) + } + + // Build eagerly so configuration errors surface here rather than at Run, + // which is where the original SDK reported them. + scanner, err := s.build() + if err != nil { + cancel() + return nil, err + } + s.scanner = scanner + s.refreshStats() + + return s, nil +} + +// build translates the options into the current SDK and creates a scanner. +func (s *SDKScanner) build() (*sdk.Scanner, error) { + o := s.opts + + options := []sdk.Option{ + sdk.WithTargets(o.Targets...), + sdk.WithConcurrency(orDefault(o.Concurrency, 25)), + sdk.WithRateLimit(orDefault(o.RateLimit, 150)), + sdk.WithTimeout(orDefault(o.Timeout, 50)), + sdk.WithMaxRespBodySize(orDefault(o.MaxRespBodySize, 2)), + + sdk.WithResultHandler(func(sdk.Result) {}), // keeps the result path warm + sdk.WithRawResultHandler(s.handleRawResult), + sdk.WithPortHandler(s.handlePort), + sdk.WithHostHandler(s.handleHost), + sdk.WithWebProbeHandler(s.handleWebProbe), + sdk.WithProgressHandler(s.handleProgress), + sdk.WithScanInfoHandler(s.handleScanInfo), + sdk.WithFailureHandler(s.handleFailure), + } + + if o.Retries >= 0 { + options = append(options, sdk.WithRetries(o.Retries)) + } + if o.MaxHostError >= 0 { + options = append(options, sdk.WithMaxHostError(o.MaxHostError)) + } + if strings.TrimSpace(o.TargetsFile) != "" { + options = append(options, sdk.WithTargetsFile(o.TargetsFile)) + } + + // PocFile keeps its exclusive meaning; PocPaths and AppendPoc append. + // Combining them used to drop AppendPoc silently. + var pocPaths []string + if v := strings.TrimSpace(o.PocFile); v != "" { + pocPaths = append(pocPaths, v) + } + pocPaths = append(pocPaths, o.PocPaths...) + pocPaths = append(pocPaths, o.AppendPoc...) + if len(pocPaths) > 0 { + options = append(options, sdk.WithPocPaths(pocPaths...)) + } + if o.PocPathsOnly || strings.TrimSpace(o.PocFile) != "" { + options = append(options, sdk.WithPocPathsOnly()) + } + + if v := strings.TrimSpace(o.Search); v != "" { + options = append(options, sdk.WithSearch(v)) + } + if v := strings.TrimSpace(o.Severity); v != "" { + options = append(options, sdk.WithSeverity(v)) + } + if len(o.ExcludePocs) > 0 { + options = append(options, sdk.WithExcludePocs(o.ExcludePocs...)) + } + if v := strings.TrimSpace(o.ExcludePocsFile); v != "" { + options = append(options, sdk.WithExcludePocsFile(v)) + } + + if o.ReqLimitPerTarget > 0 { + options = append(options, sdk.WithRequestLimitPerTarget(o.ReqLimitPerTarget)) + } + if o.AutoReqLimit { + options = append(options, sdk.WithAutoRequestLimit()) + } + if o.Polite { + options = append(options, sdk.WithPolite()) + } + if o.Balanced { + options = append(options, sdk.WithBalanced()) + } + if o.Aggressive { + options = append(options, sdk.WithAggressive()) + } + if o.Smart { + options = append(options, sdk.WithSmartConcurrency()) + } + if o.VulnerabilityScannerBreakpoint { + options = append(options, sdk.WithStopOnFirstMatch()) + } + + if o.DisableFingerprint { + options = append(options, sdk.WithFingerprintDisabled()) + } + if v := strings.TrimSpace(o.FingerprintFilterMode); v != "" { + options = append(options, sdk.WithFingerprintFilterMode(v)) + } + if o.EnableWebProbe { + options = append(options, sdk.WithWebProbe()) + } + + if v := strings.TrimSpace(o.Proxy); v != "" { + options = append(options, sdk.WithProxy(v)) + } + if len(o.Headers) > 0 { + options = append(options, sdk.WithHeaders(o.Headers...)) + } + + if o.PortScan { + options = append(options, sdk.WithPortScan(sdk.PortScanOptions{ + Ports: o.PSPorts, + RateLimit: o.PSRateLimit, + TimeoutMs: o.PSTimeout, + Retries: o.PSRetries, + SkipDiscovery: o.PSSkipDiscovery, + ChunkSize: o.PSS4Chunk, + })) + } + + if o.EnableOOB && strings.TrimSpace(o.OOB) != "" { + options = append(options, sdk.WithOOB(sdk.OOBOptions{ + Adapter: o.OOB, + Key: o.OOBKey, + Domain: o.OOBDomain, + ApiURL: o.OOBApiUrl, + HttpURL: o.OOBHttpUrl, + RateLimit: o.OOBRateLimit, + Concurrency: o.OOBConcurrency, + FinalizeTimeout: o.OOBFinalizeTimeout, + PollInterval: o.OOBPollInterval, + HitRetention: o.OOBHitRetention, + })) + } + + if o.Dingtalk { + options = append(options, sdk.WithDingtalk()) + } + if o.Wecom { + options = append(options, sdk.WithWecom()) + } + + if strings.TrimSpace(o.CuratedEnabled) != "" || strings.TrimSpace(o.CuratedEndpoint) != "" || + o.CuratedTimeout > 0 || o.CuratedForceUpdate { + options = append(options, sdk.WithCurated(sdk.CuratedOptions{ + Enabled: o.CuratedEnabled, + Endpoint: o.CuratedEndpoint, + TimeoutSec: o.CuratedTimeout, + ForceUpdate: o.CuratedForceUpdate, + })) + } + + if o.DisableRequestResponse { + options = append(options, sdk.WithRequestResponse(false)) + } + if o.MaxStoredResults > 0 { + options = append(options, sdk.WithMaxStoredResults(o.MaxStoredResults)) + } + if len(o.RedactedHeaders) > 0 { + options = append(options, sdk.WithRedactedHeaders(o.RedactedHeaders...)) + } + + if o.TaskHardTimeoutSec > 0 || o.TaskSmartTimeout { + options = append(options, sdk.WithTaskTimeout(sdk.TaskTimeoutOptions{ + HardSec: o.TaskHardTimeoutSec, + Smart: o.TaskSmartTimeout, + })) + } + if v := strings.TrimSpace(o.ResumeFile); v != "" { + options = append(options, sdk.WithCheckpoint(sdk.CheckpointOptions{Path: v})) + } + if strings.TrimSpace(o.Cyberspace) != "" { + options = append(options, sdk.WithCyberspace(sdk.CyberspaceOptions{ + Engine: o.Cyberspace, + Query: o.Query, + Count: o.QueryCount, + })) + } + if o.MonitorTargets { + options = append(options, sdk.WithTargetPreProbe()) + } + + return sdk.New(s.ctx, options...) +} + +func orDefault(v, fallback int) int { + if v <= 0 { + return fallback + } + return v +} + +// --- event plumbing --------------------------------------------------------- +// +// The callbacks and channels are public fields that callers assign after +// construction, so every handler reads them at call time rather than +// capturing them when the scanner is built. + +func (s *SDKScanner) handleRawResult(r *result.Result) { + if r == nil { + return + } + s.resultsMu.Lock() + s.results = append(s.results, r) + s.resultsMu.Unlock() + atomic.AddInt32(&s.stats.FoundVulns, 1) + + if fn := s.OnResult; fn != nil { + fn(r) + } + sendOrDrop(s.ctx, s.ResultChan, r) +} + +func (s *SDKScanner) handlePort(e sdk.PortEvent) { + if fn := s.OnPort; fn != nil { + fn(e.Host, e.Port) + } + sendOrDrop(s.ctx, s.PortChan, PortScanResult{Host: e.Host, Port: e.Port}) +} + +func (s *SDKScanner) handleHost(e sdk.HostEvent) { + sendOrDrop(s.ctx, s.HostChan, HostDiscoveryResult{Host: e.Host}) +} + +func (s *SDKScanner) handleWebProbe(e sdk.WebProbeEvent) { + r := WebProbeResult{URL: e.URL, Title: e.Title, Server: e.Server, PoweredBy: e.PoweredBy} + if fn := s.OnWebProbe; fn != nil { + fn(r) + } + sendOrDrop(s.ctx, s.WebProbeChan, r) +} + +func (s *SDKScanner) handleProgress(p sdk.PhaseProgress) { + if p.Phase == sdk.PhaseVuln { + atomic.StoreInt32(&s.stats.CompletedScans, int32(p.Finished)) + } + sendOrDrop(s.ctx, s.PhaseProgressChan, PhaseProgress{ + Phase: p.Phase, + Status: p.Status, + Finished: p.Finished, + Total: p.Total, + Percent: p.Percent, + }) +} + +func (s *SDKScanner) handleScanInfo(i sdk.ScanInfo) { + sendOrDrop(s.ctx, s.ScanInfoChan, ScanInfoUpdate{ + TotalTargets: i.TotalTargets, + Targets: i.Targets, + TotalPocs: i.TotalPocs, + TotalScans: i.TotalScans, + OOBEnabled: i.OOBEnabled, + OOBStatus: i.OOBStatus, + }) +} + +func (s *SDKScanner) handleFailure(f sdk.Failure) { + if fn := s.opts.OnFailure; fn != nil { + fn(f.Target, f.PocID, f.Err) + } +} + +// sendOrDrop mirrors the original streaming behaviour: a full channel drops +// the value instead of stalling the scan. Callers that never drain a channel +// would otherwise deadlock the whole scan, which is why this cannot become a +// blocking send. +func sendOrDrop[T any](ctx context.Context, ch chan T, v T) { + if ch == nil { + return + } + // A send racing closeChans would panic; the scan must not die for that. + defer func() { _ = recover() }() + select { + case ch <- v: + case <-ctx.Done(): + default: + } +} + +func (s *SDKScanner) closeChans() { + s.closeChansOnce.Do(func() { + for _, c := range []func(){ + func() { closeChan(s.ResultChan) }, + func() { closeChan(s.PortChan) }, + func() { closeChan(s.HostChan) }, + func() { closeChan(s.WebProbeChan) }, + func() { closeChan(s.PhaseProgressChan) }, + func() { closeChan(s.ScanInfoChan) }, + } { + c() + } + }) +} + +func closeChan[T any](ch chan T) { + if ch != nil { + close(ch) + } +} + +// --- lifecycle -------------------------------------------------------------- + +// Run executes the scan synchronously and returns when it finishes. +func (s *SDKScanner) Run() error { + sc, err := s.current() + if err != nil { + return err + } + s.started.Store(true) + s.printScanInfo() + + runErr := sc.Execute(s.ctx) + + s.refreshStats() + s.stats.EndTime = time.Now() + s.closeChans() + s.runDoneOnce.Do(func() { close(s.runDone) }) + return runErr +} + +// RunAsync starts the scan in the background and returns immediately. +func (s *SDKScanner) RunAsync() error { + sc, err := s.current() + if err != nil { + return err + } + s.started.Store(true) + go func() { + s.printScanInfo() + _ = sc.Execute(s.ctx) + s.refreshStats() + s.stats.EndTime = time.Now() + s.closeChans() + s.runDoneOnce.Do(func() { close(s.runDone) }) + }() + return nil +} + +// current returns the scanner to run, rebuilding it when a Set* call changed +// the configuration after construction. +func (s *SDKScanner) current() (*sdk.Scanner, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if !s.dirty { + return s.scanner, nil + } + rebuilt, err := s.build() + if err != nil { + return nil, err + } + if s.scanner != nil { + _ = s.scanner.Close() + } + s.scanner = rebuilt + s.dirty = false + s.refreshStatsLocked() + return s.scanner, nil +} + +// Stop asks the scan to stop and returns immediately. +func (s *SDKScanner) Stop() { + s.mu.Lock() + sc := s.scanner + s.mu.Unlock() + if sc != nil { + sc.Stop() + } + s.cancel() +} + +// Close stops the scan, waits for it to finish and releases every resource. +// It is safe to call more than once. +func (s *SDKScanner) Close() { + if s.closed.Swap(true) { + return + } + s.Stop() + if s.started.Load() { + <-s.runDone + } + s.closeChans() + + s.mu.Lock() + sc := s.scanner + s.mu.Unlock() + if sc != nil { + _ = sc.Close() + } + + s.resultsMu.Lock() + s.results = nil + s.resultsMu.Unlock() +} + +// Pause suspends task scheduling. +func (s *SDKScanner) Pause() { + if sc := s.peek(); sc != nil { + sc.Pause() + } +} + +// Resume resumes a paused scan. +func (s *SDKScanner) Resume() { + if sc := s.peek(); sc != nil { + sc.Resume() + } +} + +// IsPaused reports whether the scan is paused. +func (s *SDKScanner) IsPaused() bool { + sc := s.peek() + return sc != nil && sc.IsPaused() +} + +// IsStopping reports whether Stop has been called. +func (s *SDKScanner) IsStopping() bool { + sc := s.peek() + return sc != nil && sc.IsStopping() +} + +func (s *SDKScanner) peek() *sdk.Scanner { + s.mu.Lock() + defer s.mu.Unlock() + return s.scanner +} + +// --- results ---------------------------------------------------------------- + +// GetResults returns the findings collected so far. +func (s *SDKScanner) GetResults() []*result.Result { + s.resultsMu.Lock() + defer s.resultsMu.Unlock() + out := make([]*result.Result, len(s.results)) + copy(out, s.results) + return out +} + +// GetOpenPorts returns the open ports discovered by the port pre-scan. +func (s *SDKScanner) GetOpenPorts() map[string][]int { + if sc := s.peek(); sc != nil { + return sc.OpenPorts() + } + return map[string][]int{} +} + +// HasVulnerabilities reports whether any vulnerability was found. +func (s *SDKScanner) HasVulnerabilities() bool { + return s.GetVulnerabilityCount() > 0 +} + +// GetVulnerabilityCount returns the number of findings. +func (s *SDKScanner) GetVulnerabilityCount() int { + if sc := s.peek(); sc != nil { + return sc.ResultCount() + } + return 0 +} + +// GetStats returns a snapshot of the scan counters. +func (s *SDKScanner) GetStats() ScanStats { + s.refreshStats() + out := s.stats + out.CompletedScans = atomic.LoadInt32(&s.stats.CompletedScans) + out.FoundVulns = atomic.LoadInt32(&s.stats.FoundVulns) + return out +} + +func (s *SDKScanner) refreshStats() { + s.mu.Lock() + defer s.mu.Unlock() + s.refreshStatsLocked() +} + +func (s *SDKScanner) refreshStatsLocked() { + if s.scanner == nil { + return + } + st := s.scanner.Stats() + s.stats.StartTime = st.StartTime + if !st.EndTime.IsZero() { + s.stats.EndTime = st.EndTime + } + s.stats.TotalTargets = st.TotalTargets + s.stats.TotalPocs = st.TotalPocs + s.stats.TotalScans = st.TotalScans + atomic.StoreInt32(&s.stats.CompletedScans, int32(st.CompletedScans)) + atomic.StoreInt32(&s.stats.FoundVulns, int32(st.FoundVulns)) +} + +// GetProgress returns overall scan progress in the range [0, 100]. +func (s *SDKScanner) GetProgress() float64 { + if sc := s.peek(); sc != nil { + return sc.Progress() + } + return 0 +} + +// --- runtime adjustment ----------------------------------------------------- + +// SetProxy sets the HTTP or SOCKS5 proxy. Call it before Run. +func (s *SDKScanner) SetProxy(proxy string) { + s.mu.Lock() + s.opts.Proxy = proxy + s.dirty = true + timeout, retries := s.opts.Timeout, s.opts.Retries + maxBody, reqLimit := s.opts.MaxRespBodySize, s.opts.ReqLimitPerTarget + accept := s.opts.DefaultAccept + s.mu.Unlock() + + // The HTTP client is process-wide, so the proxy takes effect immediately + // as it did before. + _ = retryhttpclient.Init(&retryhttpclient.Options{ + Proxy: proxy, + Timeout: timeout, + Retries: retries, + MaxRespBodySize: maxBody, + ReqLimitPerTarget: reqLimit, + DefaultAccept: accept, + }) +} + +// SetRateLimit sets the requests-per-second cap. Call it before Run. +func (s *SDKScanner) SetRateLimit(rateLimit int) { + s.mu.Lock() + defer s.mu.Unlock() + s.opts.RateLimit = rateLimit + s.dirty = true +} + +// SetConcurrency sets the number of concurrent workers. Call it before Run. +func (s *SDKScanner) SetConcurrency(concurrency int) { + s.mu.Lock() + defer s.mu.Unlock() + s.opts.Concurrency = concurrency + s.dirty = true +} + +// --- out-of-band ------------------------------------------------------------ + +// IsOOBEnabled reports whether out-of-band detection is configured. +func (s *SDKScanner) IsOOBEnabled() bool { + if s.opts == nil { + return false + } + if s.opts.EnableOOB { + return true + } + return s.opts.OOB != "" && (s.opts.OOBKey != "" || s.opts.OOBDomain != "") +} + +// GetOOBStatus reports whether out-of-band detection is usable, together with +// a human-readable description. It probes the configured service. +func (s *SDKScanner) GetOOBStatus() (bool, string) { + if s.opts == nil || !s.opts.EnableOOB { + return false, "OOB未配置或未启用" + } + sc := s.peek() + if sc == nil { + return false, "扫描器未初始化" + } + return sc.OOBStatus() +} + +// Scanner exposes the underlying [sdk.Scanner] for callers that want the +// current API's streams, diagnostics and typed results without giving up this +// facade. It is nil only before construction succeeds. +func (s *SDKScanner) Scanner() *sdk.Scanner { return s.peek() } + +// printScanInfo writes the pre-scan summary. The original SDK always printed +// it, so it stays on unless SDKOptions.Silent is set. +func (s *SDKScanner) printScanInfo() { + if s.opts != nil && s.opts.Silent { + return + } + sc := s.peek() + if sc == nil { + return + } + info := sc.Info() + + fmt.Printf("\n========== 扫描信息 ==========\n") + fmt.Printf("目标数量: %d\n", info.TotalTargets) + fmt.Printf("POC数量: %d\n", info.TotalPocs) + fmt.Printf("总扫描任务: %d\n", info.TotalScans) + + if len(info.Targets) <= 5 { + fmt.Printf("扫描目标: %s\n", strings.Join(info.Targets, ", ")) + } else { + fmt.Printf("目标过多，仅显示前3个: %s...\n", strings.Join(info.Targets[:3], ", ")) + } + + if s.opts.EnableOOB { + if ok, status := info.OOBEnabled, info.OOBStatus; ok { + fmt.Printf("OOB状态: ✓ %s\n", status) + } else { + fmt.Printf("OOB状态: ✗ %s\n", status) + } + } else { + fmt.Printf("OOB状态: ✗ OOB未配置或未启用\n") + } + + fmt.Printf("=============================\n") +} diff --git a/afrog_test.go b/afrog_test.go new file mode 100644 index 000000000..ace6e5978 --- /dev/null +++ b/afrog_test.go @@ -0,0 +1,448 @@ +package afrog + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/zan8in/afrog/v3/pkg/result" +) + +const compatToken = "AFROG_COMPAT_TOKEN" + +func newCompatServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Server", "afrog-compat/1.0") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("compat" + compatToken + "")) + })) + t.Cleanup(srv.Close) + return srv +} + +func writeCompatPoc(t *testing.T, dir, id string) { + t.Helper() + body := "id: " + id + ` +info: + name: compat test poc + author: compat + severity: info +rules: + r0: + request: + method: GET + path: /probe + expression: response.status == 200 && response.body.bcontains(b"` + compatToken + `") +expression: r0() +` + if err := os.WriteFile(filepath.Join(dir, id+".yaml"), []byte(body), 0o644); err != nil { + t.Fatalf("write poc: %v", err) + } +} + +// newCompatOptions builds options the way pre-existing integrations do: take +// the defaults, then assign struct fields directly. +func newCompatOptions(t *testing.T, target string) *SDKOptions { + t.Helper() + dir := t.TempDir() + writeCompatPoc(t, dir, "compat-match") + + opts := NewSDKOptions() + opts.Targets = []string{target} + opts.PocFile = dir + opts.DisableFingerprint = true + opts.Timeout = 10 + opts.Silent = true // keeps the test output clean; not part of the old API + return opts +} + +// The original usage pattern must keep working unchanged. +func TestCompat_BasicScan(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + + results := scanner.GetResults() + if len(results) != 1 { + t.Fatalf("GetResults returned %d results, want 1", len(results)) + } + if !scanner.HasVulnerabilities() { + t.Error("HasVulnerabilities = false after a finding") + } + if got := scanner.GetVulnerabilityCount(); got != 1 { + t.Errorf("GetVulnerabilityCount = %d, want 1", got) + } + + // GetResults still hands back the engine's own type, with the request and + // response attached. + r := results[0] + if r.PocInfo == nil || r.PocInfo.Id != "compat-match" { + t.Errorf("result carries the wrong PoC: %+v", r.PocInfo) + } + if len(r.AllPocResult) == 0 { + t.Error("result carries no request/response") + } +} + +// Defaults must match the original ones exactly. +func TestCompat_NewSDKOptionsDefaults(t *testing.T) { + o := NewSDKOptions() + tests := []struct { + name string + got any + want any + }{ + {"RateLimit", o.RateLimit, 150}, + {"Concurrency", o.Concurrency, 25}, + {"Retries", o.Retries, 1}, + {"Timeout", o.Timeout, 50}, + {"MaxHostError", o.MaxHostError, 3}, + {"MaxRespBodySize", o.MaxRespBodySize, 2}, + {"BruteMaxRequests", o.BruteMaxRequests, 5000}, + {"DefaultAccept", o.DefaultAccept, true}, + {"FingerprintFilterMode", o.FingerprintFilterMode, "strict"}, + {"PSPorts", o.PSPorts, "top"}, + {"PSS4Chunk", o.PSS4Chunk, 1000}, + {"OOBRateLimit", o.OOBRateLimit, 25}, + {"OOBConcurrency", o.OOBConcurrency, 25}, + {"OOBFinalizeTimeout", o.OOBFinalizeTimeout, -1}, + } + for _, tt := range tests { + if tt.got != tt.want { + t.Errorf("%s = %v, want %v", tt.name, tt.got, tt.want) + } + } +} + +// Callbacks are public fields assigned after construction, so they must be +// read at call time rather than captured when the scanner is built. +func TestCompat_CallbacksAssignedAfterConstruction(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + var mu sync.Mutex + var got []*result.Result + scanner.OnResult = func(r *result.Result) { + mu.Lock() + got = append(got, r) + mu.Unlock() + } + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + + mu.Lock() + defer mu.Unlock() + if len(got) != 1 { + t.Fatalf("OnResult fired %d times, want 1", len(got)) + } +} + +// EnableStream allocates the channels; they must deliver and then close so a +// range loop terminates. +func TestCompat_StreamingChannels(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + opts.EnableStream = true + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if scanner.ResultChan == nil || scanner.PhaseProgressChan == nil { + t.Fatal("EnableStream did not allocate the channels") + } + + var count int + drained := make(chan struct{}) + go func() { + defer close(drained) + for range scanner.ResultChan { + count++ + } + }() + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + + select { + case <-drained: + case <-time.After(30 * time.Second): + t.Fatal("ResultChan was not closed when the scan finished") + } + if count != 1 { + t.Errorf("received %d results from ResultChan, want 1", count) + } +} + +// Without EnableStream the channels stay nil, exactly as before. +func TestCompat_ChannelsNilWithoutEnableStream(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if scanner.ResultChan != nil || scanner.PortChan != nil || scanner.HostChan != nil || + scanner.WebProbeChan != nil || scanner.PhaseProgressChan != nil || scanner.ScanInfoChan != nil { + t.Error("channels were allocated without EnableStream") + } +} + +// A caller that enables streaming but never drains must not stall the scan. +// The original implementation dropped on a full channel, and code in the wild +// relies on that. +func TestCompat_UndrainedChannelDoesNotBlockScan(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + opts.EnableStream = true + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + done := make(chan error, 1) + go func() { done <- scanner.Run() }() + + select { + case err := <-done: + if err != nil { + t.Fatalf("Run: %v", err) + } + case <-time.After(60 * time.Second): + t.Fatal("scan stalled on an undrained channel") + } +} + +func TestCompat_RunAsync(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if err := scanner.RunAsync(); err != nil { + t.Fatalf("RunAsync: %v", err) + } + + deadline := time.Now().Add(60 * time.Second) + for time.Now().Before(deadline) { + if scanner.GetProgress() >= 100 && scanner.GetVulnerabilityCount() > 0 { + break + } + time.Sleep(50 * time.Millisecond) + } + if got := scanner.GetVulnerabilityCount(); got != 1 { + t.Fatalf("found %d results after RunAsync, want 1", got) + } +} + +// The setters are called between construction and Run, so their values have to +// reach the scan. +func TestCompat_SettersApplyBeforeRun(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + scanner.SetRateLimit(42) + scanner.SetConcurrency(3) + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + if got := scanner.GetVulnerabilityCount(); got != 1 { + t.Fatalf("found %d results, want 1", got) + } + if scanner.opts.RateLimit != 42 || scanner.opts.Concurrency != 3 { + t.Errorf("setters did not update the options: rate=%d concurrency=%d", + scanner.opts.RateLimit, scanner.opts.Concurrency) + } +} + +func TestCompat_StatsAndProgress(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + before := scanner.GetStats() + if before.TotalPocs != 1 { + t.Errorf("TotalPocs = %d before the scan, want 1", before.TotalPocs) + } + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + + after := scanner.GetStats() + if after.FoundVulns != 1 { + t.Errorf("FoundVulns = %d, want 1", after.FoundVulns) + } + if after.CompletedScans < 1 { + t.Errorf("CompletedScans = %d, want at least 1", after.CompletedScans) + } + if after.StartTime.IsZero() || after.EndTime.IsZero() { + t.Error("StartTime or EndTime not set") + } + if got := scanner.GetProgress(); got != 100 { + t.Errorf("GetProgress = %.2f after a clean finish, want 100", got) + } +} + +// Close must be safe before a scan, after one, and when called twice. +func TestCompat_CloseAtEveryPoint(t *testing.T) { + srv := newCompatServer(t) + + tests := []struct { + name string + before func(s *SDKScanner) + }{ + {"never started", func(*SDKScanner) {}}, + {"after run", func(s *SDKScanner) { _ = s.Run() }}, + {"after stop", func(s *SDKScanner) { s.Stop() }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + tt.before(scanner) + + done := make(chan struct{}) + go func() { defer close(done); scanner.Close(); scanner.Close() }() + + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("Close blocked") + } + }) + } +} + +// A missing target used to be a configuration error at construction time. +func TestCompat_ConstructionErrorsSurfaceEarly(t *testing.T) { + opts := NewSDKOptions() + opts.Silent = true + if _, err := NewSDKScanner(opts); err == nil { + t.Fatal("NewSDKScanner with no targets should have failed") + } +} + +// A nil options value falls back to the defaults rather than panicking. +func TestCompat_NilOptions(t *testing.T) { + if _, err := NewSDKScanner(nil); err == nil { + t.Fatal("NewSDKScanner(nil) should fail on the missing target, not panic") + } +} + +// Specifying PocFile and AppendPoc together used to drop AppendPoc silently. +func TestCompat_PocFileAndAppendPocAreBothHonoured(t *testing.T) { + srv := newCompatServer(t) + + dirA := t.TempDir() + writeCompatPoc(t, dirA, "compat-a") + dirB := t.TempDir() + writeCompatPoc(t, dirB, "compat-b") + + opts := NewSDKOptions() + opts.Targets = []string{srv.URL} + opts.PocFile = dirA + opts.AppendPoc = []string{dirB} + opts.DisableFingerprint = true + opts.Timeout = 10 + opts.Silent = true + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if got := scanner.GetStats().TotalPocs; got != 2 { + t.Fatalf("loaded %d pocs, want 2 (PocFile + AppendPoc)", got) + } +} + +// The new options must be reachable from the old struct without changing how +// the rest of it is used. +func TestCompat_NewCapabilitiesAreReachable(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + opts.TaskHardTimeoutSec = 30 + opts.MaxStoredResults = 10 + opts.ResumeFile = filepath.Join(t.TempDir(), "compat.afg") + + var failures int + opts.OnFailure = func(string, string, error) { failures++ } + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner with the new options: %v", err) + } + defer scanner.Close() + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + if got := scanner.GetVulnerabilityCount(); got != 1 { + t.Fatalf("found %d results, want 1", got) + } + if _, err := os.Stat(opts.ResumeFile); err != nil { + t.Errorf("resume file was not written: %v", err) + } +} + +// The facade must expose the current SDK for callers that want to migrate +// gradually. +func TestCompat_ScannerAccessorExposesTheModernAPI(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + modern := scanner.Scanner() + if modern == nil { + t.Fatal("Scanner() returned nil") + } + if got := modern.PocCount(); got != 1 { + t.Errorf("PocCount = %d, want 1", got) + } +} diff --git a/docs/SDK_Usage_Guide_English.md b/docs/SDK_Usage_Guide_English.md index 5eeda26fa..2281637b6 100644 --- a/docs/SDK_Usage_Guide_English.md +++ b/docs/SDK_Usage_Guide_English.md @@ -15,6 +15,42 @@ The Afrog SDK is the Go API for embedding vulnerability scanning into your own p - **Typed errors** — failures are matched with `errors.Is` - **Deterministic cleanup** — `Close` releases every background goroutine +### Two APIs, one implementation + +| Package | Entry point | Use for | +|---|---|---| +| `github.com/zan8in/afrog/v3/pkg/sdk` | `sdk.New(ctx, opts...)` | New code. This guide covers it. | +| `github.com/zan8in/afrog/v3` | `afrog.NewSDKScanner(opts)` | Existing integrations, unchanged | + +The root package is a compatibility facade that delegates to `pkg/sdk`, so **the fixes and features described here reach both APIs**. Existing code needs no changes: + +```go +options := afrog.NewSDKOptions() +options.Targets = []string{"https://example.com"} +options.PocFile = pocPath +options.Concurrency = 10 + +scanner, err := afrog.NewSDKScanner(options) +if err != nil { + log.Fatal(err) +} +defer scanner.Close() + +scanner.OnResult = func(r *result.Result) { + log.Printf("found: %s", r.PocInfo.Id) +} +if err := scanner.Run(); err != nil { + log.Fatal(err) +} +results := scanner.GetResults() +``` + +`SDKOptions` gained optional fields that expose the newer capabilities to the old style: `PocPaths`/`PocPathsOnly` (globs and append semantics), `ResumeFile`, `TaskHardTimeoutSec`/`TaskSmartTimeout`, `Cyberspace`/`Query`/`QueryCount`, `MonitorTargets`, `OOBPollInterval`/`OOBHitRetention`, `MaxStoredResults`, `RedactedHeaders`, `OnFailure` and `Silent`. Leaving them zero reproduces the previous behaviour exactly. + +To migrate gradually, `scanner.Scanner()` returns the underlying `*sdk.Scanner`, giving access to the streams and diagnostics of the current API. + +One old behaviour was corrected: setting `PocFile` and `AppendPoc` together used to drop `AppendPoc` silently; both are now loaded. + ## Installation ```bash diff --git "a/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" "b/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" index 20375b86d..b933bd50b 100644 --- "a/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" +++ "b/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" @@ -15,6 +15,42 @@ Afrog SDK 是把漏洞扫描能力嵌入自己程序的 Go 接口，包路径为 - **类型化错误** —— 使用 `errors.Is` 判断失败原因 - **资源可控** —— `Close` 释放所有后台协程，无泄漏 +### 两套 API 共存 + +| 包 | 入口 | 适用 | +|---|---|---| +| `github.com/zan8in/afrog/v3/pkg/sdk` | `sdk.New(ctx, opts...)` | 新代码，本文档主要介绍这套 | +| `github.com/zan8in/afrog/v3` | `afrog.NewSDKScanner(opts)` | 已有集成，保持原样即可 | + +根包是旧 API 的兼容门面，内部委托给 `pkg/sdk`，因此**这次的缺陷修复和新能力对两套 API 同时生效**。旧写法无需改动： + +```go +options := afrog.NewSDKOptions() +options.Targets = []string{"https://example.com"} +options.PocFile = pocPath +options.Concurrency = 10 + +scanner, err := afrog.NewSDKScanner(options) +if err != nil { + log.Fatal(err) +} +defer scanner.Close() + +scanner.OnResult = func(r *result.Result) { + log.Printf("发现: %s", r.PocInfo.Id) +} +if err := scanner.Run(); err != nil { + log.Fatal(err) +} +results := scanner.GetResults() +``` + +`SDKOptions` 上新增了若干可选字段，把这次的新能力开放给旧写法：`PocPaths`/`PocPathsOnly`（glob 与追加语义）、`ResumeFile`（断点续扫）、`TaskHardTimeoutSec`/`TaskSmartTimeout`、`Cyberspace`/`Query`/`QueryCount`、`MonitorTargets`、`OOBPollInterval`/`OOBHitRetention`、`MaxStoredResults`、`RedactedHeaders`、`OnFailure`、`Silent`。留空则行为与以前完全一致。 + +想逐步迁移的话，`scanner.Scanner()` 会返回底层的 `*sdk.Scanner`，可以直接用新 API 的流式订阅与诊断能力。 + +有一处旧行为被修正：同时指定 `PocFile` 与 `AppendPoc` 时，旧版会静默丢弃 `AppendPoc`，现在两者都会加载。 + ## 安装 ```bash From 84eb9d3b897ce94080503a6e2bd6c76b4749e5e8 Mon Sep 17 00:00:00 2001 From: zhizhuo Date: Tue, 11 Aug 2026 22:47:40 +0800 Subject: [PATCH 4/5] =?UTF-8?q?=E6=96=87=E6=A1=A3=E8=A1=A5=E5=85=85?= =?UTF-8?q?=E5=8D=87=E7=BA=A7=E6=8F=90=E7=A4=BA=EF=BC=8C=E8=A1=A5=E9=BD=90?= =?UTF-8?q?=E4=B8=A4=E5=A5=97=20SDK=20=E7=9A=84=E6=B5=8B=E8=AF=95=E8=A6=86?= =?UTF-8?q?=E7=9B=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 文档 - 中英文指南开头加入醒目的升级提示：说明本次改动较大、建议迁移到 pkg/sdk，同时明确旧接口入口与调用方式保持不变，优化与缺陷修复 都在内部完成 兼容性验证 - 用 go doc 对改动前后的根包 API 面做逐项机器比对：旧 API 零删除、 零签名变更，只新增了 16 个可选字段与 1 个 Scanner() 方法 测试 - 旧 API 补齐 Pause/Resume/IsPaused/IsStopping/SetProxy/GetOpenPorts/ IsOOBEnabled/GetOOBStatus 的覆盖，20 个方法现已全部有测试； 另加打印行为、二次 Run 被拒绝的用例 - 新 API 补一组表驱动测试，逐个断言 51 个选项确实生效，避免某个 单行选项函数写错却无人察觉；并断言关键选项透传到引擎、SDK 模式 下不会配置任何报告文件输出 - 补齐访问器、剩余四个事件流、错误类型的 Unwrap、端口记录与 进度加权路径 覆盖率：根包 81.4%，pkg/sdk 82.7%（此前 68.8%） 全量 15 个包在 -race 下通过；改动前的 7 个真实旧示例未经修改编译通过 --- afrog_test.go | 255 +++++++++++++ docs/SDK_Usage_Guide_English.md | 10 + ...7\345\215\227_\344\270\255\346\226\207.md" | 10 + pkg/sdk/coverage_test.go | 345 ++++++++++++++++++ 4 files changed, 620 insertions(+) create mode 100644 pkg/sdk/coverage_test.go diff --git a/afrog_test.go b/afrog_test.go index ace6e5978..a673dd959 100644 --- a/afrog_test.go +++ b/afrog_test.go @@ -5,6 +5,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strings" "sync" "testing" "time" @@ -12,6 +13,42 @@ import ( "github.com/zan8in/afrog/v3/pkg/result" ) +func captureStdout(t *testing.T) (*os.File, *os.File, *os.File) { + t.Helper() + stdout := os.Stdout + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe: %v", err) + } + os.Stdout = w + return stdout, r, w +} + +func restoreStdout(t *testing.T, stdout, r, w *os.File) string { + t.Helper() + os.Stdout = stdout + _ = w.Close() + + var sb strings.Builder + done := make(chan struct{}) + go func() { + defer close(done) + buf := make([]byte, 4096) + for { + n, err := r.Read(buf) + if n > 0 { + sb.Write(buf[:n]) + } + if err != nil { + return + } + } + }() + <-done + _ = r.Close() + return sb.String() +} + const compatToken = "AFROG_COMPAT_TOKEN" func newCompatServer(t *testing.T) *httptest.Server { @@ -428,6 +465,224 @@ func TestCompat_NewCapabilitiesAreReachable(t *testing.T) { } } +// Pause, Resume and IsPaused must still gate the scan. +func TestCompat_PauseResume(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if scanner.IsPaused() { + t.Error("IsPaused = true before anything happened") + } + + scanner.Pause() + if !scanner.IsPaused() { + t.Error("IsPaused = false after Pause") + } + + scanner.Resume() + if scanner.IsPaused() { + t.Error("IsPaused = true after Resume") + } + + // The scan must still complete normally after a pause/resume cycle. + if err := scanner.Run(); err != nil { + t.Fatalf("Run after pause/resume: %v", err) + } + if got := scanner.GetVulnerabilityCount(); got != 1 { + t.Errorf("found %d results, want 1", got) + } +} + +// IsStopping reflects Stop, and a stopped scan must not hang. +func TestCompat_StopAndIsStopping(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if scanner.IsStopping() { + t.Error("IsStopping = true before Stop") + } + + if err := scanner.RunAsync(); err != nil { + t.Fatalf("RunAsync: %v", err) + } + scanner.Stop() + + if !scanner.IsStopping() { + t.Error("IsStopping = false after Stop") + } + + done := make(chan struct{}) + go func() { defer close(done); scanner.Close() }() + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatal("Close blocked after Stop") + } +} + +// GetOpenPorts returns an empty map rather than nil when no pre-scan ran, so +// callers can range over it unconditionally. +func TestCompat_GetOpenPorts(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + ports := scanner.GetOpenPorts() + if ports == nil { + t.Fatal("GetOpenPorts returned nil") + } + if len(ports) != 0 { + t.Errorf("GetOpenPorts returned %d hosts without a port pre-scan", len(ports)) + } + + if err := scanner.Run(); err != nil { + t.Fatalf("Run: %v", err) + } + if got := scanner.GetOpenPorts(); got == nil { + t.Error("GetOpenPorts returned nil after the scan") + } +} + +// SetProxy re-initialises the shared HTTP client, so it must not break a scan +// when pointed back at no proxy. +func TestCompat_SetProxy(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + scanner.SetProxy("") + if scanner.opts.Proxy != "" { + t.Errorf("Proxy = %q, want empty", scanner.opts.Proxy) + } + + if err := scanner.Run(); err != nil { + t.Fatalf("Run after SetProxy: %v", err) + } + if got := scanner.GetVulnerabilityCount(); got != 1 { + t.Errorf("found %d results, want 1", got) + } +} + +// The OOB helpers must report "not configured" without probing anything when +// OOB was never enabled. +func TestCompat_OOBReportsDisabledWhenUnconfigured(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if scanner.IsOOBEnabled() { + t.Error("IsOOBEnabled = true without any OOB configuration") + } + enabled, status := scanner.GetOOBStatus() + if enabled { + t.Error("GetOOBStatus reported enabled without configuration") + } + if strings.TrimSpace(status) == "" { + t.Error("GetOOBStatus returned an empty description") + } +} + +// IsOOBEnabled keeps its original shape: the explicit flag wins, and failing +// that a configured adapter plus a key or domain counts as enabled. +func TestCompat_IsOOBEnabledMatchesOriginalRules(t *testing.T) { + tests := []struct { + name string + opts *SDKOptions + want bool + }{ + {"nothing set", &SDKOptions{}, false}, + {"explicit flag", &SDKOptions{EnableOOB: true}, true}, + {"adapter only", &SDKOptions{OOB: "ceyeio"}, false}, + {"adapter and key", &SDKOptions{OOB: "ceyeio", OOBKey: "k"}, true}, + {"adapter and domain", &SDKOptions{OOB: "ceyeio", OOBDomain: "d"}, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := &SDKScanner{opts: tt.opts} + if got := s.IsOOBEnabled(); got != tt.want { + t.Errorf("IsOOBEnabled = %v, want %v", got, tt.want) + } + }) + } +} + +// Run prints the scan summary the way the original SDK did, and Silent turns +// that off for callers that want the quiet behaviour of the new API. +func TestCompat_PrintsSummaryUnlessSilent(t *testing.T) { + srv := newCompatServer(t) + + tests := []struct { + name string + silent bool + wantOutput bool + }{ + {"default prints", false, true}, + {"silent is quiet", true, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + opts := newCompatOptions(t, srv.URL) + opts.Silent = tt.silent + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + stdout, r, w := captureStdout(t) + runErr := scanner.Run() + output := restoreStdout(t, stdout, r, w) + + if runErr != nil { + t.Fatalf("Run: %v", runErr) + } + if tt.wantOutput && !strings.Contains(output, "扫描信息") { + t.Errorf("expected the scan summary on stdout, got:\n%s", output) + } + if !tt.wantOutput && output != "" { + t.Errorf("Silent still wrote to stdout:\n%s", output) + } + }) + } +} + +// A second Run on the same scanner must report the single-use error rather +// than silently doing nothing. +func TestCompat_SecondRunIsRejected(t *testing.T) { + srv := newCompatServer(t) + scanner, err := NewSDKScanner(newCompatOptions(t, srv.URL)) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + if err := scanner.Run(); err != nil { + t.Fatalf("first Run: %v", err) + } + if err := scanner.Run(); err == nil { + t.Error("second Run should have failed on a single-use scanner") + } +} + // The facade must expose the current SDK for callers that want to migrate // gradually. func TestCompat_ScannerAccessorExposesTheModernAPI(t *testing.T) { diff --git a/docs/SDK_Usage_Guide_English.md b/docs/SDK_Usage_Guide_English.md index 2281637b6..86369d8f8 100644 --- a/docs/SDK_Usage_Guide_English.md +++ b/docs/SDK_Usage_Guide_English.md @@ -1,5 +1,15 @@ # Afrog SDK Usage Guide +> ## ⚠️ Upgrade notice +> +> **This is a substantial change.** The SDK has been rebuilt as a dedicated `pkg/sdk` package with functional options and a `context`-driven lifecycle, and a number of defects were fixed along the way: leaked goroutines, data races, and PoCs being dropped without a word. +> +> **Migrating to `pkg/sdk` is recommended for both new and existing projects.** It is materially better on type safety, resource cleanup, error visibility and data completeness, and it is what the rest of this guide documents. See "Two APIs, one implementation" below for the mapping. +> +> **Existing code keeps working untouched.** The root package `github.com/zan8in/afrog/v3` still exposes the complete original surface (`NewSDKOptions`, `NewSDKScanner` and every method and field). Entry points and call patterns are **unchanged**; the improvements and fixes happen inside. We compared the root package's `go doc` output before and after item by item: **nothing was removed and no signature changed** — only optional fields were added. +> +> One old behaviour was corrected: setting `PocFile` and `AppendPoc` together used to drop `AppendPoc` silently, and both are now loaded. That is a bug fix, not an interface change. + ## Overview The Afrog SDK is the Go API for embedding vulnerability scanning into your own programs. The import path is `github.com/zan8in/afrog/v3/pkg/sdk`. diff --git "a/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" "b/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" index b933bd50b..67b349e3b 100644 --- "a/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" +++ "b/docs/SDK\344\275\277\347\224\250\346\214\207\345\215\227_\344\270\255\346\226\207.md" @@ -1,5 +1,15 @@ # Afrog SDK 使用指南 +> ## ⚠️ 升级提示 +> +> **本次改动幅度较大**：SDK 已重构为独立的 `pkg/sdk` 包，采用函数式选项与 `context` 驱动的生命周期，同时修复了协程泄漏、数据竞争、PoC 被静默丢弃等一批问题。 +> +> **建议新老项目都迁移到 `pkg/sdk`。** 新接口在类型安全、资源释放、错误可观测性和数据完整性上都明显更好，本文档也以它为主。迁移对照见下方「两套 API 共存」一节。 +> +> **旧代码无需改动即可继续运行。** 根包 `github.com/zan8in/afrog/v3` 保留了完整的旧接口（`NewSDKOptions`、`NewSDKScanner` 及全部方法与字段），入口函数与调用方式**一律保持原样**，优化和缺陷修复都在内部完成。我们用 `go doc` 对改动前后的根包 API 做了逐项机器比对，**旧 API 面零删除、零签名变更**，只新增了可选字段。 +> +> 唯一被修正的旧行为：同时指定 `PocFile` 与 `AppendPoc` 时，旧版会静默丢弃 `AppendPoc`，现在两者都会加载。这是缺陷修复，不是接口变更。 + ## 概述 Afrog SDK 是把漏洞扫描能力嵌入自己程序的 Go 接口，包路径为 `github.com/zan8in/afrog/v3/pkg/sdk`。 diff --git a/pkg/sdk/coverage_test.go b/pkg/sdk/coverage_test.go new file mode 100644 index 000000000..e4287a497 --- /dev/null +++ b/pkg/sdk/coverage_test.go @@ -0,0 +1,345 @@ +package sdk + +import ( + "context" + "errors" + "testing" + "time" +) + +// Every option has to actually change the configuration. A typo in one of the +// one-line option functions would otherwise be invisible: the option compiles, +// the scan runs, and the setting is silently ignored. +func TestOptions_EveryOptionApplies(t *testing.T) { + tests := []struct { + name string + option Option + check func(*Options) bool + }{ + {"WithTargets", WithTargets("a", "b"), func(o *Options) bool { return len(o.Targets) == 2 }}, + {"WithTargetsFile", WithTargetsFile("t.txt"), func(o *Options) bool { return o.TargetsFile == "t.txt" }}, + {"WithPocPaths", WithPocPaths("p1", "p2"), func(o *Options) bool { return len(o.PocPaths) == 2 }}, + {"WithPocPathsOnly", WithPocPathsOnly(), func(o *Options) bool { return o.PocPathsOnly }}, + {"WithSearch", WithSearch("tomcat"), func(o *Options) bool { return o.Search == "tomcat" }}, + {"WithSeverity", WithSeverity("high"), func(o *Options) bool { return o.Severity == "high" }}, + {"WithExcludePocs", WithExcludePocs("x", "y"), func(o *Options) bool { return len(o.ExcludePocs) == 2 }}, + {"WithExcludePocsFile", WithExcludePocsFile("e.txt"), func(o *Options) bool { return o.ExcludePocsFile == "e.txt" }}, + + {"WithConcurrency", WithConcurrency(7), func(o *Options) bool { return o.Concurrency == 7 }}, + {"WithRateLimit", WithRateLimit(11), func(o *Options) bool { return o.RateLimit == 11 }}, + {"WithTimeout", WithTimeout(13), func(o *Options) bool { return o.Timeout == 13 }}, + {"WithRetries", WithRetries(4), func(o *Options) bool { return o.Retries == 4 }}, + {"WithMaxHostError", WithMaxHostError(9), func(o *Options) bool { return o.MaxHostError == 9 }}, + {"WithMaxRespBodySize", WithMaxRespBodySize(5), func(o *Options) bool { return o.MaxRespBodySize == 5 }}, + {"WithRequestLimitPerTarget", WithRequestLimitPerTarget(6), func(o *Options) bool { return o.ReqLimitPerTarget == 6 }}, + {"WithPolite", WithPolite(), func(o *Options) bool { return o.Polite }}, + {"WithBalanced", WithBalanced(), func(o *Options) bool { return o.Balanced }}, + {"WithAggressive", WithAggressive(), func(o *Options) bool { return o.Aggressive }}, + {"WithAutoRequestLimit", WithAutoRequestLimit(), func(o *Options) bool { return o.AutoReqLimit }}, + {"WithSmartConcurrency", WithSmartConcurrency(), func(o *Options) bool { return o.Smart }}, + {"WithStopOnFirstMatch", WithStopOnFirstMatch(), func(o *Options) bool { return o.StopOnFirstMatch }}, + + {"WithFingerprintDisabled", WithFingerprintDisabled(), func(o *Options) bool { return o.DisableFingerprint }}, + {"WithFingerprintFilterMode", WithFingerprintFilterMode(FingerprintOpportunistic), + func(o *Options) bool { return o.FingerprintFilterMode == FingerprintOpportunistic }}, + {"WithWebProbe", WithWebProbe(), func(o *Options) bool { return o.EnableWebProbe }}, + + {"WithProxy", WithProxy("http://127.0.0.1:8080"), func(o *Options) bool { return o.Proxy != "" }}, + {"WithHeaders", WithHeaders("X-A: 1"), func(o *Options) bool { return len(o.Headers) == 1 }}, + + {"WithPortScan", WithPortScan(PortScanOptions{Ports: "80"}), + func(o *Options) bool { return o.EnablePortSan && o.PortScan.Ports == "80" }}, + {"WithOOB", WithOOB(OOBOptions{Adapter: "ceyeio"}), + func(o *Options) bool { return o.OOB.Enabled && o.OOB.Adapter == "ceyeio" }}, + {"WithCurated", WithCurated(CuratedOptions{Enabled: "auto"}), + func(o *Options) bool { return o.Curated.Enabled == "auto" }}, + + {"WithDingtalk", WithDingtalk(), func(o *Options) bool { return o.Dingtalk }}, + {"WithWecom", WithWecom(), func(o *Options) bool { return o.Wecom }}, + + {"WithRequestResponse", WithRequestResponse(false), func(o *Options) bool { return !o.IncludeRequestResponse }}, + {"WithMaxStoredResults", WithMaxStoredResults(3), func(o *Options) bool { return o.MaxStoredResults == 3 }}, + {"WithStreamBuffer", WithStreamBuffer(8), func(o *Options) bool { return o.StreamBuffer == 8 }}, + {"WithRedactedHeaders", WithRedactedHeaders(), func(o *Options) bool { return len(o.RedactedHeaders) > 0 }}, + {"WithVerbose", WithVerbose(), func(o *Options) bool { return o.Verbose }}, + + {"WithTaskTimeout", WithTaskTimeout(TaskTimeoutOptions{HardSec: 12}), + func(o *Options) bool { return o.TaskTimeout.HardSec == 12 }}, + {"WithExecutionMonitor", WithExecutionMonitor(ExecutionMonitorOptions{LogLimit: 2}), + func(o *Options) bool { return o.EnableMonitor && o.Monitor.LogLimit == 2 }}, + {"WithCheckpoint", WithCheckpoint(CheckpointOptions{Path: "c.afg"}), + func(o *Options) bool { return o.Checkpoint.Path == "c.afg" }}, + {"WithCyberspace", WithCyberspace(CyberspaceOptions{Engine: CyberspaceZoomEye, Query: "q"}), + func(o *Options) bool { return o.Cyberspace.Engine == CyberspaceZoomEye }}, + {"WithTargetPreProbe", WithTargetPreProbe(), func(o *Options) bool { return o.TargetPreProbe }}, + + {"WithResultHandler", WithResultHandler(func(Result) {}), + func(o *Options) bool { return len(o.handlers.result) == 1 }}, + {"WithRawResultHandler", WithRawResultHandler(nil), + func(o *Options) bool { return len(o.handlers.rawResult) == 0 }}, + {"WithFailureHandler", WithFailureHandler(func(Failure) {}), + func(o *Options) bool { return len(o.handlers.failure) == 1 }}, + {"WithPortHandler", WithPortHandler(func(PortEvent) {}), + func(o *Options) bool { return len(o.handlers.port) == 1 }}, + {"WithHostHandler", WithHostHandler(func(HostEvent) {}), + func(o *Options) bool { return len(o.handlers.host) == 1 }}, + {"WithWebProbeHandler", WithWebProbeHandler(func(WebProbeEvent) {}), + func(o *Options) bool { return len(o.handlers.webProbe) == 1 }}, + {"WithProgressHandler", WithProgressHandler(func(PhaseProgress) {}), + func(o *Options) bool { return len(o.handlers.progress) == 1 }}, + {"WithScanInfoHandler", WithScanInfoHandler(func(ScanInfo) {}), + func(o *Options) bool { return len(o.handlers.scanInfo) == 1 }}, + {"WithMonitorHandler", WithMonitorHandler(func(string) {}), + func(o *Options) bool { return len(o.handlers.monitor) == 1 }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + o := NewOptions() + if err := tt.option(o); err != nil { + t.Fatalf("%s: %v", tt.name, err) + } + if !tt.check(o) { + t.Errorf("%s did not take effect", tt.name) + } + }) + } +} + +// The options that map straight onto engine fields must survive the +// translation, not just land on the SDK Options struct. +func TestScanner_OptionsReachTheEngine(t *testing.T) { + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "reach.yaml", "sdk-test-reach") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + WithConcurrency(7), + WithRateLimit(11), + WithTimeout(13), + WithRetries(4), + WithMaxHostError(9), + WithMaxRespBodySize(5), + WithRequestLimitPerTarget(6), + WithSmartConcurrency(), + WithWebProbe(), + WithSearch("reach"), + WithHeaders("X-Test: 1"), + WithStopOnFirstMatch(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + in := scanner.internal + checks := []struct { + name string + got any + want any + }{ + {"Concurrency", in.Concurrency, 7}, + {"RateLimit", in.RateLimit, 11}, + {"Timeout", in.Timeout, 13}, + {"Retries", in.Retries, 4}, + {"MaxHostError", in.MaxHostError, 9}, + {"MaxRespBodySize", in.MaxRespBodySize, 5}, + {"ReqLimitPerTarget", in.ReqLimitPerTarget, 6}, + {"Smart", in.Smart, true}, + {"EnableWebProbe", in.EnableWebProbe, true}, + {"Search", in.Search, "reach"}, + {"DisableFingerprint", in.DisableFingerprint, true}, + {"StopOnFirstMatch", in.VulnerabilityScannerBreakpoint, true}, + {"SDKMode", in.SDKMode, true}, + {"Silent", in.Silent, true}, + {"DisableUpdateCheck", in.DisableUpdateCheck, true}, + } + for _, c := range checks { + if c.got != c.want { + t.Errorf("engine %s = %v, want %v", c.name, c.got, c.want) + } + } + if len(in.Header) != 1 { + t.Errorf("engine Header has %d entries, want 1", len(in.Header)) + } + // A library must never be configured to write report files. + if in.Json != "" || in.JsonAll != "" || in.Output != "" { + t.Error("engine is configured to write report files") + } +} + +// The accessors and the less-used streams must behave sensibly before, during +// and after a scan. +func TestScanner_AccessorsAndRemainingStreams(t *testing.T) { + scanner, _ := newTestScanner(t) + + if scanner.IsRunning() { + t.Error("IsRunning = true before Start") + } + if scanner.HasResults() { + t.Error("HasResults = true before the scan") + } + if scanner.CuratedError() != nil { + t.Errorf("CuratedError = %v with curated disabled", scanner.CuratedError()) + } + if enabled, status := scanner.OOBStatus(); enabled || status != "disabled" { + t.Errorf("OOBStatus = (%v, %q), want (false, \"disabled\")", enabled, status) + } + + scanner.Pause() + if !scanner.IsPaused() { + t.Error("IsPaused = false after Pause") + } + scanner.Resume() + if scanner.IsPaused() { + t.Error("IsPaused = true after Resume") + } + + // Subscribing before the scan means each stream must close when it ends, + // otherwise these range loops would never finish. + ports := scanner.PortStream() + hosts := scanner.HostStream() + probes := scanner.WebProbeStream() + progress := scanner.ProgressStream() + infos := scanner.ScanInfoStream() + + drained := make(chan struct{}) + go func() { + defer close(drained) + for range ports { + } + for range hosts { + } + for range probes { + } + for range progress { + } + for range infos { + } + }() + + if err := scanner.Execute(context.Background()); err != nil { + t.Fatalf("Execute: %v", err) + } + + select { + case <-drained: + case <-time.After(30 * time.Second): + t.Fatal("a stream was left open after the scan finished") + } + + if !scanner.HasResults() { + t.Error("HasResults = false after a finding") + } + if scanner.IsRunning() { + t.Error("IsRunning = true after the scan finished") + } + if scanner.OpenPorts() == nil { + t.Error("OpenPorts returned nil") + } +} + +// The error types have to unwrap so callers can use errors.Is and errors.As. +func TestErrors_WrapAndUnwrap(t *testing.T) { + base := errors.New("boom") + + mount := &CuratedMountError{Err: base} + if !errors.Is(mount, base) { + t.Error("CuratedMountError does not unwrap to its cause") + } + if mount.Error() == "" { + t.Error("CuratedMountError.Error() is empty") + } + + f := Failure{Target: "t", PocID: "p", Err: base} + if !errors.Is(f, base) { + t.Error("Failure does not unwrap to its cause") + } + if f.Error() != "boom" { + t.Errorf("Failure.Error() = %q, want %q", f.Error(), "boom") + } + + // A Failure with no cause must not panic and must read as empty. + var empty Failure + if empty.Error() != "" { + t.Errorf("empty Failure.Error() = %q, want empty", empty.Error()) + } +} + +// Progress before the scan finishes goes through the weighted path, which is +// skipped entirely by the "finished cleanly is 100%" shortcut. +func TestScanner_ProgressBeforeFinishIsWeighted(t *testing.T) { + srv := newTestServer(t) + dir := t.TempDir() + writePoc(t, dir, "weighted.yaml", "sdk-test-weighted") + + scanner, err := New(context.Background(), + WithTargets(srv.URL), + WithPocPaths(dir), + WithPocPathsOnly(), + WithFingerprintDisabled(), + // Both stages carry weight, so the port and web-probe phases are read. + WithPortScan(PortScanOptions{Ports: "80", SkipDiscovery: true}), + WithWebProbe(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = scanner.Close() }) + + got := scanner.Progress() + if got < 0 || got > 100 { + t.Fatalf("Progress = %.2f before the scan, want within [0, 100]", got) + } +} + +// recordOpenPort builds a nested map, so the same host reported twice must +// collect both ports rather than replacing the first. +func TestScanner_RecordOpenPort(t *testing.T) { + scanner, _ := newTestScanner(t) + + scanner.recordOpenPort("10.0.0.1", 80) + scanner.recordOpenPort("10.0.0.1", 443) + scanner.recordOpenPort("10.0.0.1", 80) // duplicate + scanner.recordOpenPort("10.0.0.2", 22) + + ports := scanner.OpenPorts() + if len(ports) != 2 { + t.Fatalf("OpenPorts has %d hosts, want 2", len(ports)) + } + if len(ports["10.0.0.1"]) != 2 { + t.Errorf("10.0.0.1 has %v, want two distinct ports", ports["10.0.0.1"]) + } + if len(ports["10.0.0.2"]) != 1 { + t.Errorf("10.0.0.2 has %v, want one port", ports["10.0.0.2"]) + } +} + +func TestClampFloat(t *testing.T) { + tests := []struct { + in float64 + want float64 + }{ + {-1, 0}, {0, 0}, {50.5, 50.5}, {100, 100}, {101, 100}, + } + for _, tt := range tests { + if got := clampFloat(tt.in); got != tt.want { + t.Errorf("clampFloat(%v) = %v, want %v", tt.in, got, tt.want) + } + } +} + +// Fingerprint hits are turned into results, which is a separate path from +// ordinary PoC matches. +func TestScanner_FingerprintHitsBecomeResults(t *testing.T) { + scanner, _ := newTestScanner(t) + + scanner.handleFingerprint("key", nil) + if got := scanner.ResultCount(); got != 0 { + t.Fatalf("an empty hit list produced %d results", got) + } +} From ef76d51748ce52764874f5715e60f02a677582ec Mon Sep 17 00:00:00 2001 From: zhizhuo Date: Tue, 11 Aug 2026 22:58:28 +0800 Subject: [PATCH 5/5] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=85=BC=E5=AE=B9?= =?UTF-8?q?=E5=B1=82=E4=B8=A2=E5=A4=B1=20BruteMaxRequests=20=E4=B8=8E=20De?= =?UTF-8?q?faultAccept=20=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 审查兼容层每个字段的接线情况时发现：SDKOptions.BruteMaxRequests 和 DefaultAccept 从未被传递到引擎。旧代码设置这两个字段会被静默忽略， 是编译期发现不了的兼容性破坏。 根因是 pkg/sdk 有这两个 Options 字段却没有对应的选项函数，兼容层 无从设置。现补上 WithBruteMaxRequests 与 WithDefaultAccept，并在 兼容层接线。 同时补一组守护测试，逐字段断言 SDKOptions 的每一项都真正到达引擎， 覆盖基础配置、端口预扫描、OOB 与新增能力四个块。回退接线后测试会 报 "BruteMaxRequests did not reach the engine: got 5000, want 123"， 确认不是空测。 新增 sdk.Scanner.EngineOptions() 作为只读逃生舱，让兼容层可以验证 配置确实落到引擎，而不是靠肉眼比对。 覆盖率：根包 81.4% -> 85.4%，pkg/sdk 82.7% --- afrog.go | 4 + afrog_test.go | 196 +++++++++++++++++++++++++++++++++++++++++++++ pkg/sdk/options.go | 22 +++++ pkg/sdk/scanner.go | 10 +++ 4 files changed, 232 insertions(+) diff --git a/afrog.go b/afrog.go index f8525b505..8ffc60611 100644 --- a/afrog.go +++ b/afrog.go @@ -338,6 +338,10 @@ func (s *SDKScanner) build() (*sdk.Scanner, error) { if o.MaxHostError >= 0 { options = append(options, sdk.WithMaxHostError(o.MaxHostError)) } + if o.BruteMaxRequests >= 0 { + options = append(options, sdk.WithBruteMaxRequests(o.BruteMaxRequests)) + } + options = append(options, sdk.WithDefaultAccept(o.DefaultAccept)) if strings.TrimSpace(o.TargetsFile) != "" { options = append(options, sdk.WithTargetsFile(o.TargetsFile)) } diff --git a/afrog_test.go b/afrog_test.go index a673dd959..144836941 100644 --- a/afrog_test.go +++ b/afrog_test.go @@ -683,6 +683,202 @@ func TestCompat_SecondRunIsRejected(t *testing.T) { } } +// A field can sit on SDKOptions, compile fine, and still never reach the +// engine. That failure is invisible: the caller sets it, the scan runs, and +// the setting is ignored. This pins every field that maps onto an engine +// setting to the value the engine actually ends up with. +func TestCompat_EveryOptionReachesTheEngine(t *testing.T) { + srv := newCompatServer(t) + dir := t.TempDir() + writeCompatPoc(t, dir, "reach") + + opts := NewSDKOptions() + opts.Silent = true + opts.Targets = []string{srv.URL} + opts.PocFile = dir + opts.DisableFingerprint = true + + opts.Concurrency = 7 + opts.RateLimit = 11 + opts.Timeout = 13 + opts.Retries = 4 + opts.MaxHostError = 9 + opts.MaxRespBodySize = 5 + opts.BruteMaxRequests = 123 + opts.DefaultAccept = false + opts.ReqLimitPerTarget = 6 + opts.Smart = true + opts.EnableWebProbe = true + opts.Search = "reach" + opts.Severity = "info" + opts.ExcludePocs = []string{"nope"} + opts.Headers = []string{"X-Test: 1"} + opts.VulnerabilityScannerBreakpoint = true + opts.FingerprintFilterMode = "opportunistic" + opts.Proxy = "" + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + in := scanner.Scanner().EngineOptions() + checks := []struct { + field string + got any + want any + }{ + {"Concurrency", in.Concurrency, 7}, + {"RateLimit", in.RateLimit, 11}, + {"Timeout", in.Timeout, 13}, + {"Retries", in.Retries, 4}, + {"MaxHostError", in.MaxHostError, 9}, + {"MaxRespBodySize", in.MaxRespBodySize, 5}, + {"BruteMaxRequests", in.BruteMaxRequests, 123}, + {"DefaultAccept", in.DefaultAccept, false}, + {"ReqLimitPerTarget", in.ReqLimitPerTarget, 6}, + {"Smart", in.Smart, true}, + {"EnableWebProbe", in.EnableWebProbe, true}, + {"Search", in.Search, "reach"}, + {"Severity", in.Severity, "info"}, + {"DisableFingerprint", in.DisableFingerprint, true}, + {"FingerprintFilterMode", in.FingerprintFilterMode, "opportunistic"}, + {"VulnerabilityScannerBreakpoint", in.VulnerabilityScannerBreakpoint, true}, + {"PocPathsOnly", in.PocPathsOnly, true}, // PocFile keeps its exclusive meaning + } + for _, c := range checks { + if c.got != c.want { + t.Errorf("SDKOptions.%s did not reach the engine: got %v, want %v", c.field, c.got, c.want) + } + } + if len(in.Header) != 1 { + t.Errorf("Headers did not reach the engine: %v", in.Header) + } + if len(in.ExcludePocs) != 1 { + t.Errorf("ExcludePocs did not reach the engine: %v", in.ExcludePocs) + } +} + +// The port pre-scan block is a nested struct in the old options, so each of +// its fields needs its own mapping. +func TestCompat_PortScanOptionsReachTheEngine(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + opts.PortScan = true + opts.PSPorts = "80,443" + opts.PSRateLimit = 33 + opts.PSTimeout = 700 + opts.PSRetries = 2 + opts.PSSkipDiscovery = true + opts.PSS4Chunk = 250 + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + in := scanner.Scanner().EngineOptions() + checks := []struct { + field string + got any + want any + }{ + {"PortScan", in.PortScan, true}, + {"PSPorts", in.PSPorts, "80,443"}, + {"PSRateLimit", in.PSRateLimit, 33}, + {"PSTimeout", in.PSTimeout, 700}, + {"PSRetries", in.PSRetries, 2}, + {"PSSkipDiscovery", in.PSSkipDiscovery, true}, + {"PSS4Chunk", in.PSS4Chunk, 250}, + } + for _, c := range checks { + if c.got != c.want { + t.Errorf("SDKOptions.%s did not reach the engine: got %v, want %v", c.field, c.got, c.want) + } + } +} + +// The OOB block likewise has to arrive field by field. +func TestCompat_OOBOptionsReachTheEngine(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + opts.EnableOOB = true + opts.OOB = "ceyeio" + opts.OOBKey = "test-key" + opts.OOBDomain = "test.example.com" + opts.OOBRateLimit = 8 + opts.OOBConcurrency = 9 + opts.OOBFinalizeTimeout = 15 + opts.OOBPollInterval = 4 + opts.OOBHitRetention = 6 + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + in := scanner.Scanner().EngineOptions() + checks := []struct { + field string + got any + want any + }{ + {"EnableOOB", in.EnableOOB, true}, + {"OOB", in.OOB, "ceyeio"}, + {"OOBKey", in.OOBKey, "test-key"}, + {"OOBDomain", in.OOBDomain, "test.example.com"}, + {"OOBRateLimit", in.OOBRateLimit, 8}, + {"OOBConcurrency", in.OOBConcurrency, 9}, + {"OOBFinalizeTimeout", in.OOBFinalizeTimeout, 15}, + {"OOBPollInterval", in.OOBPollInterval, 4}, + {"OOBHitRetention", in.OOBHitRetention, 6}, + } + for _, c := range checks { + if c.got != c.want { + t.Errorf("SDKOptions.%s did not reach the engine: got %v, want %v", c.field, c.got, c.want) + } + } + if !scanner.IsOOBEnabled() { + t.Error("IsOOBEnabled = false with OOB configured") + } +} + +// The capabilities added after the facade must reach the engine too. +func TestCompat_NewOptionsReachTheEngine(t *testing.T) { + srv := newCompatServer(t) + opts := newCompatOptions(t, srv.URL) + opts.TaskHardTimeoutSec = 45 + opts.TaskSmartTimeout = true + opts.MonitorTargets = true + opts.ResumeFile = filepath.Join(t.TempDir(), "reach.afg") + + scanner, err := NewSDKScanner(opts) + if err != nil { + t.Fatalf("NewSDKScanner: %v", err) + } + defer scanner.Close() + + in := scanner.Scanner().EngineOptions() + checks := []struct { + field string + got any + want any + }{ + {"TaskHardTimeoutSec", in.TaskHardTimeoutSec, 45}, + {"TaskSmartTimeout", in.TaskSmartTimeout, true}, + {"MonitorTargets", in.MonitorTargets, true}, + {"Resume", in.Resume, opts.ResumeFile}, + } + for _, c := range checks { + if c.got != c.want { + t.Errorf("SDKOptions.%s did not reach the engine: got %v, want %v", c.field, c.got, c.want) + } + } +} + // The facade must expose the current SDK for callers that want to migrate // gradually. func TestCompat_ScannerAccessorExposesTheModernAPI(t *testing.T) { diff --git a/pkg/sdk/options.go b/pkg/sdk/options.go index f09483a53..211bdd73a 100644 --- a/pkg/sdk/options.go +++ b/pkg/sdk/options.go @@ -686,6 +686,28 @@ func WithMaxRespBodySize(mb int) Option { } } +// WithBruteMaxRequests caps how many requests a brute-force rule may send. +// Zero disables the cap. Exceeding it sets Exchange.BruteTruncated. +func WithBruteMaxRequests(n int) Option { + return func(o *Options) error { + if n < 0 { + return fmt.Errorf("%w: brute max requests must be >= 0, got %d", ErrInvalidOptions, n) + } + o.BruteMaxRequests = n + return nil + } +} + +// WithDefaultAccept controls whether requests carry a default Accept header. +// It is enabled by default; disable it for targets that behave differently +// when one is present. +func WithDefaultAccept(enabled bool) Option { + return func(o *Options) error { + o.DefaultAccept = enabled + return nil + } +} + // WithRequestLimitPerTarget caps concurrent requests per target. func WithRequestLimitPerTarget(n int) Option { return func(o *Options) error { diff --git a/pkg/sdk/scanner.go b/pkg/sdk/scanner.go index 8716af31b..731b80c3e 100644 --- a/pkg/sdk/scanner.go +++ b/pkg/sdk/scanner.go @@ -899,6 +899,16 @@ func (s *Scanner) PocDiagnostics() []config.PocLoadError { return out } +// EngineOptions returns the engine configuration this scanner resolved from +// its options. +// +// This is an escape hatch for advanced integrations — notably the +// backward-compatible facade in the root package, which uses it to verify that +// every legacy option still reaches the engine. Treat the result as read-only: +// the scanner keeps using it, and config.Options is an internal detail whose +// shape may change between releases. +func (s *Scanner) EngineOptions() *config.Options { return s.internal } + // CuratedError reports why the optional curated PoC source failed to mount. // It is never fatal; nil means the source mounted or was disabled. func (s *Scanner) CuratedError() error { return s.curatedErr }