JS injection support

runner/headless.go

@@ -110,11 +110,27 @@ func NewBrowser(proxy string, useLocal bool, optionalArgs map[string]string) (*B
 	return engine, nil
 }
 
-func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, idle time.Duration, headers []string, fullPage bool) ([]byte, string, []NetworkRequest, error) {
-	page, err := b.engine.Page(proto.TargetCreateTarget{})
+func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, idle time.Duration, headers []string, fullPage bool, jsCodes []string) ([]byte, string, []NetworkRequest, error) {
+	page, networkRequests, err := b.setupPageAndNavigate(url, timeout, headers, jsCodes)
 	if err != nil {
 		return nil, "", []NetworkRequest{}, err
 	}
+	defer b.closePage(page)
+
+	screenshot, body, err := b.takeScreenshotAndGetBody(page, idle, fullPage)
+	if err != nil {
+		return nil, "", networkRequests, err
+	}
+
+	return screenshot, body, networkRequests, nil
+}
+
+// setupPageAndNavigate opens a page, performs all adaptive actions including JS injection
+func (b *Browser) setupPageAndNavigate(url string, timeout time.Duration, headers []string, jsCodes []string) (*rod.Page, []NetworkRequest, error) {
+	page, err := b.engine.Page(proto.TargetCreateTarget{})
+	if err != nil {
+		return nil, []NetworkRequest{}, err
+	}
 
 	// Enable network
 	page.EnableDomain(proto.NetworkEnable{})
@@ -184,32 +200,46 @@ func (b *Browser) ScreenshotWithBody(url string, timeout time.Duration, idle tim
 	}
 
 	page = page.Timeout(timeout)
-	defer func() {
-		_ = page.Close()
-	}()
 
 	if err := page.Navigate(url); err != nil {
-		return nil, "", networkRequests.Slice, err
+		return page, networkRequests.Slice, err
+	}
+
+	if len(jsCodes) > 0 {
+		_, err := b.ExecuteJavascriptCodesWithPage(page, jsCodes)
+		if err != nil {
+			return page, networkRequests.Slice, err
+		}
 	}
 
 	page.Timeout(5 * time.Second).WaitNavigation(proto.PageLifecycleEventNameFirstMeaningfulPaint)()
 
+	return page, networkRequests.Slice, nil
+}
+
+// takeScreenshotAndGetBody performs the screenshot actions
+func (b *Browser) takeScreenshotAndGetBody(page *rod.Page, idle time.Duration, fullPage bool) ([]byte, string, error) {
 	if err := page.WaitLoad(); err != nil {
-		return nil, "", networkRequests.Slice, err
+		return nil, "", err
 	}
 	_ = page.WaitIdle(idle)
 
 	screenshot, err := page.Screenshot(fullPage, &proto.PageCaptureScreenshot{})
 	if err != nil {
-		return nil, "", networkRequests.Slice, err
+		return nil, "", err
 	}
 
 	body, err := page.HTML()
 	if err != nil {
-		return screenshot, "", networkRequests.Slice, err
+		return screenshot, "", err
 	}
 
-	return screenshot, body, networkRequests.Slice, nil
+	return screenshot, body, nil
+}
+
+// closePage closes the page and performs cleanup
+func (b *Browser) closePage(page *rod.Page) {
+	_ = page.Close()
 }
 
 func (b *Browser) Close() {
@@ -268,3 +298,18 @@ func getSimpleErrorType(errorText, errorType, blockedReason string) string {
 	}
 	return "UNKNOWN"
 }
+
+func (b *Browser) ExecuteJavascriptCodesWithPage(page *rod.Page, jsc []string) ([]*proto.RuntimeRemoteObject, error) {
+	outputs := make([]*proto.RuntimeRemoteObject, 0, len(jsc))
+	for _, js := range jsc {
+		if js == "" {
+			continue
+		}
+		output, err := page.Eval(js)
+		if err != nil {
+			return nil, err
+		}
+		outputs = append(outputs, output)
+	}
+	return outputs, nil
+}

runner/options.go

@@ -334,6 +334,9 @@ type Options struct {
 	Protocol                  string
 	OutputFilterErrorPagePath string
 	DisableStdout             bool
+
+	JavascriptCodes goflags.StringSlice
+
 	// AssetUpload
 	AssetUpload bool
 	// AssetName
@@ -404,6 +407,7 @@ func ParseOptions() *Options {
 		flagSet.BoolVar(&options.NoScreenshotFullPage, "no-screenshot-full-page", false, "disable saving full page screenshot"),
 		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.StringSliceVarP(&options.JavascriptCodes, "javascript-code", "jsc", nil, "execute JavaScript code after navigation", goflags.StringSliceOptions),
 	)
 
 	flagSet.CreateGroup("matchers", "Matchers",

runner/runner.go

@@ -39,7 +39,6 @@ import (
 	"github.com/projectdiscovery/httpx/static"
 	"github.com/projectdiscovery/mapcidr/asn"
 	"github.com/projectdiscovery/networkpolicy"
-	errorutil "github.com/projectdiscovery/utils/errors" //nolint
 	osutil "github.com/projectdiscovery/utils/os"
 	"github.com/projectdiscovery/utils/structs"
 
@@ -69,6 +68,7 @@ import (
 	"github.com/projectdiscovery/mapcidr"
 	"github.com/projectdiscovery/rawhttp"
 	converstionutil "github.com/projectdiscovery/utils/conversion"
+	errkit "github.com/projectdiscovery/utils/errkit"
 	fileutil "github.com/projectdiscovery/utils/file"
 	pdhttputil "github.com/projectdiscovery/utils/http"
 	iputil "github.com/projectdiscovery/utils/ip"
@@ -2243,7 +2243,14 @@ retry:
 	var pHash uint64
 	if scanopts.Screenshot {
 		var err error
-		screenshotBytes, headlessBody, linkRequest, err = r.browser.ScreenshotWithBody(fullURL, scanopts.ScreenshotTimeout, scanopts.ScreenshotIdle, r.options.CustomHeaders, scanopts.IsScreenshotFullPage())
+		screenshotBytes, headlessBody, linkRequest, err = r.browser.ScreenshotWithBody(
+			fullURL,
+			scanopts.ScreenshotTimeout,
+			scanopts.ScreenshotIdle,
+			r.options.CustomHeaders,
+			scanopts.IsScreenshotFullPage(),
+			r.options.JavascriptCodes,
+		)
 		if err != nil {
 			gologger.Warning().Msgf("Could not take screenshot '%s': %s", fullURL, err)
 		} else {
@@ -2507,7 +2514,7 @@ func (r *Runner) HandleFaviconHash(hp *httpx.HTTPX, req *retryablehttp.Request,
 func (r *Runner) calculateFaviconHashWithRaw(data []byte) (string, string, error) {
 	hashNum, md5Hash, err := stringz.FaviconHash(data)
 	if err != nil {
-		return "", "", errorutil.NewWithTag("favicon", "could not calculate favicon hash").Wrap(err) //nolint
+		return "", "", errkit.Wrapf(err, "could not calculate favicon hash")
 	}
 	return fmt.Sprintf("%d", hashNum), md5Hash, nil
 }