feat: add idle screenshot option

projectdiscovery · Sep 15, 2024 · 4e33c0c · 4e33c0c
1 parent d58ad9d
commit 4e33c0c
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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 {