Add Headers support

Author: Chocapikk

cmd/scan.go

@@ -35,6 +35,8 @@ var scanCmd = &cobra.Command{
 		outputFile := cmd.Flag("output").Value.String()
 		outputFormat := utils.DetectOutputFormat(outputFile)
 
+		headers, _ := cmd.Flags().GetStringArray("header")
+
 		opts := scanner.ScanOptions{
 			URL:            cmd.Flag("url").Value.String(),
 			File:           cmd.Flag("file").Value.String(),
@@ -45,6 +47,7 @@ var scanCmd = &cobra.Command{
 			Verbose:        mustBool(cmd.Flags().GetBool("verbose")),
 			ScanMode:       cmd.Flag("mode").Value.String(),
 			PluginList:     cmd.Flag("plugin-list").Value.String(),
+			Headers:        headers,
 		}
 
 		if opts.URL == "" && opts.File == "" {
@@ -66,6 +69,8 @@ func init() {
 	scanCmd.Flags().StringP("mode", "m", "stealthy", "Scan mode: stealthy, bruteforce, or hybrid")
 	scanCmd.Flags().
 		StringP("plugin-list", "p", "", "Path to a custom plugin list file for bruteforce mode")
+	scanCmd.Flags().
+		StringArrayP("header", "H", []string{}, "HTTP header to include in requests. Can be specified multiple times.")
 }
 
 func mustBool(value bool, err error) bool {

internal/scanner/bruteforce.go

@@ -55,6 +55,7 @@ func BruteforcePlugins(
 	plugins []string,
 	threads int,
 	progress *utils.ProgressManager,
+	headers []string,
 ) []string {
 	if len(plugins) == 0 {
 		utils.DefaultLogger.Warning("No plugins provided for brute-force scan")
@@ -88,7 +89,7 @@ func BruteforcePlugins(
 				progress.SetMessage(fmt.Sprintf("🔎 Bruteforcing plugin %-30.30s", p))
 			}
 
-			version := utils.GetPluginVersion(normalized, p, threads)
+			version := utils.GetPluginVersion(normalized, p, threads, headers)
 			if version != "" && version != "unknown" {
 				if progress != nil {
 					progress.ClearLine()
@@ -120,9 +121,10 @@ func HybridScan(
 	bruteforcePlugins []string,
 	threads int,
 	progress *utils.ProgressManager,
+	headers []string,
 ) []string {
 	if len(stealthyPlugins) == 0 {
-		return BruteforcePlugins(target, bruteforcePlugins, threads, progress)
+		return BruteforcePlugins(target, bruteforcePlugins, threads, progress, headers)
 	}
 
 	detectedMap := make(map[string]bool, len(stealthyPlugins))
@@ -137,7 +139,7 @@ func HybridScan(
 		}
 	}
 
-	brutefound := BruteforcePlugins(target, remaining, threads, progress)
+	brutefound := BruteforcePlugins(target, remaining, threads, progress, headers)
 	result := make([]string, len(stealthyPlugins), len(stealthyPlugins)+len(brutefound))
 	copy(result, stealthyPlugins)
 	return append(result, brutefound...)

internal/scanner/endpoints.go

@@ -51,8 +51,8 @@ func fetchEndpointsFromPath(target, path string, httpClient *utils.HTTPClientMan
 	return endpoints
 }
 
-func FetchEndpoints(target string) []string {
-	httpClient := utils.NewHTTPClient(10 * time.Second)
+func FetchEndpoints(target string, headers []string) []string {
+	httpClient := utils.NewHTTPClient(10*time.Second, headers)
 
 	endpointsChan := make(chan []string, 2)
 	var wg sync.WaitGroup

internal/scanner/endpoints_test.go

@@ -31,53 +31,71 @@ import (
 func TestFetchEndpoints(t *testing.T) {
 	tests := []struct {
 		name       string
-		target     string
 		mockServer func(w http.ResponseWriter, r *http.Request)
+		headers    []string
 		want       []string
 	}{
 		{
-			name: "Valid response with routes",
+			name: "Valid response with routes and header present",
 			mockServer: func(w http.ResponseWriter, r *http.Request) {
+				if r.Header.Get("X-Test") != "value" {
+					w.WriteHeader(http.StatusForbidden)
+					return
+				}
 				response := map[string]interface{}{
 					"routes": map[string]interface{}{
 						"/wp/v2/posts":      nil,
 						"/wp/v2/comments":   nil,
 						"/wp/v2/categories": nil,
 					},
 				}
-				if err := json.NewEncoder(w).Encode(response); err != nil {
-					t.Errorf("Failed to encode JSON response: %v", err)
+				_ = json.NewEncoder(w).Encode(response)
+			},
+			headers: []string{"X-Test: value"},
+			want:    []string{"/wp/v2/posts", "/wp/v2/comments", "/wp/v2/categories"},
+		},
+		{
+			name: "Valid response with routes but header missing",
+			mockServer: func(w http.ResponseWriter, r *http.Request) {
+				if r.Header.Get("X-Test") == "" {
+					response := map[string]interface{}{
+						"routes": map[string]interface{}{
+							"/wp/v2/posts": nil,
+						},
+					}
+					_ = json.NewEncoder(w).Encode(response)
+					return
 				}
+				w.WriteHeader(http.StatusInternalServerError)
 			},
-			want: []string{"/wp/v2/posts", "/wp/v2/comments", "/wp/v2/categories"},
+			headers: nil,
+			want:    []string{"/wp/v2/posts"},
 		},
 		{
 			name: "Response without routes",
 			mockServer: func(w http.ResponseWriter, r *http.Request) {
-				response := map[string]interface{}{
+				_ = json.NewEncoder(w).Encode(map[string]interface{}{
 					"data": "No routes here",
-				}
-				if err := json.NewEncoder(w).Encode(response); err != nil {
-					t.Errorf("Failed to encode JSON response: %v", err)
-				}
+				})
 			},
-			want: []string{},
+			headers: nil,
+			want:    []string{},
 		},
 		{
 			name: "Invalid JSON response",
 			mockServer: func(w http.ResponseWriter, r *http.Request) {
-				if _, err := w.Write([]byte("{invalid-json")); err != nil {
-					t.Errorf("Failed to write invalid JSON: %v", err)
-				}
+				_, _ = w.Write([]byte("{invalid-json"))
 			},
-			want: []string{},
+			headers: nil,
+			want:    []string{},
 		},
 		{
 			name: "HTTP error response",
 			mockServer: func(w http.ResponseWriter, r *http.Request) {
 				w.WriteHeader(http.StatusInternalServerError)
 			},
-			want: []string{},
+			headers: nil,
+			want:    []string{},
 		},
 	}
 
@@ -86,7 +104,7 @@ func TestFetchEndpoints(t *testing.T) {
 			server := httptest.NewServer(http.HandlerFunc(tt.mockServer))
 			defer server.Close()
 
-			got := FetchEndpoints(server.URL)
+			got := FetchEndpoints(server.URL, tt.headers)
 
 			sort.Strings(got)
 			sort.Strings(tt.want)

internal/scanner/scan.go

@@ -37,6 +37,7 @@ type ScanOptions struct {
 	Verbose        bool
 	ScanMode       string
 	PluginList     string
+	Headers        []string
 }
 
 func ScanTargets(opts ScanOptions) {
@@ -114,7 +115,7 @@ func performStealthyScan(
 		return nil, PluginDetectionResult{}
 	}
 
-	endpoints := FetchEndpoints(target)
+	endpoints := FetchEndpoints(target, opts.Headers)
 	if len(endpoints) == 0 {
 		if opts.File == "" {
 			utils.DefaultLogger.Warning("No REST endpoints found on " + target)
@@ -153,7 +154,7 @@ func performBruteforceScan(
 		}
 	}
 
-	detected := BruteforcePlugins(target, plugins, threads, pb)
+	detected := BruteforcePlugins(target, plugins, threads, pb, opts.Headers)
 
 	pr := PluginDetectionResult{
 		Plugins:  make(map[string]*PluginData, len(detected)),
@@ -209,7 +210,7 @@ func performHybridScan(
 		bruteBar = utils.NewProgressBar(len(remaining), "🔎 Bruteforcing remaining")
 	}
 
-	brutefound := BruteforcePlugins(target, remaining, threads, bruteBar)
+	brutefound := BruteforcePlugins(target, remaining, threads, bruteBar, opts.Headers)
 
 	if bruteBar != nil {
 		bruteBar.Finish()
@@ -306,7 +307,7 @@ func ScanSite(
 
 			version := "unknown"
 			if !opts.NoCheckVersion {
-				version = utils.GetPluginVersion(target, pl, opts.Threads)
+				version = utils.GetPluginVersion(target, pl, opts.Threads, opts.Headers)
 			}
 
 			var matched []wordfence.Vulnerability

internal/utils/http.go

@@ -41,9 +41,10 @@ const maxRedirects = 10
 type HTTPClientManager struct {
 	client    *http.Client
 	userAgent string
+	headers   []string
 }
 
-func NewHTTPClient(timeout time.Duration) *HTTPClientManager {
+func NewHTTPClient(timeout time.Duration, headers []string) *HTTPClientManager {
 	transport := &http.Transport{
 		TLSClientConfig:   &tls.Config{InsecureSkipVerify: true},
 		DisableKeepAlives: true,
@@ -63,6 +64,7 @@ func NewHTTPClient(timeout time.Duration) *HTTPClientManager {
 	return &HTTPClientManager{
 		client:    client,
 		userAgent: uarand.GetRandom(),
+		headers:   headers,
 	}
 }
 
@@ -72,7 +74,35 @@ func (h *HTTPClientManager) Get(url string) (string, error) {
 		return "", errors.New("failed to create request: " + err.Error())
 	}
 
-	req.Header.Set("User-Agent", h.userAgent)
+	hasUA := false
+	for _, hdr := range h.headers {
+		parts := strings.SplitN(hdr, ":", 2)
+		if len(parts) != 2 {
+			continue
+		}
+		key := strings.TrimSpace(parts[0])
+		if strings.EqualFold(key, "User-Agent") {
+			req.Header.Add("User-Agent", strings.TrimSpace(parts[1]))
+			hasUA = true
+		}
+	}
+
+	if !hasUA {
+		req.Header.Set("User-Agent", h.userAgent)
+	}
+
+	for _, hdr := range h.headers {
+		parts := strings.SplitN(hdr, ":", 2)
+		if len(parts) != 2 {
+			continue
+		}
+		key := strings.TrimSpace(parts[0])
+		if strings.EqualFold(key, "User-Agent") {
+			continue
+		}
+		value := strings.TrimSpace(parts[1])
+		req.Header.Add(key, value)
+	}
 
 	resp, err := h.client.Do(req)
 	if err != nil {
@@ -117,20 +147,14 @@ func (h *HTTPClientManager) Get(url string) (string, error) {
 	return string(data), nil
 }
 
-// NormalizeURL ensures the URL has the correct format (removes trailing slash)
 func NormalizeURL(url string) string {
-	// Remove trailing slash if present
 	url = strings.TrimSuffix(url, "/")
-
-	// Ensure URL has http:// or https:// prefix
 	if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
 		url = "https://" + url
 	}
-
 	return url
 }
 
-// SplitLines splits a byte array into lines
 func SplitLines(data []byte) []string {
 	var lines []string
 	scanner := bufio.NewScanner(bytes.NewReader(data))

internal/utils/http_test.go

@@ -79,7 +79,7 @@ func TestHTTPClientManager_Get(t *testing.T) {
 			mockServer := httptest.NewServer(tt.serverFunc)
 			defer mockServer.Close()
 
-			client := NewHTTPClient(5 * time.Second)
+			client := NewHTTPClient(5*time.Second, nil)
 
 			got, err := client.Get(mockServer.URL)
 

internal/utils/logger_test.go

@@ -17,6 +17,25 @@
 // IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 // CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+// Copyright (c) 2025 Valentin Lobstein (Chocapikk) <balgogan@protonmail.com>
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of
+// this software and associated documentation files (the "Software"), to deal in
+// the Software without restriction, including without limitation the rights to
+// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+// the Software, and to permit persons to whom the Software is furnished to do so,
+// subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
 package utils
 
 import (
@@ -45,7 +64,7 @@ func TestLogger_Info(t *testing.T) {
 	var buf bytes.Buffer
 	originalLogger := DefaultLogger.Logger
 	DefaultLogger.Logger = log.New(&buf, "", 0)
-	defer func() { DefaultLogger.Logger = originalLogger }() // Restaure l'ancien Logger
+	defer func() { DefaultLogger.Logger = originalLogger }()
 
 	msg := "This is an info message"
 	DefaultLogger.Info(msg)
@@ -105,14 +124,13 @@ func TestLogger_PrintBanner(t *testing.T) {
 	version := "v1.0.0"
 	isLatest := true
 	DefaultLogger.PrintBanner(version, isLatest)
-	defer func() { _ = w.Close() }()
+	_ = w.Close()
 
 	var outBuf bytes.Buffer
 	_, _ = outBuf.ReadFrom(r)
 	os.Stdout = originalStdout
 
 	output := outBuf.String()
-
 	if !strings.Contains(output, version) || !strings.Contains(output, "latest") {
 		t.Errorf("PrintBanner() output = %v, want version %v and 'latest'", output, version)
 	}
@@ -121,14 +139,13 @@ func TestLogger_PrintBanner(t *testing.T) {
 	os.Stdout = w
 
 	DefaultLogger.PrintBanner(version, false)
+	_ = w.Close()
 
-	defer func() { _ = w.Close() }()
 	outBuf.Reset()
 	_, _ = outBuf.ReadFrom(r)
 	os.Stdout = originalStdout
 
 	output = outBuf.String()
-
 	if !strings.Contains(output, "outdated") {
 		t.Errorf("PrintBanner() output = %v, want 'outdated'", output)
 	}

internal/utils/version.go

@@ -70,8 +70,8 @@ func CheckLatestVersion(currentVersion string) (string, bool) {
 	return latest.String(), curr.Compare(latest) >= 0
 }
 
-func GetPluginVersion(target, plugin string, _ int) string {
-	httpClient := NewHTTPClient(10 * time.Second)
+func GetPluginVersion(target, plugin string, _ int, headers []string) string {
+	httpClient := NewHTTPClient(10*time.Second, headers)
 	return fetchVersionFromReadme(httpClient, target, plugin)
 }