Adding support for custom headers

README.md

@@ -44,7 +44,9 @@ dirstalk scan http://someaddress.url/ \
 --socks5 127.0.0.1:9150 \
 --cookie name=value \
 --use-cookie-jar \
---user-agent my_user_agent
+--user-agent my_user_agent \
+--header "Authorization: Bearer 123"
+
 ```
 
 
@@ -58,6 +60,7 @@ dirstalk scan http://someaddress.url/ \
 - `--cookie` cookie to add to each request; eg name=value (can be specified multiple times)
 - `--use-cookie-jar` enables the use of a cookie jar: it will retain any cookie sent from the server and send them for the following requests
 - `--user-agent` user agent to use for http requests
+- `--header` header to add to each request; eg name=value (can be specified multiple times)
 
 ##### Useful resources
 - [here](https://github.com/dustyfresh/dictionaries/tree/master/DirBuster-Lists) you can find dictionaries that can be used with dirstalk

pkg/cmd/config.go

@@ -63,9 +63,34 @@ func scanConfigFromCmd(cmd *cobra.Command) (*scan.Config, error) {
 		return nil, errors.Wrap(err, "failed to convert rawCookies to objects")
 	}
 
+	rawHeaders, err := cmd.Flags().GetStringArray(flagHeader)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read cookies flag")
+	}
+
+	c.Headers, err = rawHeadersToHeaders(rawHeaders)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to convert rawCookies to objects")
+	}
+
 	return c, nil
 }
 
+func rawHeadersToHeaders(rawHeaders []string) (map[string]string, error) {
+	headers := make(map[string]string, len(rawHeaders)*2)
+
+	for _, rawHeader := range rawHeaders {
+		parts := strings.Split(rawHeader, ":")
+		if len(parts) != 2 {
+			return nil, errors.Errorf("header is in invalid format")
+		}
+
+		headers[parts[0]] = parts[1]
+	}
+
+	return headers, nil
+}
+
 func rawCookiesToCookies(rawCookies []string) ([]*http.Cookie, error) {
 	cookies := make([]*http.Cookie, 0, len(rawCookies))
 

pkg/cmd/flags.go

@@ -17,7 +17,7 @@ const (
 	flagUserAgent       = "user-agent"
 	flagCookieJar       = "use-cookie-jar"
 	flagCookie          = "cookie"
-	flagHeaders         = "header"
+	flagHeader          = "header"
 
 	// Generate dictionary flags
 	flagOutput           = "out"

pkg/cmd/root_integration_test.go

@@ -122,7 +122,7 @@ func TestScanWithUserAgentFlag(t *testing.T) {
 }
 
 func TestScanWithCookies(t *testing.T) {
-	logger, _ := test.NewLogger()
+	logger, loggerBuffer := test.NewLogger()
 
 	c, err := createCommand(logger)
 	assert.NoError(t, err)
@@ -155,6 +155,10 @@ func TestScanWithCookies(t *testing.T) {
 		assert.Equal(t, r.Cookies()[1].Name, "name2")
 		assert.Equal(t, r.Cookies()[1].Value, "val2")
 	})
+
+	// to ensure we print the headers to the cli
+	assert.Contains(t, loggerBuffer.String(), "name1=val1")
+	assert.Contains(t, loggerBuffer.String(), "name2=val2")
 }
 
 func TestWhenProvidingCookiesInWrongFormatShouldErr(t *testing.T) {
@@ -242,7 +246,7 @@ func TestScanWithCookieJar(t *testing.T) {
 	})
 }
 
-func TestScanWithUnknownHeaderShouldErr(t *testing.T) {
+func TestScanWithUnknownFlagShouldErr(t *testing.T) {
 	logger, _ := test.NewLogger()
 
 	c, err := createCommand(logger)
@@ -268,6 +272,44 @@ func TestScanWithUnknownHeaderShouldErr(t *testing.T) {
 	assert.Equal(t, 0, serverAssertion.Len())
 }
 
+func TestScanWithHeaders(t *testing.T) {
+	logger, loggerBuffer := test.NewLogger()
+
+	c, err := createCommand(logger)
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	testServer, serverAssertion := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
+	)
+	defer testServer.Close()
+
+	_, _, err = executeCommand(
+		c,
+		"scan",
+		testServer.URL,
+		"--header",
+		"Accept-Language: en-US,en;q=0.5",
+		"--header",
+		`"Authorization: Bearer 123"`,
+		"--dictionary",
+		"testdata/dict.txt",
+	)
+	assert.NoError(t, err)
+
+	serverAssertion.Range(func(_ int, r http.Request) {
+		assert.Equal(t, 2, len(r.Header))
+
+		assert.Equal(t, "en-US,en;q=0.5", r.Header.Get("Accept-Language"))
+		assert.Equal(t, "Bearer 123", r.Header.Get("Authorization"))
+	})
+
+	// to ensure we print the headers to the cli
+	assert.Contains(t, loggerBuffer.String(), "Accept-Language")
+	assert.Contains(t, loggerBuffer.String(), "Authorization")
+	assert.Contains(t, loggerBuffer.String(), "Bearer 123")
+}
+
 func TestDictionaryGenerateCommand(t *testing.T) {
 	logger, _ := test.NewLogger()
 

pkg/cmd/scan.go

@@ -89,6 +89,12 @@ func NewScanCommand(logger *logrus.Logger) (*cobra.Command, error) {
 		"cookie to add to each request; eg name=value (can be specified multiple times)",
 	)
 
+	cmd.Flags().StringArray(
+		flagHeader,
+		[]string{},
+		"header to add to each request; eg name=value (can be specified multiple times)",
+	)
+
 	return cmd, nil
 }
 

pkg/scan/bootstrap.go

@@ -1,6 +1,7 @@
 package scan
 
 import (
+	"fmt"
 	"net/http"
 	"net/url"
 
@@ -19,6 +20,7 @@ func StartScan(logger *logrus.Logger, eventManager *emission.Emitter, cnf *Confi
 		cnf.UserAgent,
 		cnf.UseCookieJar,
 		cnf.Cookies,
+		cnf.Headers,
 		u,
 	)
 	if err != nil {
@@ -54,10 +56,14 @@ func StartScan(logger *logrus.Logger, eventManager *emission.Emitter, cnf *Confi
 	logger.WithFields(logrus.Fields{
 		"url":               u.String(),
 		"threads":           cnf.Threads,
-		"dictionary.length": len(dict),
+		"dictionary-length": len(dict),
+		"scan-depth":        cnf.ScanDepth,
+		"timeout":           cnf.TimeoutInMilliseconds,
+		"socks5":            cnf.Socks5Url,
 		"cookies":           strigifyCookies(cnf.Cookies),
+		"cookie-jar":        cnf.UseCookieJar,
+		"headers":           stringyfyHeaders(cnf.Headers),
 		"user-agent":        cnf.UserAgent,
-		"socks5":            cnf.Socks5Url,
 	}).Info("Starting scan")
 
 	s.Scan(u, cnf.Threads)
@@ -71,7 +77,17 @@ func strigifyCookies(cookies []*http.Cookie) string {
 	result := ""
 
 	for _, cookie := range cookies {
-		result += cookie.Name + "=" + cookie.Value + ";"
+		result += fmt.Sprintf("{%s=%s}", cookie.Name, cookie.Value)
+	}
+
+	return result
+}
+
+func stringyfyHeaders(headers map[string]string) string {
+	result := ""
+
+	for name, value := range headers {
+		result += fmt.Sprintf("{%s:%s}", name, value)
 	}
 
 	return result

pkg/scan/client/client.go

@@ -19,6 +19,7 @@ func NewClientFromConfig(
 	userAgent string,
 	useCookieJar bool,
 	cookies []*http.Cookie,
+	headers map[string]string,
 	u *url.URL,
 ) (*http.Client, error) {
 	transport := http.Transport{
@@ -61,10 +62,18 @@ func NewClientFromConfig(
 	}
 
 	var err error
+
 	c.Transport, err = decorateTransportWithUserAgentDecorator(c.Transport, userAgent)
 	if err != nil {
 		return nil, errors.Wrap(err, "NewClientFromConfig: failed to decorate transport")
 	}
 
+	if len(headers) > 0 {
+		c.Transport, err = decorateTransportWithHeadersDecorator(c.Transport, headers)
+		if err != nil {
+			return nil, errors.Wrap(err, "NewClientFromConfig: failed to decorate transport")
+		}
+	}
+
 	return c, nil
 }

pkg/scan/client/client_test.go

@@ -26,6 +26,7 @@ func TestWhenRemoteIsTooSlowClientShouldTimeout(t *testing.T) {
 		false,
 		nil,
 		nil,
+		nil,
 	)
 	assert.NoError(t, err)
 
@@ -72,6 +73,7 @@ func TestShouldForwardProvidedCookiesWhenUsingJar(t *testing.T) {
 		"",
 		true,
 		cookies,
+		map[string]string{},
 		u,
 	)
 	assert.NoError(t, err)
@@ -130,6 +132,7 @@ func TestShouldForwardCookiesWhenJarIsDisabled(t *testing.T) {
 		"",
 		false,
 		cookies,
+		map[string]string{},
 		u,
 	)
 	assert.NoError(t, err)
@@ -149,6 +152,42 @@ func TestShouldForwardCookiesWhenJarIsDisabled(t *testing.T) {
 	})
 }
 
+func TestShouldForwardProvidedHeader(t *testing.T) {
+	const (
+		headerName  = "my_header_name"
+		headerValue = "my_header_value_123"
+	)
+	testServer, serverAssertion := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
+	)
+	defer testServer.Close()
+
+	u, err := url.Parse(testServer.URL)
+	assert.NoError(t, err)
+
+	c, err := client.NewClientFromConfig(
+		100,
+		nil,
+		"",
+		false,
+		nil,
+		map[string]string{headerName: headerValue},
+		u,
+	)
+	assert.NoError(t, err)
+
+	res, err := c.Get(testServer.URL)
+	assert.NoError(t, err)
+	assert.NotNil(t, res)
+
+	assert.Equal(t, 1, serverAssertion.Len())
+
+	serverAssertion.At(0, func(r http.Request) {
+		assert.Equal(t, headerValue, r.Header.Get(headerName))
+
+	})
+}
+
 func TestShouldFailToCreateAClientWithInvalidSocks5Url(t *testing.T) {
 	u := url.URL{Scheme: "potatoscheme"}
 
@@ -158,6 +197,7 @@ func TestShouldFailToCreateAClientWithInvalidSocks5Url(t *testing.T) {
 		"",
 		false,
 		nil,
+		map[string]string{},
 		nil,
 	)
 	assert.Nil(t, c)

pkg/scan/client/headers.go

@@ -0,0 +1,31 @@
+package client
+
+import (
+	"errors"
+	"net/http"
+)
+
+func decorateTransportWithHeadersDecorator(decorated http.RoundTripper, headers map[string]string) (*headersTransportDecorator, error) {
+	if decorated == nil {
+		return nil, errors.New("decorated round tripper is nil")
+	}
+
+	if headers == nil {
+		return nil, errors.New("headers is nil")
+	}
+
+	return &headersTransportDecorator{decorated: decorated, headers: headers}, nil
+}
+
+type headersTransportDecorator struct {
+	decorated http.RoundTripper
+	headers   map[string]string
+}
+
+func (h *headersTransportDecorator) RoundTrip(r *http.Request) (*http.Response, error) {
+	for key, value := range h.headers {
+		r.Header.Set(key, value)
+	}
+
+	return h.decorated.RoundTrip(r)
+}

pkg/scan/client/headers_internal_test.go

@@ -0,0 +1,20 @@
+package client
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDecorateTransportHeaderShouldFailWithNilDecorated(t *testing.T) {
+	transport, err := decorateTransportWithHeadersDecorator(nil, map[string]string{})
+	assert.Nil(t, transport)
+	assert.Error(t, err)
+}
+
+func TestDecorateTransportHeaderShouldFailWithNilHeaderMap(t *testing.T) {
+	transport, err := decorateTransportWithHeadersDecorator(http.DefaultTransport, nil)
+	assert.Nil(t, transport)
+	assert.Error(t, err)
+}

pkg/scan/config.go

@@ -16,4 +16,5 @@ type Config struct {
 	UserAgent             string
 	UseCookieJar          bool
 	Cookies               []*http.Cookie
+	Headers               map[string]string
 }