Merge pull request #86 from projectdiscovery/feature-global-headers

custom header via cli

internal/runner/options.go

@@ -5,23 +5,25 @@ import (
 	"os"
 
 	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/nuclei/pkg/requests"
 )
 
 // Options contains the configuration options for tuning
 // the template requesting process.
 type Options struct {
-	Templates     string // Signature specifies the template/templates to use
-	Targets       string // Targets specifies the targets to scan using templates.
-	Threads       int    // Thread controls the number of concurrent requests to make.
-	Timeout       int    // Timeout is the seconds to wait for a response from the server.
-	Retries       int    // Retries is the number of times to retry the request
-	Output        string // Output is the file to write found subdomains to.
-	ProxyURL      string // ProxyURL is the URL for the proxy server
-	ProxySocksURL string // ProxySocksURL is the URL for the proxy socks server
-	Silent        bool   // Silent suppresses any extra text and only writes found URLs on screen.
-	Version       bool   // Version specifies if we should just show version and exit
-	Verbose       bool   // Verbose flag indicates whether to show verbose output or not
-	NoColor       bool   // No-Color disables the colored output.
+	Templates     string                 // Signature specifies the template/templates to use
+	Targets       string                 // Targets specifies the targets to scan using templates.
+	Threads       int                    // Thread controls the number of concurrent requests to make.
+	Timeout       int                    // Timeout is the seconds to wait for a response from the server.
+	Retries       int                    // Retries is the number of times to retry the request
+	Output        string                 // Output is the file to write found subdomains to.
+	ProxyURL      string                 // ProxyURL is the URL for the proxy server
+	ProxySocksURL string                 // ProxySocksURL is the URL for the proxy socks server
+	Silent        bool                   // Silent suppresses any extra text and only writes found URLs on screen.
+	Version       bool                   // Version specifies if we should just show version and exit
+	Verbose       bool                   // Verbose flag indicates whether to show verbose output or not
+	NoColor       bool                   // No-Color disables the colored output.
+	CustomHeaders requests.CustomHeaders // Custom global headers
 
 	Stdin bool // Stdin specifies whether stdin input was given to the process
 }
@@ -42,6 +44,7 @@ func ParseOptions() *Options {
 	flag.IntVar(&options.Threads, "c", 10, "Number of concurrent requests to make")
 	flag.IntVar(&options.Timeout, "timeout", 5, "Time to wait in seconds before timeout")
 	flag.IntVar(&options.Retries, "retries", 1, "Number of times to retry a failed request")
+	flag.Var(&options.CustomHeaders, "H", "Custom Header.")
 
 	flag.Parse()
 

internal/runner/runner.go

@@ -202,6 +202,7 @@ func (r *Runner) processTemplateWithList(template *templates.Template, request i
 			Retries:       r.options.Retries,
 			ProxyURL:      r.options.ProxyURL,
 			ProxySocksURL: r.options.ProxySocksURL,
+			CustomHeaders: r.options.CustomHeaders,
 		})
 	}
 	if err != nil {

pkg/executor/executer_http.go

@@ -8,6 +8,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -23,11 +24,12 @@ import (
 // HTTPExecutor is client for performing HTTP requests
 // for a template.
 type HTTPExecutor struct {
-	httpClient  *retryablehttp.Client
-	template    *templates.Template
-	httpRequest *requests.HTTPRequest
-	writer      *bufio.Writer
-	outputMutex *sync.Mutex
+	httpClient    *retryablehttp.Client
+	template      *templates.Template
+	httpRequest   *requests.HTTPRequest
+	writer        *bufio.Writer
+	outputMutex   *sync.Mutex
+	customHeaders requests.CustomHeaders
 }
 
 // HTTPOptions contains configuration options for the HTTP executor.
@@ -39,6 +41,7 @@ type HTTPOptions struct {
 	Retries       int
 	ProxyURL      string
 	ProxySocksURL string
+	CustomHeaders requests.CustomHeaders
 }
 
 // NewHTTPExecutor creates a new HTTP executor from a template
@@ -59,11 +62,12 @@ func NewHTTPExecutor(options *HTTPOptions) (*HTTPExecutor, error) {
 	client.CheckRetry = retryablehttp.HostSprayRetryPolicy()
 
 	executer := &HTTPExecutor{
-		httpClient:  client,
-		template:    options.Template,
-		httpRequest: options.HTTPRequest,
-		outputMutex: &sync.Mutex{},
-		writer:      options.Writer,
+		httpClient:    client,
+		template:      options.Template,
+		httpRequest:   options.HTTPRequest,
+		outputMutex:   &sync.Mutex{},
+		writer:        options.Writer,
+		customHeaders: options.CustomHeaders,
 	}
 	return executer, nil
 }
@@ -82,6 +86,7 @@ mainLoop:
 		if compiledRequest.Error != nil {
 			return errors.Wrap(err, "could not make http request")
 		}
+		e.setCustomHeaders(compiledRequest)
 		req := compiledRequest.Request
 		resp, err := e.httpClient.Do(req)
 		if err != nil {
@@ -223,3 +228,19 @@ func makeCheckRedirectFunc(followRedirects bool, maxRedirects int) checkRedirect
 		return nil
 	}
 }
+
+func (e *HTTPExecutor) setCustomHeaders(r *requests.CompiledHTTP) {
+	for _, customHeader := range e.customHeaders {
+		// This should be pre-computed somewhere and done only once
+		tokens := strings.Split(customHeader, ":")
+		// if it's an invalid header skip it
+		if len(tokens) < 2 {
+			continue
+		}
+
+		headerName, headerValue := tokens[0], strings.Join(tokens[1:], "")
+		headerName = strings.TrimSpace(headerName)
+		headerValue = strings.TrimSpace(headerValue)
+		r.Request.Header.Set(headerName, headerValue)
+	}
+}

pkg/requests/http-request.go

@@ -303,3 +303,17 @@ type CompiledHTTP struct {
 	Error   error
 	Meta    map[string]interface{}
 }
+
+// CustomHeaders valid for all requests
+type CustomHeaders []string
+
+// String returns just a label
+func (c *CustomHeaders) String() string {
+	return "Custom Global Headers"
+}
+
+// Set a new global header
+func (c *CustomHeaders) Set(value string) error {
+	*c = append(*c, value)
+	return nil
+}