fix: add user agent handler

projectdiscovery · Apr 27, 2024 · 650b0e4 · 650b0e4
1 parent 9ba3bb8
commit 650b0e4
Showing 1 changed file with 20 additions and 5 deletions.
diff --git a/pkg/engine/hybrid/crawl.go b/pkg/engine/hybrid/crawl.go
@@ -237,14 +237,29 @@ func (c *Crawler) addHeadersToPage(page *rod.Page) {
 	if len(c.Headers) == 0 {
 		return
 	}
+
 	var arr []string
+	var userAgentSet bool
+
 	for k, v := range c.Headers {
-		arr = append(arr, k, v)
+		if k == "User-Agent" {
+			userAgentParams := &proto.NetworkSetUserAgentOverride{
+				UserAgent: v,
+			}
+			if err := page.SetUserAgent(userAgentParams); err != nil {
+				gologger.Error().Msgf("headless: could not set user agent: %v", err)
+			}
+			userAgentSet = true
+		} else {
+			arr = append(arr, k, v)
+		}
 	}
-	// ignore cleanup callback
-	_, err := page.SetExtraHeaders(arr)
-	if err != nil {
-		gologger.Error().Msgf("headless: could not set extra headers: %v", err)
+
+	if !userAgentSet || len(arr) > 0 {
+		_, err := page.SetExtraHeaders(arr)
+		if err != nil {
+			gologger.Error().Msgf("headless: could not set extra headers: %v", err)
+		}
 	}
 }