Added configurable redirects per request

internal/runner/runner.go

@@ -16,14 +16,13 @@ import (
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/pkg/extractors"
 	"github.com/projectdiscovery/nuclei/pkg/matchers"
+	"github.com/projectdiscovery/nuclei/pkg/requests"
 	"github.com/projectdiscovery/nuclei/pkg/templates"
 	retryablehttp "github.com/projectdiscovery/retryablehttp-go"
 )
 
 // Runner is a client for running the enumeration process.
 type Runner struct {
-	// client is the http client with retries
-	client *retryablehttp.Client
 	// output is the output file to write if any
 	output      *os.File
 	outputMutex *sync.Mutex
@@ -38,29 +37,6 @@ func New(options *Options) (*Runner, error) {
 		options:     options,
 	}
 
-	retryablehttpOptions := retryablehttp.DefaultOptionsSpraying
-	retryablehttpOptions.RetryWaitMax = 10 * time.Second
-	retryablehttpOptions.RetryMax = options.Retries
-
-	// Create the HTTP Client
-	client := retryablehttp.NewWithHTTPClient(&http.Client{
-		Transport: &http.Transport{
-			MaxIdleConnsPerHost: -1,
-			TLSClientConfig: &tls.Config{
-				Renegotiation:      tls.RenegotiateOnceAsClient,
-				InsecureSkipVerify: true,
-			},
-			DisableKeepAlives: true,
-		},
-		Timeout: time.Duration(options.Timeout) * time.Second,
-		CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
-			return http.ErrUseLastResponse
-		},
-	}, retryablehttpOptions)
-	client.CheckRetry = retryablehttp.HostSprayRetryPolicy()
-
-	runner.client = client
-
 	// Create the output file if asked
 	if options.Output != "" {
 		output, err := os.Create(options.Output)
@@ -86,7 +62,14 @@ func (r *Runner) RunEnumeration() {
 
 	// If the template path is a single template and not a glob, use that.
 	if !strings.Contains(r.options.Templates, "*") {
-		r.processTemplate(r.options.Templates)
+		template, err := templates.ParseTemplate(r.options.Templates)
+		if err != nil {
+			gologger.Errorf("Could not parse template file '%s': %s\n", r.options.Templates, err)
+			return
+		}
+		for _, request := range template.Requests {
+			r.processTemplateRequest(template, request)
+		}
 		return
 	}
 
@@ -97,37 +80,38 @@ func (r *Runner) RunEnumeration() {
 	}
 
 	for _, match := range matches {
-		r.processTemplate(match)
+		template, err := templates.ParseTemplate(match)
+		if err != nil {
+			gologger.Errorf("Could not parse template file '%s': %s\n", match, err)
+			return
+		}
+		for _, request := range template.Requests {
+			r.processTemplateRequest(template, request)
+		}
 	}
 }
 
 // processTemplate processes a template and runs the enumeration on all the targets
-func (r *Runner) processTemplate(templatePath string) {
+func (r *Runner) processTemplateRequest(template *templates.Template, request *requests.Request) {
-	template, err := templates.ParseTemplate(templatePath)
-	if err != nil {
-		gologger.Errorf("Could not parse template file '%s': %s\n", templatePath, err)
-		return
-	}
-
 	// Handle a list of hosts as argument
 	if r.options.Targets != "" {
 		file, err := os.Open(r.options.Targets)
 		if err != nil {
 			gologger.Fatalf("Could not open targets file '%s': %s\n", r.options.Targets, err)
 		}
-		r.processTemplateWithList(template, file)
+		r.processTemplateWithList(template, request, file)
 		file.Close()
 		return
 	}
 
 	// Handle stdin input
 	if r.options.Stdin {
-		r.processTemplateWithList(template, os.Stdin)
+		r.processTemplateWithList(template, request, os.Stdin)
 	}
 }
 
 // processDomain processes the list with a template
-func (r *Runner) processTemplateWithList(template *templates.Template, reader io.Reader) {
+func (r *Runner) processTemplateWithList(template *templates.Template, request *requests.Request, reader io.Reader) {
 	// Display the message for the template
 	message := fmt.Sprintf("[%s] Loaded template %s (@%s)", template.ID, template.Info.Name, template.Info.Author)
 	if template.Info.Severity != "" {
@@ -144,6 +128,8 @@ func (r *Runner) processTemplateWithList(template *templates.Template, reader io
 		defer writer.Flush()
 	}
 
+	client := r.makeHTTPClient(request.Redirects, request.MaxRedirects)
+
 	scanner := bufio.NewScanner(reader)
 	for scanner.Scan() {
 		text := scanner.Text()
@@ -154,7 +140,7 @@ func (r *Runner) processTemplateWithList(template *templates.Template, reader io
 		wg.Add(1)
 
 		go func(URL string) {
-			r.sendRequest(template, URL, writer)
+			r.sendRequest(template, request, URL, writer, client)
 			<-limiter
 			wg.Done()
 		}(text)
@@ -164,71 +150,69 @@ func (r *Runner) processTemplateWithList(template *templates.Template, reader io
 }
 
 // sendRequest sends a request to the target based on a template
-func (r *Runner) sendRequest(template *templates.Template, URL string, writer *bufio.Writer) {
+func (r *Runner) sendRequest(template *templates.Template, request *requests.Request, URL string, writer *bufio.Writer, client *retryablehttp.Client) {
-	for _, request := range template.Requests {
+	// Compile each request for the template based on the URL
-		// Compile each request for the template based on the URL
+	compiledRequest, err := request.MakeRequest(URL)
-		compiledRequest, err := request.MakeRequest(URL)
+	if err != nil {
+		gologger.Warningf("[%s] Could not make request %s: %s\n", template.ID, URL, err)
+		return
+	}
+
+	// Send the request to the target servers
+reqLoop:
+	for _, req := range compiledRequest {
+		resp, err := client.Do(req)
 		if err != nil {
-			gologger.Warningf("[%s] Could not make request %s: %s\n", template.ID, URL, err)
+			if resp != nil {
+				resp.Body.Close()
+			}
+			gologger.Warningf("[%s] Could not send request %s: %s\n", template.ID, URL, err)
+			return
+		}
+
+		data, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			io.Copy(ioutil.Discard, resp.Body)
+			resp.Body.Close()
+			gologger.Warningf("[%s] Could not read body %s: %s\n", template.ID, URL, err)
 			continue
 		}
+		resp.Body.Close()
 
-		// Send the request to the target servers
+		body := unsafeToString(data)
-	reqLoop:
+
-		for _, req := range compiledRequest {
+		var headers string
-			resp, err := r.client.Do(req)
+		for _, matcher := range request.Matchers {
-			if err != nil {
+			// Only build the headers string if the matcher asks for it
-				if resp != nil {
+			part := matcher.GetPart()
-					resp.Body.Close()
+			if part == matchers.AllPart || part == matchers.HeaderPart && headers == "" {
-				}
+				headers = headersToString(resp.Header)
-				gologger.Warningf("[%s] Could not send request %s: %s\n", template.ID, URL, err)
-				return
 			}
 
-			data, err := ioutil.ReadAll(resp.Body)
+			// Check if the matcher matched
-			if err != nil {
+			if matcher.Match(resp, body, headers) {
-				io.Copy(ioutil.Discard, resp.Body)
+				// If there is an extractor, run it.
-				resp.Body.Close()
+				var extractorResults []string
-				gologger.Warningf("[%s] Could not read body %s: %s\n", template.ID, URL, err)
+				for _, extractor := range request.Extractors {
-				continue
+					part := extractor.GetPart()
-			}
+					if part == extractors.AllPart || part == extractors.HeaderPart && headers == "" {
-			resp.Body.Close()
+						headers = headersToString(resp.Header)
-
-			body := unsafeToString(data)
-
-			var headers string
-			for _, matcher := range request.Matchers {
-				// Only build the headers string if the matcher asks for it
-				part := matcher.GetPart()
-				if part == matchers.AllPart || part == matchers.HeaderPart && headers == "" {
-					headers = headersToString(resp.Header)
-				}
-
-				// Check if the matcher matched
-				if matcher.Match(resp, body, headers) {
-					// If there is an extractor, run it.
-					var extractorResults []string
-					for _, extractor := range request.Extractors {
-						part := extractor.GetPart()
-						if part == extractors.AllPart || part == extractors.HeaderPart && headers == "" {
-							headers = headersToString(resp.Header)
-						}
-						extractorResults = append(extractorResults, extractor.Extract(body, headers)...)
 					}
+					extractorResults = append(extractorResults, extractor.Extract(body, headers)...)
+				}
 
-					// All the matchers matched, print the output on the screen
+				// All the matchers matched, print the output on the screen
-					output := buildOutput(template, req, extractorResults, matcher)
+				output := buildOutput(template, req, extractorResults, matcher)
-					gologger.Silentf("%s", output)
+				gologger.Silentf("%s", output)
 
-					if writer != nil {
+				if writer != nil {
-						r.outputMutex.Lock()
+					r.outputMutex.Lock()
-						writer.WriteString(output)
+					writer.WriteString(output)
-						r.outputMutex.Unlock()
+					r.outputMutex.Unlock()
-					}
 				}
 			}
-			continue reqLoop
 		}
+		continue reqLoop
 	}
 }
 
@@ -263,3 +247,40 @@ func buildOutput(template *templates.Template, req *retryablehttp.Request, extra
 
 	return builder.String()
 }
+
+// makeHTTPClient creates a HTTP client with configurable redirect field
+func (r *Runner) makeHTTPClient(redirects bool, maxRedirects int) *retryablehttp.Client {
+	retryablehttpOptions := retryablehttp.DefaultOptionsSpraying
+	retryablehttpOptions.RetryWaitMax = 10 * time.Second
+	retryablehttpOptions.RetryMax = r.options.Retries
+
+	// Create the HTTP Client
+	client := retryablehttp.NewWithHTTPClient(&http.Client{
+		Transport: &http.Transport{
+			MaxIdleConnsPerHost: -1,
+			TLSClientConfig: &tls.Config{
+				Renegotiation:      tls.RenegotiateOnceAsClient,
+				InsecureSkipVerify: true,
+			},
+			DisableKeepAlives: true,
+		},
+		Timeout: time.Duration(r.options.Timeout) * time.Second,
+		CheckRedirect: func(_ *http.Request, requests []*http.Request) error {
+			if !redirects {
+				return http.ErrUseLastResponse
+			}
+			if maxRedirects == 0 {
+				if len(requests) > 10 {
+					return http.ErrUseLastResponse
+				}
+				return nil
+			}
+			if len(requests) > maxRedirects {
+				return http.ErrUseLastResponse
+			}
+			return nil
+		},
+	}, retryablehttpOptions)
+	client.CheckRetry = retryablehttp.HostSprayRetryPolicy()
+	return client
+}

pkg/requests/requests.go

@@ -28,6 +28,10 @@ type Request struct {
 	// Extractors contains the extraction mechanism for the request to identify
 	// and extract parts of the response.
 	Extractors []*extractors.Extractor `yaml:"extractors,omitempty"`
+	// Redirects specifies whether redirects should be followed.
+	Redirects bool `yaml:"redirects,omitempty"`
+	// MaxRedirects is the maximum number of redirects that should be followed.
+	MaxRedirects int `yaml:"max-redirects,omitempty"`
 }
 
 // MakeRequest creates a *http.Request from a request template