From 53d652d9c700a8f109c3eab950efe90b42492d6d Mon Sep 17 00:00:00 2001 From: Stefano Gabryel Date: Mon, 10 Jun 2019 09:32:13 +0200 Subject: [PATCH 1/2] Refactoring cookie header flag --- README.md | 7 +++++-- pkg/cmd/config.go | 2 +- pkg/cmd/flags.go | 3 ++- pkg/cmd/root_integration_test.go | 34 +++++++++++++++++++++++++++++--- pkg/cmd/scan.go | 6 +++--- 5 files changed, 42 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index dd7aec4..f6f5a1e 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,10 @@ 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 ``` @@ -52,7 +55,7 @@ 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 diff --git a/pkg/cmd/config.go b/pkg/cmd/config.go index 6dac3b2..95a2412 100644 --- a/pkg/cmd/config.go +++ b/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") } diff --git a/pkg/cmd/flags.go b/pkg/cmd/flags.go index 75fe354..47b42f7 100644 --- a/pkg/cmd/flags.go +++ b/pkg/cmd/flags.go @@ -16,7 +16,8 @@ const ( flagSocks5Host = "socks5" flagUserAgent = "user-agent" flagCookieJar = "use-cookie-jar" - flagCookies = "cookies" + flagCookie = "cookie" + flagHeaders = "header" // Generate dictionary flags flagOutput = "out" diff --git a/pkg/cmd/root_integration_test.go b/pkg/cmd/root_integration_test.go index 04f4f17..60157bb 100644 --- a/pkg/cmd/root_integration_test.go +++ b/pkg/cmd/root_integration_test.go @@ -137,8 +137,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", ) @@ -175,7 +177,7 @@ func TestWhenProvidingCookiesInWrongFormatShouldErr(t *testing.T) { c, "scan", testServer.URL, - "--cookies", + "--cookie", malformedCookie, "--dictionary", "testdata/dict.txt", @@ -240,6 +242,32 @@ func TestScanWithCookieJar(t *testing.T) { }) } +func TestScanWithUnknownHeaderShouldErr(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 TestDictionaryGenerateCommand(t *testing.T) { logger, _ := test.NewLogger() diff --git a/pkg/cmd/scan.go b/pkg/cmd/scan.go index db36870..eb87ad5 100644 --- a/pkg/cmd/scan.go +++ b/pkg/cmd/scan.go @@ -83,10 +83,10 @@ 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{}, - "comma separated list of cookies to add to each request; eg name=value,name2=value2", + "cookie to add to each request; eg name=value (can be specified multiple times)", ) return cmd, nil From 9636bbf435ac9b980bbccb10ce1d36f8645db29d Mon Sep 17 00:00:00 2001 From: Stefano Gabryel Date: Mon, 10 Jun 2019 10:26:18 +0200 Subject: [PATCH 2/2] Adding support for custom headers --- README.md | 5 +- pkg/cmd/config.go | 25 ++++++ pkg/cmd/flags.go | 2 +- pkg/cmd/root_integration_test.go | 83 ++++++++++++++++++- pkg/cmd/scan.go | 6 ++ pkg/scan/bootstrap.go | 22 ++++- pkg/scan/client/client.go | 9 ++ pkg/scan/client/client_test.go | 40 +++++++++ pkg/scan/client/headers.go | 31 +++++++ pkg/scan/client/headers_internal_test.go | 20 +++++ ...al_test.go => user_agent_internal_test.go} | 0 pkg/scan/config.go | 1 + 12 files changed, 236 insertions(+), 8 deletions(-) create mode 100644 pkg/scan/client/headers.go create mode 100644 pkg/scan/client/headers_internal_test.go rename pkg/scan/client/{client_internal_test.go => user_agent_internal_test.go} (100%) diff --git a/README.md b/README.md index f6f5a1e..c55637b 100644 --- a/README.md +++ b/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 diff --git a/pkg/cmd/config.go b/pkg/cmd/config.go index 95a2412..4f14338 100644 --- a/pkg/cmd/config.go +++ b/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 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)) diff --git a/pkg/cmd/flags.go b/pkg/cmd/flags.go index 47b42f7..10bfcd1 100644 --- a/pkg/cmd/flags.go +++ b/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" diff --git a/pkg/cmd/root_integration_test.go b/pkg/cmd/root_integration_test.go index 60157bb..4ffcabf 100644 --- a/pkg/cmd/root_integration_test.go +++ b/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) @@ -155,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) { @@ -242,7 +249,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 +275,76 @@ 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 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() diff --git a/pkg/cmd/scan.go b/pkg/cmd/scan.go index eb87ad5..6029797 100644 --- a/pkg/cmd/scan.go +++ b/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 } diff --git a/pkg/scan/bootstrap.go b/pkg/scan/bootstrap.go index 2f1cd37..5885843 100644 --- a/pkg/scan/bootstrap.go +++ b/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 diff --git a/pkg/scan/client/client.go b/pkg/scan/client/client.go index b6e3119..f18b6f8 100644 --- a/pkg/scan/client/client.go +++ b/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 } diff --git a/pkg/scan/client/client_test.go b/pkg/scan/client/client_test.go index ae62c3b..10c63c8 100644 --- a/pkg/scan/client/client_test.go +++ b/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) diff --git a/pkg/scan/client/headers.go b/pkg/scan/client/headers.go new file mode 100644 index 0000000..e979cf2 --- /dev/null +++ b/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) +} diff --git a/pkg/scan/client/headers_internal_test.go b/pkg/scan/client/headers_internal_test.go new file mode 100644 index 0000000..9f52705 --- /dev/null +++ b/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) +} diff --git a/pkg/scan/client/client_internal_test.go b/pkg/scan/client/user_agent_internal_test.go similarity index 100% rename from pkg/scan/client/client_internal_test.go rename to pkg/scan/client/user_agent_internal_test.go diff --git a/pkg/scan/config.go b/pkg/scan/config.go index 048101d..c189e7f 100644 --- a/pkg/scan/config.go +++ b/pkg/scan/config.go @@ -16,4 +16,5 @@ type Config struct { UserAgent string UseCookieJar bool Cookies []*http.Cookie + Headers map[string]string }