Adding headers support

README.md

@@ -41,7 +41,12 @@ dirstalk scan http://someaddress.url/ \
 --http-timeout 10000 \
 --scan-depth 10 \
 --threads 10 \
---socks5 127.0.0.1:9150
+--socks5 127.0.0.1:9150 \
+--cookie name=value \
+--use-cookie-jar \
+--user-agent my_user_agent \
+--header "Authorization: Bearer 123"
+
 ```
 
 
@@ -52,9 +57,10 @@ dirstalk scan http://someaddress.url/ \
 - `--scan-depth` the maximum recursion depth
 - `--threads` the number of threads performing concurrent requests
 - `--socks5` SOCKS5 server to connect to (all the requests including the one to fetch the remote dictionary will go through it)
-- `--cookies` comma separated list of cookies to add to each request; eg name=value,name2=value2
+- `--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

@@ -53,7 +53,7 @@ func scanConfigFromCmd(cmd *cobra.Command) (*scan.Config, error) {
 		return nil, errors.Wrap(err, "cookie jar flag is invalid")
 	}
 
-	rawCookies, err := cmd.Flags().GetStringSlice(flagCookies)
+	rawCookies, err := cmd.Flags().GetStringArray(flagCookie)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to read cookies flag")
 	}
@@ -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 rawHeaders")
+	}
+
 	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: %s", rawHeader)
+		}
+
+		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

@@ -16,7 +16,8 @@ const (
 	flagSocks5Host      = "socks5"
 	flagUserAgent       = "user-agent"
 	flagCookieJar       = "use-cookie-jar"
-	flagCookies         = "cookies"
+	flagCookie          = "cookie"
+	flagHeader          = "header"
 
 	// Generate dictionary flags
 	flagOutput           = "out"

pkg/cmd/root_integration_test.go

@@ -91,7 +91,7 @@ blabla
 func TestScanWithUserAgentFlag(t *testing.T) {
 	const testUserAgent = "my_test_user_agent"
 
-	logger, _ := test.NewLogger()
+	logger, loggerBuffer := test.NewLogger()
 
 	c, err := createCommand(logger)
 	assert.NoError(t, err)
@@ -119,10 +119,13 @@ func TestScanWithUserAgentFlag(t *testing.T) {
 	serverAssertion.Range(func(_ int, r http.Request) {
 		assert.Equal(t, testUserAgent, r.Header.Get("User-Agent"))
 	})
+
+	// to ensure we print the user agent to the cli
+	assert.Contains(t, loggerBuffer.String(), testUserAgent)
 }
 
 func TestScanWithCookies(t *testing.T) {
-	logger, _ := test.NewLogger()
+	logger, loggerBuffer := test.NewLogger()
 
 	c, err := createCommand(logger)
 	assert.NoError(t, err)
@@ -137,8 +140,10 @@ func TestScanWithCookies(t *testing.T) {
 		c,
 		"scan",
 		testServer.URL,
-		"--cookies",
-		"name1=val1,name2=val2",
+		"--cookie",
+		"name1=val1",
+		"--cookie",
+		"name2=val2",
 		"--dictionary",
 		"testdata/dict.txt",
 	)
@@ -153,6 +158,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 cookies to the cli
+	assert.Contains(t, loggerBuffer.String(), "name1=val1")
+	assert.Contains(t, loggerBuffer.String(), "name2=val2")
 }
 
 func TestWhenProvidingCookiesInWrongFormatShouldErr(t *testing.T) {
@@ -175,7 +184,7 @@ func TestWhenProvidingCookiesInWrongFormatShouldErr(t *testing.T) {
 		c,
 		"scan",
 		testServer.URL,
-		"--cookies",
+		"--cookie",
 		malformedCookie,
 		"--dictionary",
 		"testdata/dict.txt",
@@ -240,6 +249,102 @@ func TestScanWithCookieJar(t *testing.T) {
 	})
 }
 
+func TestScanWithUnknownFlagShouldErr(t *testing.T) {
+	logger, _ := 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,
+		"--gibberishflag",
+		"--dictionary",
+		"testdata/dict.txt",
+	)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "unknown flag")
+
+	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 TestScanWithMalformedHeaderShouldErr(t *testing.T) {
+	const malformedHeader = "gibberish"
+
+	logger, _ := 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",
+		malformedHeader,
+		"--dictionary",
+		"testdata/dict.txt",
+	)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), malformedHeader)
+	assert.Contains(t, err.Error(), "header is in invalid format")
+
+	assert.Equal(t, 0, serverAssertion.Len())
+}
+
 func TestDictionaryGenerateCommand(t *testing.T) {
 	logger, _ := test.NewLogger()
 

pkg/cmd/scan.go

@@ -83,10 +83,16 @@ func NewScanCommand(logger *logrus.Logger) (*cobra.Command, error) {
 			"from the server and send them for the following requests",
 	)
 
-	cmd.Flags().StringSlice(
-		flagCookies,
+	cmd.Flags().StringArray(
+		flagCookie,
+		[]string{},
+		"cookie to add to each request; eg name=value (can be specified multiple times)",
+	)
+
+	cmd.Flags().StringArray(
+		flagHeader,
 		[]string{},
-		"comma separated list of cookies to add to each request; eg name=value,name2=value2",
+		"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
 }