using duration var

projectdiscovery · Sep 16, 2024 · 8e48bde · 8e48bde
1 parent 4e33c0c
commit 8e48bde
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/runner/options.go b/runner/options.go
@@ -100,8 +100,8 @@ type ScanOptions struct {
 	DisableStdin              bool
 	NoScreenshotBytes         bool
 	NoHeadlessBody            bool
-	ScreenshotTimeout         int
-	ScreenshotIdle            int
+	ScreenshotTimeout         time.Duration
+	ScreenshotIdle            time.Duration
 }
 
 func (s *ScanOptions) Clone() *ScanOptions {
@@ -306,8 +306,8 @@ type Options struct {
 	HttpApiEndpoint    string
 	NoScreenshotBytes  bool
 	NoHeadlessBody     bool
-	ScreenshotTimeout  int
-	ScreenshotIdle     int
+	ScreenshotTimeout  time.Duration
+	ScreenshotIdle     time.Duration
 	// HeadlessOptionalArguments specifies optional arguments to pass to Chrome
 	HeadlessOptionalArguments goflags.StringSlice
 	Protocol                  string
@@ -376,8 +376,8 @@ func ParseOptions() *Options {
 		flagSet.StringSliceVarP(&options.HeadlessOptionalArguments, "headless-options", "ho", nil, "start headless chrome with additional options", goflags.FileCommaSeparatedStringSliceOptions),
 		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.DurationVarP(&options.ScreenshotTimeout, "screenshot-timeout", "st", 10*time.Second, "set timeout for screenshot in seconds"),
+		flagSet.DurationVarP(&options.ScreenshotIdle, "screenshot-idle", "sid", 1*time.Second, "set idle time before taking screenshot in seconds"),
 	)
 
 	flagSet.CreateGroup("matchers", "Matchers",
@@ -629,7 +629,7 @@ func (options *Options) ValidateOptions() error {
 				msg += fmt.Sprintf("%s flag is", last)
 			}
 			msg += " incompatible with silent flag"
-			return fmt.Errorf(msg)
+			return errors.New(msg)
 		}
 	}
 

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, time.Duration(scanopts.ScreenshotIdle)*time.Second, r.options.CustomHeaders)
+		screenshotBytes, headlessBody, err = r.browser.ScreenshotWithBody(fullURL, scanopts.ScreenshotTimeout, scanopts.ScreenshotIdle, r.options.CustomHeaders)
 		if err != nil {
 			gologger.Warning().Msgf("Could not take screenshot '%s': %s", fullURL, err)
 		} else {