add assume-404-regex flag

cmd/testserver/main.go

@@ -1,17 +1,39 @@
 package main
 
 import (
+	"fmt"
+	"log"
 	"net/http"
 )
 
 func main() {
 	handlerFunc := func(writer http.ResponseWriter, request *http.Request) {}
+	handlerWith404BodyFunc := func(writer http.ResponseWriter, request *http.Request) {
+		pageName := request.URL.Path
+		// respond with 200 intentionally
+		writer.WriteHeader(http.StatusOK)
 
+		_, err := writer.Write([]byte(fmt.Sprintf("404: page %s was not found", pageName)))
+		if err != nil {
+			log.Printf("Error writing response: %v", err)
+		}
+	}
+	handlerWith403BodyFunc := func(writer http.ResponseWriter, request *http.Request) {
+		// respond with 200 intentionally
+		writer.WriteHeader(http.StatusOK)
+		_, err := writer.Write([]byte("403: forbidden"))
+		if err != nil {
+			log.Printf("Error writing response: %v", err)
+		}
+	}
+
+	http.HandleFunc("/noHome", handlerWith404BodyFunc)
+	http.HandleFunc("/forbiddenHome", handlerWith403BodyFunc)
 	http.HandleFunc("/home", handlerFunc)
 	http.HandleFunc("/index", handlerFunc)
 	http.HandleFunc("/index/home", handlerFunc)
 
-	if err := http.ListenAndServe(":8080", nil); err != nil {
+	if err := http.ListenAndServe("127.0.0.1:7999", nil); err != nil {
 		panic(err)
 	}
 }

functional-tests.sh

@@ -40,7 +40,7 @@ function assert_not_contains {
 
 ## Starting test server running on the 8080 port
 echo "Starting test server"
-./dist/testserver&
+./dist/testserver &
 SERVER_PID=$!
 sleep 1
 echo "Done"
@@ -66,14 +66,33 @@ assert_contains "$VERSION_RESULT" "Built" "the build time is expected to be prin
 SCAN_RESULT=$(./dist/dirstalk scan 2>&1 || true);
 assert_contains "$SCAN_RESULT" "error" "an error is expected when no argument is passed"
 
-SCAN_RESULT=$(./dist/dirstalk scan -d resources/tests/dictionary.txt http://localhost:8080 2>&1);
+SCAN_RESULT=$(./dist/dirstalk scan -d resources/tests/dictionary.txt http://localhost:7999 --assume-status-regex=404="404: page .* was not " --assume-status-regex=403="403: forbidden" 2>&1);
+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" "4 results found" "a recap was expected when performing a scan"
+assert_contains "$SCAN_RESULT" "├── forbiddenHome" "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"
+
+SCAN_RESULT=$(./dist/dirstalk scan -d resources/tests/dictionary.txt http://localhost:7999 --http-statuses-to-ignore=403,404 --assume-status-regex=404="404: page .* was not " --assume-status-regex=403="403: forbidden" 2>&1);
 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" "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"
 
+SCAN_RESULT=$(./dist/dirstalk scan -d resources/tests/dictionary.txt http://localhost:7999 2>&1);
+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" "5 results found" "a recap was expected when performing a scan"
+assert_contains "$SCAN_RESULT" "├── forbiddenHome" "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"
+assert_contains "$SCAN_RESULT" "└── noHome" "a recap was expected when performing a scan"
+
 assert_not_contains "$SCAN_RESULT" "error" "no error is expected for a successful scan"
 
 SCAN_RESULT=$(./dist/dirstalk scan -h 2>&1);

pkg/cmd/config.go

@@ -1,8 +1,11 @@
 package cmd
 
 import (
+	"fmt"
 	"net/http"
 	"net/url"
+	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -75,6 +78,33 @@ func scanConfigFromCmd(cmd *cobra.Command) (*scan.Config, error) {
 		return nil, errors.Wrapf(err, failedToReadPropertyError, flagScanHeader)
 	}
 
+	statusStrings, err := cmd.Flags().GetStringArray(flagAssumeStatusRegex)
+	if err != nil {
+		return nil, errors.Wrapf(err, failedToReadPropertyError, flagAssumeStatusRegex)
+	}
+	if len(statusStrings) > 0 {
+		c.AssumeStatusCodeRegex = make(map[int]regexp.Regexp)
+		for _, statusString := range statusStrings {
+			index := strings.Index(statusString, "=")
+			if index == -1 {
+				return nil, errors.New(
+					fmt.Sprintf("Failed to parse option '%s': '=' was not found in string.\n"+
+						"Usage: --%s=HTTP_CODE=REGEX", statusString, flagAssumeStatusRegex))
+			}
+			httpCodeString := statusString[:index]
+			httpCode, err := strconv.Atoi(httpCodeString)
+			if err != nil {
+				return nil, errors.Wrapf(err, failedToReadPropertyError, flagAssumeStatusRegex)
+			}
+			regexString := statusString[index+1:]
+			regex, err := regexp.Compile(regexString)
+			if err != nil {
+				return nil, errors.Wrapf(err, failedToReadPropertyError, flagAssumeStatusRegex)
+			}
+			c.AssumeStatusCodeRegex[httpCode] = *regex
+		}
+	}
+
 	if c.Headers, err = rawHeadersToHeaders(rawHeaders); err != nil {
 		return nil, errors.Wrapf(err, "failed to convert rawHeaders (%v)", rawHeaders)
 	}

pkg/cmd/flags.go

@@ -23,6 +23,7 @@ const (
 	flagScanHeader                          = "header"
 	flagScanResultOutput                    = "out"
 	flagShouldSkipSSLCertificatesValidation = "no-check-certificate"
+	flagAssumeStatusRegex                   = "assume-status-regex"
 
 	flagIgnore20xWithEmptyBody = "ignore-empty-body"
 

pkg/cmd/scan.go

@@ -7,6 +7,7 @@ import (
 	"net/url"
 	"os"
 	"os/signal"
+	"regexp"
 
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -139,6 +140,14 @@ func NewScanCommand(logger *logrus.Logger) *cobra.Command {
 		"ignore HTTP 20x responses with empty body",
 	)
 
+	cmd.Flags().StringArray(
+		flagAssumeStatusRegex,
+		[]string{},
+		fmt.Sprintf("Assume 404 response code if body matches the regex. "+
+			"Useful when server replies with 200 on Not found page.\n"+
+			"Usage: --%s=HTTP_CODE=REGEX", flagAssumeStatusRegex),
+	)
+
 	return cmd
 }
 
@@ -198,6 +207,7 @@ func startScan(logger *logrus.Logger, cnf *scan.Config, u *url.URL) error {
 		"cookie-jar":        cnf.UseCookieJar,
 		"headers":           stringifyHeaders(cnf.Headers),
 		"user-agent":        cnf.UserAgent,
+		"assumeStatusCodes": stringifyAssumptions(cnf.AssumeStatusCodeRegex),
 	}).Info("Starting scan")
 
 	resultSummarizer := summarizer.NewResultSummarizer(tree.NewResultTreeProducer(), logger)
@@ -263,7 +273,10 @@ func buildScanner(cnf *scan.Config, dict []string, u *url.URL, logger *logrus.Lo
 	targetProducer := producer.NewDictionaryProducer(cnf.HTTPMethods, dict, cnf.ScanDepth)
 	reproducer := producer.NewReProducer(targetProducer)
 
-	resultFilter := filter.NewHTTPStatusResultFilter(cnf.HTTPStatusesToIgnore, cnf.IgnoreEmpty20xResponses)
+	resultFilter, err := filter.NewHTTPStatusResultFilter(cnf.HTTPStatusesToIgnore, cnf.IgnoreEmpty20xResponses, cnf.AssumeStatusCodeRegex)
+	if err != nil {
+		return nil, err
+	}
 
 	scannerClient, err := buildScannerClient(cnf, u)
 	if err != nil {
@@ -360,3 +373,13 @@ func stringifyHeaders(headers map[string]string) string {
 
 	return result
 }
+
+func stringifyAssumptions(assumption map[int]regexp.Regexp) string {
+	result := ""
+
+	for code, value := range assumption {
+		result += fmt.Sprintf("{%d:%s}", code, value.String())
+	}
+
+	return result
+}

pkg/scan/config.go

@@ -3,6 +3,7 @@ package scan
 import (
 	"net/http"
 	"net/url"
+	"regexp"
 )
 
 // Config represents the configuration needed to perform a scan.
@@ -23,4 +24,5 @@ type Config struct {
 	Out                                 string
 	ShouldSkipSSLCertificatesValidation bool
 	IgnoreEmpty20xResponses             bool
+	AssumeStatusCodeRegex               map[int]regexp.Regexp
 }

pkg/scan/filter.go

@@ -2,4 +2,5 @@ package scan
 
 type ResultFilter interface {
 	ShouldIgnore(Result) bool
+	ShouldReadBody() bool
 }

pkg/scan/filter/http.go

@@ -2,28 +2,42 @@ package filter
 
 import (
 	"github.com/stefanoj3/dirstalk/pkg/scan"
+	"regexp"
 )
 
-func NewHTTPStatusResultFilter(httpStatusesToIgnore []int, ignoreEmptyBody bool) HTTPStatusResultFilter {
+func NewHTTPStatusResultFilter(httpStatusesToIgnore []int, ignoreEmptyBody bool, assumeStatusRegex map[int]regexp.Regexp) (*HTTPStatusResultFilter, error) {
 	httpStatusesToIgnoreMap := make(map[int]struct{}, len(httpStatusesToIgnore))
 	for _, statusToIgnore := range httpStatusesToIgnore {
 		httpStatusesToIgnoreMap[statusToIgnore] = struct{}{}
 	}
 
-	return HTTPStatusResultFilter{httpStatusesToIgnoreMap: httpStatusesToIgnoreMap, ignoreEmptyBody: ignoreEmptyBody}
+	return &HTTPStatusResultFilter{httpStatusesToIgnoreMap: httpStatusesToIgnoreMap, ignoreEmptyBody: ignoreEmptyBody, assumeStatusRegex: assumeStatusRegex}, nil
 }
 
 type HTTPStatusResultFilter struct {
 	httpStatusesToIgnoreMap map[int]struct{}
 	ignoreEmptyBody         bool
+	assumeStatusRegex       map[int]regexp.Regexp
 }
 
 func (f HTTPStatusResultFilter) ShouldIgnore(result scan.Result) bool {
 	if f.ignoreEmptyBody && result.StatusCode/100 == 2 && result.ContentLength == 0 {
 		return true
 	}
 
+	if f.assumeStatusRegex != nil {
+		for code, regex := range f.assumeStatusRegex {
+			if regex.Match(result.Body) {
+				_, found := f.httpStatusesToIgnoreMap[code]
+				return found
+			}
+		}
+	}
 	_, found := f.httpStatusesToIgnoreMap[result.StatusCode]
 
 	return found
 }
+
+func (f HTTPStatusResultFilter) ShouldReadBody() bool {
+	return len(f.assumeStatusRegex) > 0
+}

pkg/scan/filter/http_test.go

@@ -47,14 +47,22 @@ func TestHTTPStatusResultFilter(t *testing.T) {
 		t.Run(scenario, func(t *testing.T) {
 			t.Parallel()
 
-			actual := filter.NewHTTPStatusResultFilter(tc.statusCodesToIgnore, false).ShouldIgnore(tc.result)
+			resultFilter, err := filter.NewHTTPStatusResultFilter(tc.statusCodesToIgnore, false, nil)
+			if err != nil {
+				panic(err)
+			}
+			actual := resultFilter.ShouldIgnore(tc.result)
 			assert.Equal(t, tc.expectedResult, actual)
 		})
 	}
 }
 
 func TestHTTPStatusResultFilterShouldWorkConcurrently(_ *testing.T) {
-	sut := filter.NewHTTPStatusResultFilter(nil, false)
+	sut, err := filter.NewHTTPStatusResultFilter(nil, false, nil)
+
+	if err != nil {
+		panic(err)
+	}
 
 	wg := sync.WaitGroup{}
 

pkg/scan/output/saver_integration_test.go

@@ -55,7 +55,7 @@ func TestFileSaverShouldWriteResults(t *testing.T) {
 
 	assert.NoError(t, file.Close())
 
-	expected := `{"Target":{"Path":"","Method":"","Depth":0},"StatusCode":0,"URL":{"Scheme":"","Opaque":"","User":null,"Host":"","Path":"","RawPath":"","ForceQuery":false,"RawQuery":"","Fragment":"","RawFragment":""},"ContentLength":0}
+	expected := `{"Target":{"Path":"","Method":"","Depth":0},"StatusCode":0,"URL":{"Scheme":"","Opaque":"","User":null,"Host":"","Path":"","RawPath":"","OmitHost":false,"ForceQuery":false,"RawQuery":"","Fragment":"","RawFragment":""},"ContentLength":0,"Body":null}
 `
 	assert.Equal(
 		t,

pkg/scan/producer/reproducer_test.go

@@ -35,8 +35,7 @@ func TestNewReProducer(t *testing.T) {
 			Request: &http.Request{
 				URL: test.MustParseURL(t, "http://mysite/contacts"),
 			},
-		},
-	)
+		}, nil)
 
 	reproducerFunc := sut.Reproduce(context.Background())
 	reproducerChannel := reproducerFunc(result)
@@ -108,8 +107,7 @@ func TestReProducerShouldProduceNothingForDepthZero(t *testing.T) {
 			Request: &http.Request{
 				URL: test.MustParseURL(t, "http://mysite/contacts"),
 			},
-		},
-	)
+		}, nil)
 
 	reproducerFunc := sut.Reproduce(context.Background())
 	reproducerChannel := reproducerFunc(result)
@@ -141,8 +139,7 @@ func BenchmarkReProducer(b *testing.B) {
 			Request: &http.Request{
 				URL: test.MustParseURL(b, "http://mysite/contacts"),
 			},
-		},
-	)
+		}, nil)
 
 	b.ResetTimer()
 

pkg/scan/scanner.go

@@ -2,6 +2,8 @@ package scan
 
 import (
 	"context"
+	"github.com/pkg/errors"
+	"io"
 	"net/http"
 	"net/url"
 	"strings"
@@ -25,15 +27,17 @@ type Result struct {
 	StatusCode    int
 	URL           url.URL
 	ContentLength int64
+	Body          []byte
 }
 
 // NewResult creates a new instance of the Result entity based on the Target and Response.
-func NewResult(target Target, response *http.Response) Result {
+func NewResult(target Target, response *http.Response, body []byte) Result {
 	return Result{
 		Target:        target,
 		StatusCode:    response.StatusCode,
 		URL:           *response.Request.URL,
 		ContentLength: response.ContentLength,
+		Body:          body,
 	}
 }
 
@@ -151,11 +155,25 @@ func (s *Scanner) processRequest(
 		return
 	}
 
-	if err := res.Body.Close(); err != nil {
+	body := res.Body
+	bodyBytes := make([]byte, 0, 0)
+	if s.resultFilter.ShouldReadBody() && res.ContentLength > 0 {
+		bodyBytes, err = io.ReadAll(body)
+		if err != nil && err != io.EOF {
+			l.WithError(err).Warn("failed to read response body")
+			return
+		}
+
+		if int64(len(bodyBytes)) != res.ContentLength {
+			l.WithError(errors.New("actual body length does not match Content-Length header")).Warn("possible body read mismatch")
+		}
+	}
+
+	if err := body.Close(); err != nil {
 		l.WithError(err).Warn("failed to close response body")
 	}
 
-	result := NewResult(target, res)
+	result := NewResult(target, res, bodyBytes)
 
 	if s.resultFilter.ShouldIgnore(result) {
 		return