From 046c2b0e9012def8f02b20fcb3cac185d0fed9d8 Mon Sep 17 00:00:00 2001 From: Stefano Gabryel Date: Mon, 10 Jun 2019 10:26:18 +0200 Subject: [PATCH] Adding support for custom headers --- README.md | 5 +- pkg/cmd/config.go | 25 ++++++++++ pkg/cmd/flags.go | 2 +- pkg/cmd/root_integration_test.go | 46 ++++++++++++++++++- 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, 200 insertions(+), 7 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..6c0b857 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 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)) 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..563e75b 100644 --- a/pkg/cmd/root_integration_test.go +++ b/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() 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 }