Adding flag to specify where to save scan output & optimizing scan logic

functional-tests.sh

@@ -64,7 +64,7 @@ assert_contains "$SCAN_RESULT" "error" "an error is expected when no argument is
 SCAN_RESULT=$(./dist/dirstalk scan -d resources/tests/dictionary.txt http://localhost:8080 2>&1 || true);
 assert_contains "$SCAN_RESULT" "/index" "result expected when performing scan"
 assert_contains "$SCAN_RESULT" "/index/home" "result expected when performing scan"
-assert_contains "$SCAN_RESULT" "8 requests made, 3 results found" "a recap was expected when performing a scan"
+assert_contains "$SCAN_RESULT" "3 results found" "a recap was expected when performing a scan"
 assert_contains "$SCAN_RESULT" "├── home" "a recap was expected when performing a scan"
 assert_contains "$SCAN_RESULT" "└── index" "a recap was expected when performing a scan"
 assert_contains "$SCAN_RESULT" "    └── home" "a recap was expected when performing a scan"

pkg/cmd/config.go

@@ -83,6 +83,8 @@ func scanConfigFromCmd(cmd *cobra.Command) (*scan.Config, error) {
 		return nil, errors.Wrap(err, "failed to convert rawHeaders")
 	}
 
+	c.Out = cmd.Flag(flagOutput).Value.String()
+
 	return c, nil
 }
 

pkg/cmd/flags.go

@@ -20,6 +20,7 @@ const (
 	flagCookieJar            = "use-cookie-jar"
 	flagCookie               = "cookie"
 	flagHeader               = "header"
+	flagResultOutput         = "out"
 
 	// Generate dictionary flags
 	flagOutput           = "out"

pkg/cmd/output.go

@@ -0,0 +1,8 @@
+package cmd
+
+import "github.com/stefanoj3/dirstalk/pkg/scan"
+
+type OutputSaver interface {
+	Save(scan.Result) error
+	Close() error
+}

pkg/cmd/scan.go

@@ -13,6 +13,8 @@ import (
 	"github.com/stefanoj3/dirstalk/pkg/dictionary"
 	"github.com/stefanoj3/dirstalk/pkg/scan"
 	"github.com/stefanoj3/dirstalk/pkg/scan/client"
+	"github.com/stefanoj3/dirstalk/pkg/scan/filter"
+	"github.com/stefanoj3/dirstalk/pkg/scan/output"
 	"github.com/stefanoj3/dirstalk/pkg/scan/producer"
 	"github.com/stefanoj3/dirstalk/pkg/scan/summarizer"
 )
@@ -115,6 +117,12 @@ func NewScanCommand(logger *logrus.Logger) (*cobra.Command, error) {
 		"header to add to each request; eg name=value (can be specified multiple times)",
 	)
 
+	cmd.Flags().String(
+		flagResultOutput,
+		"",
+		"path where to store result output",
+	)
+
 	return cmd, nil
 }
 
@@ -173,12 +181,15 @@ func startScan(logger *logrus.Logger, cnf *scan.Config, u *url.URL) error {
 	}
 
 	targetProducer := producer.NewDictionaryProducer(cnf.HTTPMethods, dict, cnf.ScanDepth)
-	reproducer := producer.NewReProducer(targetProducer, cnf.HTTPStatusesToIgnore)
+	reproducer := producer.NewReProducer(targetProducer)
+
+	resultFilter := filter.NewHTTPStatusResultFilter(cnf.HTTPStatusesToIgnore)
 
 	s := scan.NewScanner(
 		c,
 		targetProducer,
 		reproducer,
+		resultFilter,
 		logger,
 	)
 
@@ -189,19 +200,28 @@ func startScan(logger *logrus.Logger, cnf *scan.Config, u *url.URL) error {
 		"scan-depth":        cnf.ScanDepth,
 		"timeout":           cnf.TimeoutInMilliseconds,
 		"socks5":            cnf.Socks5Url,
-		"cookies":           strigifyCookies(cnf.Cookies),
+		"cookies":           stringifyCookies(cnf.Cookies),
 		"cookie-jar":        cnf.UseCookieJar,
-		"headers":           stringyfyHeaders(cnf.Headers),
+		"headers":           stringifyHeaders(cnf.Headers),
 		"user-agent":        cnf.UserAgent,
 	}).Info("Starting scan")
 
-	resultSummarizer := summarizer.NewResultSummarizer(cnf.HTTPStatusesToIgnore, logger)
+	resultSummarizer := summarizer.NewResultSummarizer(logger)
 
 	osSigint := make(chan os.Signal, 1)
 	signal.Notify(osSigint, os.Interrupt)
 
+	outputSaver, err := newOutputSaver(cnf.Out)
+	if err != nil {
+		return errors.Wrap(err, "failed to create output saver")
+	}
+
 	finishFunc := func() {
 		resultSummarizer.Summarize()
+		err := outputSaver.Close()
+		if err != nil {
+			logger.WithError(err).Error("failed to close output file")
+		}
 		logger.Info("Finished scan")
 	}
 
@@ -213,6 +233,10 @@ func startScan(logger *logrus.Logger, cnf *scan.Config, u *url.URL) error {
 			return nil
 		default:
 			resultSummarizer.Add(result)
+			err := outputSaver.Save(result)
+			if err != nil {
+				return errors.Wrap(err, "failed to add output to file")
+			}
 		}
 	}
 
@@ -221,7 +245,15 @@ func startScan(logger *logrus.Logger, cnf *scan.Config, u *url.URL) error {
 	return nil
 }
 
-func strigifyCookies(cookies []*http.Cookie) string {
+func newOutputSaver(path string) (OutputSaver, error) {
+	if path == "" {
+		return output.NewNullSaver(), nil
+	}
+
+	return output.NewFileSaver(path)
+}
+
+func stringifyCookies(cookies []*http.Cookie) string {
 	result := ""
 
 	for _, cookie := range cookies {
@@ -231,7 +263,7 @@ func strigifyCookies(cookies []*http.Cookie) string {
 	return result
 }
 
-func stringyfyHeaders(headers map[string]string) string {
+func stringifyHeaders(headers map[string]string) string {
 	result := ""
 
 	for name, value := range headers {

pkg/cmd/scan_integration_test.go

@@ -1,9 +1,11 @@
 package cmd_test
 
 import (
+	"io/ioutil"
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"sync"
 	"syscall"
 	"testing"
@@ -102,6 +104,93 @@ func TestScanCommand(t *testing.T) {
 	assert.Contains(t, loggerBuffer.String(), expectedResultTree)
 }
 
+func TestScanShouldWriteOutput(t *testing.T) {
+	logger, _ := test.NewLogger()
+
+	c, err := createCommand(logger)
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	testServer, _ := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/home" {
+				w.WriteHeader(http.StatusOK)
+				return
+			}
+
+			w.WriteHeader(http.StatusNotFound)
+		}),
+	)
+	defer testServer.Close()
+
+	outputFilename := test.RandStringRunes(10)
+	outputFilename = "testdata/out/" + outputFilename + ".txt"
+
+	defer func() {
+		err := os.Remove(outputFilename)
+		if err != nil {
+			panic("failed to remove file create during test: " + err.Error())
+		}
+	}()
+
+	_, _, err = executeCommand(
+		c,
+		"scan",
+		testServer.URL,
+		"--dictionary",
+		"testdata/dict2.txt",
+		"--out",
+		outputFilename,
+	)
+	assert.NoError(t, err)
+
+	file, err := os.Open(outputFilename)
+	assert.NoError(t, err)
+
+	b, err := ioutil.ReadAll(file)
+	assert.NoError(t, err, "failed to read file content")
+
+	expected := `{"Target":{"Path":"home","Method":"GET","Depth":3},"StatusCode":200,"URL":{"Scheme":"http","Opaque":"","User":null,"Host":"` +
+		testServer.Listener.Addr().String() +
+		`","Path":"/home","RawPath":"","ForceQuery":false,"RawQuery":"","Fragment":""}}
+`
+	assert.Equal(t, expected, string(b))
+
+	assert.NoError(t, file.Close(), "failed to close file")
+}
+
+func TestScanInvalidOutputFileShouldErr(t *testing.T) {
+	logger, _ := test.NewLogger()
+
+	c, err := createCommand(logger)
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	testServer, _ := test.NewServerWithAssertion(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/home" {
+				w.WriteHeader(http.StatusOK)
+				return
+			}
+
+			w.WriteHeader(http.StatusNotFound)
+		}),
+	)
+	defer testServer.Close()
+
+	_, _, err = executeCommand(
+		c,
+		"scan",
+		testServer.URL,
+		"--dictionary",
+		"testdata/dict2.txt",
+		"--out",
+		"/root/blabla/123/gibberish/123",
+	)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to create output saver")
+}
+
 func TestScanWithInvalidStatusesToIgnoreShouldErr(t *testing.T) {
 	logger, _ := test.NewLogger()
 

pkg/cmd/testdata/.gitignore

@@ -1,2 +1,2 @@
+out/*
 !out/.gitkeep
-out/*

pkg/common/test/rand.go

@@ -2,11 +2,14 @@ package test
 
 import (
 	"math/rand"
+	"time"
 )
 
 var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890")
 
 func RandStringRunes(n int) string {
+	rand.Seed(time.Now().UnixNano())
+
 	b := make([]rune, n)
 	for i := range b {
 		b[i] = letterRunes[rand.Intn(len(letterRunes))]

pkg/common/urlpath/join_test.go

@@ -49,7 +49,8 @@ func TestJoin(t *testing.T) {
 	}
 
 	for _, tc := range testCases {
-		tc := tc
+		tc := tc // Pinning ranged variable, more info: https://github.com/kyoh86/scopelint
+
 		scenario := fmt.Sprintf(
 			"Input: `%s`, Expected output: `%s`",
 			strings.Join(tc.input, ","),

pkg/scan/config.go

@@ -19,4 +19,5 @@ type Config struct {
 	UseCookieJar          bool
 	Cookies               []*http.Cookie
 	Headers               map[string]string
+	Out                   string
 }

pkg/scan/filter.go

@@ -0,0 +1,5 @@
+package scan
+
+type ResultFilter interface {
+	ShouldIgnore(Result) bool
+}

pkg/scan/filter/http.go

@@ -0,0 +1,23 @@
+package filter
+
+import (
+	"github.com/stefanoj3/dirstalk/pkg/scan"
+)
+
+func NewHTTPStatusResultFilter(httpStatusesToIgnore []int) HTTPStatusResultFilter {
+	httpStatusesToIgnoreMap := make(map[int]struct{}, len(httpStatusesToIgnore))
+	for _, statusToIgnore := range httpStatusesToIgnore {
+		httpStatusesToIgnoreMap[statusToIgnore] = struct{}{}
+	}
+
+	return HTTPStatusResultFilter{httpStatusesToIgnoreMap: httpStatusesToIgnoreMap}
+}
+
+type HTTPStatusResultFilter struct {
+	httpStatusesToIgnoreMap map[int]struct{}
+}
+
+func (f HTTPStatusResultFilter) ShouldIgnore(result scan.Result) bool {
+	_, found := f.httpStatusesToIgnoreMap[result.StatusCode]
+	return found
+}

pkg/scan/filter/http_test.go

@@ -0,0 +1,70 @@
+package filter_test
+
+import (
+	"fmt"
+	"net/http"
+	"sync"
+	"testing"
+
+	"github.com/stefanoj3/dirstalk/pkg/scan"
+	"github.com/stefanoj3/dirstalk/pkg/scan/filter"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHTTPStatusResultFilter(t *testing.T) {
+	testCases := []struct {
+		statusCodesToIgnore []int
+		result              scan.Result
+		expectedResult      bool
+	}{
+		{
+			statusCodesToIgnore: []int{http.StatusCreated, http.StatusNotFound},
+			result:              scan.Result{StatusCode: http.StatusOK},
+			expectedResult:      false,
+		},
+		{
+			statusCodesToIgnore: []int{http.StatusCreated, http.StatusNotFound},
+			result:              scan.Result{StatusCode: http.StatusNotFound},
+			expectedResult:      true,
+		},
+		{
+			statusCodesToIgnore: []int{},
+			result:              scan.Result{StatusCode: http.StatusNotFound},
+			expectedResult:      false,
+		},
+		{
+			statusCodesToIgnore: []int{},
+			result:              scan.Result{StatusCode: http.StatusOK},
+			expectedResult:      false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc // Pinning ranged variable, more info: https://github.com/kyoh86/scopelint
+
+		scenario := fmt.Sprintf("ignored: %v, result: %d", tc.statusCodesToIgnore, tc.result.StatusCode)
+
+		t.Run(scenario, func(t *testing.T) {
+			t.Parallel()
+
+			actual := filter.NewHTTPStatusResultFilter(tc.statusCodesToIgnore).ShouldIgnore(tc.result)
+			assert.Equal(t, tc.expectedResult, actual)
+		})
+	}
+}
+
+func TestHTTPStatusResultFilterShouldWorkConcurrently(t *testing.T) {
+	sut := filter.NewHTTPStatusResultFilter(nil)
+
+	wg := sync.WaitGroup{}
+	for i := 0; i < 1000; i++ {
+		wg.Add(1)
+
+		go func(i int) {
+			sut.ShouldIgnore(scan.Result{StatusCode: i})
+			wg.Done()
+		}(i)
+	}
+
+	wg.Wait()
+}

pkg/scan/output/conversion.go

@@ -0,0 +1,11 @@
+package output
+
+import (
+	"encoding/json"
+
+	"github.com/stefanoj3/dirstalk/pkg/scan"
+)
+
+func convertResultToRawData(r scan.Result) ([]byte, error) {
+	return json.Marshal(r)
+}