Adding possibility to specity cookies to forward

README.md

@@ -52,6 +52,9 @@ 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
+- `--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
 
 ##### 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

@@ -1,7 +1,10 @@
 package cmd
 
 import (
+	"net/http"
 	"net/url"
+	"strings"
+	"time"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
@@ -50,5 +53,37 @@ func scanConfigFromCmd(cmd *cobra.Command) (*scan.Config, error) {
 		return nil, errors.Wrap(err, "cookie jar flag is invalid")
 	}
 
+	rawCookies, err := cmd.Flags().GetStringSlice(flagCookies)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read cookies flag")
+	}
+
+	c.Cookies, err = rawCookiesToCookies(rawCookies)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to convert rawCookies to objects")
+	}
+
 	return c, nil
 }
+
+func rawCookiesToCookies(rawCookies []string) ([]*http.Cookie, error) {
+	cookies := make([]*http.Cookie, 0, len(rawCookies))
+
+	for _, rawCookie := range rawCookies {
+		parts := strings.Split(rawCookie, "=")
+		if len(parts) != 2 {
+			return nil, errors.Errorf("cookie format is invalid: %s", rawCookie)
+		}
+
+		cookies = append(
+			cookies,
+			&http.Cookie{
+				Name:    parts[0],
+				Value:   parts[1],
+				Expires: time.Now().AddDate(0, 0, 2),
+			},
+		)
+	}
+
+	return cookies, nil
+}

pkg/cmd/flags.go

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

pkg/cmd/root_integration_test.go

@@ -7,8 +7,9 @@ import (
 	"net/http/httptest"
 	"os"
 	"strings"
-	"sync/atomic"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -42,19 +43,17 @@ func TestScanCommand(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, c)
 
-	var calls int32
-	srv := httptest.NewServer(
+	testServer, serverAssertion := test.NewServerWithAssertion(
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			atomic.AddInt32(&calls, 1)
 			w.WriteHeader(http.StatusNotFound)
 		}),
 	)
-	defer srv.Close()
+	defer testServer.Close()
 
-	_, _, err = executeCommand(c, "scan", srv.URL, "--dictionary", "testdata/dict.txt")
+	_, _, err = executeCommand(c, "scan", testServer.URL, "--dictionary", "testdata/dict.txt", "-v")
 	assert.NoError(t, err)
 
-	assert.Equal(t, int32(3), calls)
+	assert.Equal(t, 3, serverAssertion.Len())
 }
 
 func TestScanWithRemoteDictionary(t *testing.T) {
@@ -76,19 +75,17 @@ blabla
 	)
 	defer dictionaryServer.Close()
 
-	var calls int32
-	srv := httptest.NewServer(
+	testServer, serverAssertion := test.NewServerWithAssertion(
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			atomic.AddInt32(&calls, 1)
 			w.WriteHeader(http.StatusNotFound)
 		}),
 	)
-	defer srv.Close()
+	defer testServer.Close()
 
-	_, _, err = executeCommand(c, "scan", srv.URL, "--dictionary", dictionaryServer.URL)
+	_, _, err = executeCommand(c, "scan", testServer.URL, "--dictionary", dictionaryServer.URL)
 	assert.NoError(t, err)
 
-	assert.Equal(t, int32(3), calls)
+	assert.Equal(t, 3, serverAssertion.Len())
 }
 
 func TestScanWithUserAgentFlag(t *testing.T) {
@@ -100,35 +97,147 @@ func TestScanWithUserAgentFlag(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, c)
 
-	var callsWithMatchingUserAgent int32
-	var callsWithNonMatchingUserAgent int32
-
-	srv := httptest.NewServer(
+	testServer, serverAssertion := test.NewServerWithAssertion(
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			if r.Header.Get("User-Agent") == testUserAgent {
-				atomic.AddInt32(&callsWithMatchingUserAgent, 1)
-			} else {
-				atomic.AddInt32(&callsWithNonMatchingUserAgent, 1)
-			}
-
 			w.WriteHeader(http.StatusNotFound)
 		}),
 	)
-	defer srv.Close()
+	defer testServer.Close()
 
 	_, _, err = executeCommand(
 		c,
 		"scan",
-		srv.URL,
+		testServer.URL,
 		"--user-agent",
 		testUserAgent,
 		"--dictionary",
 		"testdata/dict.txt",
 	)
 	assert.NoError(t, err)
 
-	assert.Equal(t, int32(3), callsWithMatchingUserAgent)
-	assert.Equal(t, int32(0), callsWithNonMatchingUserAgent)
+	assert.Equal(t, 3, serverAssertion.Len())
+	serverAssertion.Range(func(_ int, r http.Request) {
+		assert.Equal(t, testUserAgent, r.Header.Get("User-Agent"))
+	})
+}
+
+func TestScanWithCookies(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,
+		"--cookies",
+		"name1=val1,name2=val2",
+		"--dictionary",
+		"testdata/dict.txt",
+	)
+	assert.NoError(t, err)
+
+	serverAssertion.Range(func(_ int, r http.Request) {
+		assert.Equal(t, 2, len(r.Cookies()))
+
+		assert.Equal(t, r.Cookies()[0].Name, "name1")
+		assert.Equal(t, r.Cookies()[0].Value, "val1")
+
+		assert.Equal(t, r.Cookies()[1].Name, "name2")
+		assert.Equal(t, r.Cookies()[1].Value, "val2")
+	})
+}
+
+func TestWhenProvidingCookiesInWrongFormatShouldErr(t *testing.T) {
+	const malformedCookie = "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) {
+			w.WriteHeader(http.StatusNotFound)
+		}),
+	)
+	defer testServer.Close()
+
+	_, _, err = executeCommand(
+		c,
+		"scan",
+		testServer.URL,
+		"--cookies",
+		malformedCookie,
+		"--dictionary",
+		"testdata/dict.txt",
+	)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "cookie format is invalid")
+	assert.Contains(t, err.Error(), malformedCookie)
+
+	assert.Equal(t, 0, serverAssertion.Len())
+}
+
+func TestScanWithCookieJar(t *testing.T) {
+	const (
+		serverCookieName  = "server_cookie_name"
+		serverCookieValue = "server_cookie_value"
+	)
+
+	logger, _ := test.NewLogger()
+
+	c, err := createCommand(logger)
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	once := sync.Once{}
+	testServer, serverAssertion := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			once.Do(func() {
+				http.SetCookie(
+					w,
+					&http.Cookie{
+						Name:    serverCookieName,
+						Value:   serverCookieValue,
+						Expires: time.Now().AddDate(0, 1, 0),
+					},
+				)
+			})
+		}),
+	)
+	defer testServer.Close()
+
+	_, _, err = executeCommand(
+		c,
+		"scan",
+		testServer.URL,
+		"--use-cookie-jar",
+		"--dictionary",
+		"testdata/dict.txt",
+		"-t",
+		"1",
+	)
+	assert.NoError(t, err)
+
+	serverAssertion.Range(func(index int, r http.Request) {
+		if index == 0 { // first request should have no cookies
+			assert.Equal(t, 0, len(r.Cookies()))
+			return
+		}
+
+		assert.Equal(t, 1, len(r.Cookies()))
+		assert.Equal(t, r.Cookies()[0].Name, serverCookieName)
+		assert.Equal(t, r.Cookies()[0].Value, serverCookieValue)
+	})
 }
 
 func TestDictionaryGenerateCommand(t *testing.T) {

pkg/cmd/scan.go

@@ -83,6 +83,12 @@ func NewScanCommand(logger *logrus.Logger) (*cobra.Command, error) {
 			"from the server and send them for the following requests",
 	)
 
+	cmd.Flags().StringSlice(
+		flagCookies,
+		[]string{},
+		"comma separated list of cookies to add to each request; eg name=value,name2=value2",
+	)
+
 	return cmd, nil
 }
 

pkg/common/test/server.go

@@ -40,6 +40,13 @@ func (s *ServerAssertion) Range(fn func(index int, r http.Request)) {
 	}
 }
 
+func (s *ServerAssertion) At(index int, fn func(r http.Request)) {
+	s.requestsMx.RLock()
+	defer s.requestsMx.RUnlock()
+
+	fn(s.requests[index])
+}
+
 func (s *ServerAssertion) Len() int {
 	s.requestsMx.RLock()
 	defer s.requestsMx.RUnlock()

pkg/scan/bootstrap.go

@@ -1,23 +1,26 @@
 package scan
 
 import (
-	"context"
-	"net"
 	"net/http"
-	"net/http/cookiejar"
 	"net/url"
-	"time"
 
 	"github.com/chuckpreslar/emission"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/stefanoj3/dirstalk/pkg/dictionary"
-	"golang.org/x/net/proxy"
+	"github.com/stefanoj3/dirstalk/pkg/scan/client"
 )
 
 // StartScan is a convenience method that wires together all the dependencies needed to start a scan
 func StartScan(logger *logrus.Logger, eventManager *emission.Emitter, cnf *Config, u *url.URL) error {
-	c, err := buildClientFrom(cnf)
+	c, err := client.NewClientFromConfig(
+		cnf.TimeoutInMilliseconds,
+		cnf.Socks5Url,
+		cnf.UserAgent,
+		cnf.UseCookieJar,
+		cnf.Cookies,
+		u,
+	)
 	if err != nil {
 		return errors.Wrap(err, "failed to build client")
 	}
@@ -52,6 +55,9 @@ func StartScan(logger *logrus.Logger, eventManager *emission.Emitter, cnf *Confi
 		"url":               u.String(),
 		"threads":           cnf.Threads,
 		"dictionary.length": len(dict),
+		"cookies":           strigifyCookies(cnf.Cookies),
+		"user-agent":        cnf.UserAgent,
+		"socks5":            cnf.Socks5Url,
 	}).Info("Starting scan")
 
 	s.Scan(u, cnf.Threads)
@@ -61,32 +67,12 @@ func StartScan(logger *logrus.Logger, eventManager *emission.Emitter, cnf *Confi
 	return nil
 }
 
-func buildClientFrom(cnf *Config) (Doer, error) {
-	c := &http.Client{
-		Timeout: time.Millisecond * time.Duration(cnf.TimeoutInMilliseconds),
-	}
-
-	if cnf.UseCookieJar {
-		jar, err := cookiejar.New(nil)
-		if err != nil {
-			return nil, err
-		}
-		c.Jar = jar
-	}
+func strigifyCookies(cookies []*http.Cookie) string {
+	result := ""
 
-	if cnf.Socks5Url != nil {
-		tbDialer, err := proxy.FromURL(cnf.Socks5Url, proxy.Direct)
-		if err != nil {
-			return nil, err
-		}
-
-		tbTransport := &http.Transport{
-			DialContext: func(ctx context.Context, network, addr string) (conn net.Conn, e error) {
-				return tbDialer.Dial(network, addr)
-			},
-		}
-		c.Transport = tbTransport
+	for _, cookie := range cookies {
+		result += cookie.Name + "=" + cookie.Value + ";"
 	}
 
-	return newUserAgentDoerDecorator(c, cnf.UserAgent), nil
+	return result
 }