diff --git a/runner/headless.go b/runner/headless.go index c298cb76..481fd1cb 100644 --- a/runner/headless.go +++ b/runner/headless.go @@ -99,7 +99,7 @@ func NewBrowser(proxy string, useLocal bool, optionalArgs map[string]string) (*B return engine, nil } -func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, headers []string) ([]byte, string, error) { +func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, idle time.Duration, headers []string) ([]byte, string, error) { page, err := b.engine.Page(proto.TargetCreateTarget{}) if err != nil { return nil, "", err @@ -126,7 +126,7 @@ func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, headers if err := page.WaitLoad(); err != nil { return nil, "", err } - _ = page.WaitIdle(1 * time.Second) + _ = page.WaitIdle(idle) screenshot, err := page.Screenshot(true, &proto.PageCaptureScreenshot{}) if err != nil { diff --git a/runner/options.go b/runner/options.go index 1a0265b9..ce85e7b5 100644 --- a/runner/options.go +++ b/runner/options.go @@ -101,6 +101,7 @@ type ScanOptions struct { NoScreenshotBytes bool NoHeadlessBody bool ScreenshotTimeout int + ScreenshotIdle int } func (s *ScanOptions) Clone() *ScanOptions { @@ -154,6 +155,7 @@ func (s *ScanOptions) Clone() *ScanOptions { NoScreenshotBytes: s.NoScreenshotBytes, NoHeadlessBody: s.NoHeadlessBody, ScreenshotTimeout: s.ScreenshotTimeout, + ScreenshotIdle: s.ScreenshotIdle, } } @@ -305,6 +307,7 @@ type Options struct { NoScreenshotBytes bool NoHeadlessBody bool ScreenshotTimeout int + ScreenshotIdle int // HeadlessOptionalArguments specifies optional arguments to pass to Chrome HeadlessOptionalArguments goflags.StringSlice Protocol string @@ -374,6 +377,7 @@ func ParseOptions() *Options { flagSet.BoolVarP(&options.NoScreenshotBytes, "exclude-screenshot-bytes", "esb", false, "enable excluding screenshot bytes from json output"), flagSet.BoolVarP(&options.NoHeadlessBody, "exclude-headless-body", "ehb", false, "enable excluding headless header from json output"), flagSet.IntVarP(&options.ScreenshotTimeout, "screenshot-timeout", "st", 10, "set timeout for screenshot in seconds"), + flagSet.IntVarP(&options.ScreenshotIdle, "screenshot-idle", "sid", 1, "set idle time before taking screenshot in seconds"), ) flagSet.CreateGroup("matchers", "Matchers", diff --git a/runner/runner.go b/runner/runner.go index 6a3511a0..49b2eed9 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -2144,7 +2144,7 @@ retry: var pHash uint64 if scanopts.Screenshot { var err error - screenshotBytes, headlessBody, err = r.browser.ScreenshotWithBody(fullURL, time.Duration(scanopts.ScreenshotTimeout)*time.Second, r.options.CustomHeaders) + screenshotBytes, headlessBody, err = r.browser.ScreenshotWithBody(fullURL, time.Duration(scanopts.ScreenshotTimeout)*time.Second, time.Duration(scanopts.ScreenshotIdle)*time.Second, r.options.CustomHeaders) if err != nil { gologger.Warning().Msgf("Could not take screenshot '%s': %s", fullURL, err) } else {