Adding support for cookie jar (#35)

* Adding some missing docs

* Adding cookie jar flag & functionality

* Migrating integration tests to new test server

* Adding tests for the dictionary generator cmd

pkg/cmd/config.go

@@ -45,5 +45,10 @@ func scanConfigFromCmd(cmd *cobra.Command) (*scan.Config, error) {
 
 	c.UserAgent = cmd.Flag(flagUserAgent).Value.String()
 
+	c.UseCookieJar, err = cmd.Flags().GetBool(flagCookieJar)
+	if err != nil {
+		return nil, errors.Wrap(err, "cookie jar flag is invalid")
+	}
+
 	return c, nil
 }

pkg/cmd/flags.go

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

pkg/cmd/root_integration_test.go

@@ -151,6 +151,32 @@ func TestDictionaryGenerateCommand(t *testing.T) {
 	assert.Contains(t, string(content), "root_integration_test.go")
 }
 
+func TestGenerateDictionaryWithoutOutputPath(t *testing.T) {
+	logger, _ := test.NewLogger()
+
+	c, err := createCommand(logger)
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	_, _, err = executeCommand(c, "dictionary.generate", ".")
+	assert.NoError(t, err)
+}
+
+func TestGenerateDictionaryWithInvalidDirectory(t *testing.T) {
+	logger, _ := test.NewLogger()
+
+	fakePath := "./" + test.RandStringRunes(10)
+	c, err := createCommand(logger)
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	_, _, err = executeCommand(c, "dictionary.generate", fakePath)
+	assert.Error(t, err)
+
+	assert.Contains(t, err.Error(), "unable to use the provided path")
+	assert.Contains(t, err.Error(), fakePath)
+}
+
 func TestVersionCommand(t *testing.T) {
 	logger, buf := test.NewLogger()
 

pkg/cmd/scan.go

@@ -75,6 +75,14 @@ func NewScanCommand(logger *logrus.Logger) (*cobra.Command, error) {
 		"user agent to use for http requests",
 	)
 
+	cmd.Flags().BoolP(
+		flagCookieJar,
+		"",
+		false,
+		"enables the use of a cookie jar: it will retain any cookie sent "+
+			"from the server and send them for the following requests",
+	)
+
 	return cmd, nil
 }
 

pkg/common/test/server.go

@@ -0,0 +1,48 @@
+package test
+
+import (
+	"sync"
+
+	"net/http"
+	"net/http/httptest"
+)
+
+func NewServerWithAssertion(handler http.HandlerFunc) (*httptest.Server, *ServerAssertion) {
+	serverAssertion := &ServerAssertion{}
+
+	server := httptest.NewServer(serverAssertion.wrap(handler))
+
+	return server, serverAssertion
+}
+
+type ServerAssertion struct {
+	requests   []http.Request
+	requestsMx sync.RWMutex
+}
+
+func (s *ServerAssertion) wrap(handler http.HandlerFunc) http.HandlerFunc {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		handler(w, r)
+
+		s.requestsMx.Lock()
+		defer s.requestsMx.Unlock()
+
+		s.requests = append(s.requests, *r)
+	})
+}
+
+func (s *ServerAssertion) Range(fn func(index int, r http.Request)) {
+	s.requestsMx.RLock()
+	defer s.requestsMx.RUnlock()
+
+	for i, r := range s.requests {
+		fn(i, r)
+	}
+}
+
+func (s *ServerAssertion) Len() int {
+	s.requestsMx.RLock()
+	defer s.requestsMx.RUnlock()
+
+	return len(s.requests)
+}

pkg/scan/bootstrap.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net"
 	"net/http"
+	"net/http/cookiejar"
 	"net/url"
 	"time"
 
@@ -14,6 +15,7 @@ import (
 	"golang.org/x/net/proxy"
 )
 
+// 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)
 	if err != nil {
@@ -64,6 +66,14 @@ func buildClientFrom(cnf *Config) (Doer, error) {
 		Timeout: time.Millisecond * time.Duration(cnf.TimeoutInMilliseconds),
 	}
 
+	if cnf.UseCookieJar {
+		jar, err := cookiejar.New(nil)
+		if err != nil {
+			return nil, err
+		}
+		c.Jar = jar
+	}
+
 	if cnf.Socks5Url != nil {
 		tbDialer, err := proxy.FromURL(cnf.Socks5Url, proxy.Direct)
 		if err != nil {

pkg/scan/bootstrap_integration_test.go

@@ -10,7 +10,7 @@ import (
 	"testing"
 	"time"
 
-	socks5 "github.com/armon/go-socks5"
+	"github.com/armon/go-socks5"
 	"github.com/chuckpreslar/emission"
 	"github.com/stefanoj3/dirstalk/pkg/common/test"
 	"github.com/stefanoj3/dirstalk/pkg/scan"
@@ -133,14 +133,8 @@ func TestShouldUseTheSpecifiedUserAgent(t *testing.T) {
 
 	logger, _ := test.NewLogger()
 
-	var request *http.Request
-	doneChannel := make(chan bool)
-
-	testServer := httptest.NewServer(
-		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			request = r
-			doneChannel <- true
-		}),
+	testServer, serverAssertion := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
 	)
 	defer testServer.Close()
 
@@ -160,19 +154,18 @@ func TestShouldUseTheSpecifiedUserAgent(t *testing.T) {
 	err = scan.StartScan(logger, eventManager, config, u)
 	assert.NoError(t, err)
 
-	select {
-	case <-doneChannel:
-		assert.Equal(t, testUserAgent, request.Header.Get("User-Agent"))
-	case <-time.After(time.Second * 1):
-		t.Fatal("failed to receive request")
-	}
+	assert.True(t, serverAssertion.Len() > 0)
+	serverAssertion.Range(func(_ int, r http.Request) {
+		assert.Equal(t, testUserAgent, r.Header.Get("User-Agent"))
+	})
 }
 
 func TestShouldFailToScanWithAnUnreachableSocks5Server(t *testing.T) {
 	logger, loggerBuffer := test.NewLogger()
 
-	requestMap := &sync.Map{}
-	testServer := buildTestServer(requestMap)
+	testServer, serverAssertion := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
+	)
 	defer testServer.Close()
 
 	u, err := url.Parse(testServer.URL)
@@ -212,9 +205,87 @@ func TestShouldFailToScanWithAnUnreachableSocks5Server(t *testing.T) {
 	assert.Len(t, actualResults, 0)
 	assert.Contains(t, loggerBuffer.String(), "connection refused")
 
-	requestMap.Range(func(key, value interface{}) bool {
-		t.Fatal("no request was supposed to be recorded: socks5 is down, the server should remain unreachable")
-		return true
+	assert.True(t, serverAssertion.Len() == 0)
+}
+
+func TestShouldRetainCookiesSetByTheServerWhenCookieJarIsEnabled(t *testing.T) {
+	const (
+		cookieName  = "my_cookies_name"
+		cookieValue = "my_cookie_value_123"
+	)
+	logger, _ := test.NewLogger()
+
+	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:    cookieName,
+						Value:   cookieValue,
+						Expires: time.Now().AddDate(0, 1, 0),
+					},
+				)
+			})
+		}),
+	)
+
+	u, err := url.Parse(testServer.URL)
+	assert.NoError(t, err)
+
+	config := &scan.Config{
+		Threads:               1,
+		DictionaryPath:        "testdata/dictionary1.txt",
+		HTTPMethods:           []string{http.MethodGet},
+		TimeoutInMilliseconds: 400,
+		ScanDepth:             2,
+		UseCookieJar:          true,
+	}
+	eventManager := emission.NewEmitter()
+
+	err = scan.StartScan(logger, eventManager, config, u)
+	assert.NoError(t, err)
+
+	assert.Equal(t, 14, serverAssertion.Len())
+	serverAssertion.Range(func(index int, r http.Request) {
+		if index == 0 { // the first request should have no cookies
+			assert.Equal(t, 0, len(r.Cookies()))
+			return
+		}
+
+		assert.Equal(t, 1, len(r.Cookies()), "Only one cookie expected, got: %v", r.Cookies())
+		assert.Equal(t, cookieName, r.Cookies()[0].Name)
+		assert.Equal(t, cookieValue, r.Cookies()[0].Value)
+	})
+}
+
+func TestShouldNotSendAnyCookieIfServerSetNoneWhenUsingCookieJar(t *testing.T) {
+	logger, _ := test.NewLogger()
+
+	testServer, serverAssertion := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
+	)
+
+	u, err := url.Parse(testServer.URL)
+	assert.NoError(t, err)
+
+	config := &scan.Config{
+		Threads:               1,
+		DictionaryPath:        "testdata/dictionary1.txt",
+		HTTPMethods:           []string{http.MethodGet},
+		TimeoutInMilliseconds: 400,
+		ScanDepth:             2,
+		UseCookieJar:          true,
+	}
+	eventManager := emission.NewEmitter()
+
+	err = scan.StartScan(logger, eventManager, config, u)
+	assert.NoError(t, err)
+
+	assert.Equal(t, 14, serverAssertion.Len())
+	serverAssertion.Range(func(index int, r http.Request) {
+		assert.Equal(t, 0, len(r.Cookies()), "No cookies expected, got: %v", r.Cookies())
 	})
 }
 

pkg/scan/config.go

@@ -2,6 +2,7 @@ package scan
 
 import "net/url"
 
+// Config represents the configuration needed to perform a scan
 type Config struct {
 	DictionaryPath        string
 	HTTPMethods           []string
@@ -10,4 +11,5 @@ type Config struct {
 	ScanDepth             int
 	Socks5Url             *url.URL
 	UserAgent             string
+	UseCookieJar          bool
 }