--screenshot-filter option

cmd/scan.go

@@ -149,6 +149,7 @@ func init() {
 	scanCmd.PersistentFlags().BoolVar(&opts.Scan.SaveContent, "save-content", false, "Save content from network requests to the configured writers. WARNING: This flag has the potential to make your storage explode in size")
 	scanCmd.PersistentFlags().BoolVar(&opts.Scan.SkipHTML, "skip-html", false, "Don't include the first request's HTML response when writing results")
 	scanCmd.PersistentFlags().BoolVar(&opts.Scan.ScreenshotToWriter, "write-screenshots", false, "Store screenshots with writers in addition to filesystem storage")
+	scanCmd.PersistentFlags().IntSliceVar(&opts.Scan.ScreenshotCodes, "screenshot-filter", []int{}, "Http response codes to screenshot. this is a filter (by default all codes are screenshotted)")
 
 	// Chrome options
 	scanCmd.PersistentFlags().StringVar(&opts.Chrome.Path, "chrome-path", "", "The path to a Google Chrome binary to use (downloads a platform-appropriate binary by default)")

pkg/runner/drivers/chromedp.go

@@ -52,6 +52,16 @@ func (b *browserInstance) Close() {
 	os.RemoveAll(b.userData)
 }
 
+// SliceContainsInt ... returns true/false
+func SliceContainsInt(slice []int, num int) bool {
+    for _, v := range slice {
+        if v == num {
+            return true
+        }
+    }
+    return false
+}
+
 // getChromedpAllocator is a helper function to get a chrome allocation context.
 //
 // see Witness for more information on why we're explicitly not using tabs
@@ -385,6 +395,13 @@ func (run *Chromedp) Witness(target string, thisRunner *runner.Runner) (*models.
 		}
 	}
 
+	// check if the preflight returned a code to process.
+	// an empty slice implies no filtering
+	if (len(run.options.Scan.ScreenshotCodes) > 0) &&
+		!SliceContainsInt(run.options.Scan.ScreenshotCodes, result.ResponseCode) {
+		return nil, fmt.Errorf("response code (%w) not in allowed screenshot http response codes.", result.ResponseCode)
+	}
+
 	// grab the title
 	if err := chromedp.Run(navigationCtx, chromedp.Title(&result.Title)); err != nil {
 		if run.options.Logging.LogScanErrors {

pkg/runner/options.go

@@ -89,6 +89,9 @@ type Scan struct {
 	// Save content stores content from network requests (warning) this
 	// could make written artefacts huge
 	SaveContent bool
+	// ScreenshotCodes are http response codes to screenshot. this is a filter. 
+	// by default all codes are screenshotted
+	ScreenshotCodes []int
 }
 
 // NewDefaultOptions returns Options with some default values
@@ -105,6 +108,7 @@ func NewDefaultOptions() *Options {
 			Timeout:          60,
 			UriFilter:        []string{"http", "https"},
 			ScreenshotFormat: "jpeg",
+			ScreenshotCodes:  []int{},
 		},
 		Logging: Logging{
 			Debug:         true,