Merge pull request #160 from projectdiscovery/feature-workflow-unleashed

Advanced Workflows

v2/internal/runner/runner.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net/http/cookiejar"
 	"os"
 	"strings"
 	"sync"
@@ -259,6 +260,7 @@ func (r *Runner) processTemplateWithList(template *templates.Template, request i
 			ProxySocksURL: r.options.ProxySocksURL,
 			CustomHeaders: r.options.CustomHeaders,
 			JSON:          r.options.JSON,
+			CookieReuse:   value.CookieReuse,
 		})
 	}
 	if err != nil {
@@ -341,6 +343,14 @@ func (r *Runner) ProcessWorkflowWithList(workflow *workflows.Workflow) {
 func (r *Runner) ProcessWorkflow(workflow *workflows.Workflow, URL string) error {
 	script := tengo.NewScript([]byte(workflow.Logic))
 	script.SetImports(stdlib.GetModuleMap(stdlib.AllModuleNames()...))
+	var jar *cookiejar.Jar
+	if workflow.CookieReuse {
+		var err error
+		jar, err = cookiejar.New(nil)
+		if err != nil {
+			return err
+		}
+	}
 	for name, value := range workflow.Variables {
 		var writer *bufio.Writer
 		if r.output != nil {
@@ -376,6 +386,7 @@ func (r *Runner) ProcessWorkflow(workflow *workflows.Workflow, URL string) error
 					ProxyURL:      r.options.ProxyURL,
 					ProxySocksURL: r.options.ProxySocksURL,
 					CustomHeaders: r.options.CustomHeaders,
+					CookieJar:     jar,
 				}
 			} else if len(t.RequestsDNS) > 0 {
 				template.DNSOptions = &executer.DNSOptions{
@@ -426,6 +437,7 @@ func (r *Runner) ProcessWorkflow(workflow *workflows.Workflow, URL string) error
 						ProxyURL:      r.options.ProxyURL,
 						ProxySocksURL: r.options.ProxySocksURL,
 						CustomHeaders: r.options.CustomHeaders,
+						CookieJar:     jar,
 					}
 				} else if len(t.RequestsDNS) > 0 {
 					template.DNSOptions = &executer.DNSOptions{

v2/pkg/executer/executer_http.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"net/http/cookiejar"
 	"net/http/httputil"
 	"net/url"
 	"os"
@@ -36,6 +37,7 @@ type HTTPExecuter struct {
 	writer        *bufio.Writer
 	outputMutex   *sync.Mutex
 	customHeaders requests.CustomHeaders
+	CookieJar     *cookiejar.Jar
 }
 
 // HTTPOptions contains configuration options for the HTTP executer.
@@ -50,6 +52,8 @@ type HTTPOptions struct {
 	Debug         bool
 	JSON          bool
 	CustomHeaders requests.CustomHeaders
+	CookieReuse   bool
+	CookieJar     *cookiejar.Jar
 }
 
 // NewHTTPExecuter creates a new HTTP executer from a template
@@ -68,6 +72,15 @@ func NewHTTPExecuter(options *HTTPOptions) (*HTTPExecuter, error) {
 	// Create the HTTP Client
 	client := makeHTTPClient(proxyURL, options)
 	client.CheckRetry = retryablehttp.HostSprayRetryPolicy()
+	if options.CookieJar != nil {
+		client.HTTPClient.Jar = options.CookieJar
+	} else if options.CookieReuse {
+		jar, err := cookiejar.New(nil)
+		if err != nil {
+			return nil, err
+		}
+		client.HTTPClient.Jar = jar
+	}
 
 	executer := &HTTPExecuter{
 		debug:         options.Debug,
@@ -79,6 +92,7 @@ func NewHTTPExecuter(options *HTTPOptions) (*HTTPExecuter, error) {
 		outputMutex:   &sync.Mutex{},
 		writer:        options.Writer,
 		customHeaders: options.CustomHeaders,
+		CookieJar:     options.CookieJar,
 	}
 	return executer, nil
 }
@@ -95,6 +109,7 @@ func (e *HTTPExecuter) GotResults() bool {
 func (e *HTTPExecuter) ExecuteHTTP(URL string) (result Result) {
 	result.Matches = make(map[string]interface{})
 	result.Extractions = make(map[string]interface{})
+	dynamicvalues := make(map[string]string)
 	// Compile each request for the template based on the URL
 	compiledRequest, err := e.httpRequest.MakeHTTPRequest(URL)
 	if err != nil {
@@ -110,6 +125,7 @@ mainLoop:
 			return
 		}
 		e.setCustomHeaders(compiledRequest)
+		e.setDynamicValues(compiledRequest, dynamicvalues)
 		req := compiledRequest.Request
 
 		if e.debug {
@@ -188,6 +204,9 @@ mainLoop:
 		var extractorResults []string
 		for _, extractor := range e.httpRequest.Extractors {
 			for match := range extractor.Extract(resp, body, headers) {
+				if _, ok := dynamicvalues[extractor.Name]; !ok {
+					dynamicvalues[extractor.Name] = match
+				}
 				extractorResults = append(extractorResults, match)
 			}
 			// probably redundant but ensures we snapshot current payload values when extractors are valid
@@ -289,7 +308,22 @@ func (e *HTTPExecuter) setCustomHeaders(r *requests.CompiledHTTP) {
 		headerName, headerValue := tokens[0], strings.Join(tokens[1:], "")
 		headerName = strings.TrimSpace(headerName)
 		headerValue = strings.TrimSpace(headerValue)
-		r.Request.Header.Set(headerName, headerValue)
+		r.Request.Header[headerName] = []string{headerValue}
+	}
+}
+
+// for now supports only headers
+func (e *HTTPExecuter) setDynamicValues(r *requests.CompiledHTTP, dynamicValues map[string]string) {
+	for dk, dv := range dynamicValues {
+		// replace within header values
+		for k, v := range r.Request.Header {
+			for i, vv := range v {
+				if strings.Contains(vv, "{{"+dk+"}}") {
+					// coerce values to string and picks only the first value
+					r.Request.Header[k][i] = strings.ReplaceAll(r.Request.Header[k][i], "{{"+dk+"}}", dv)
+				}
+			}
+		}
 	}
 }
 

v2/pkg/requests/http-request.go

@@ -34,6 +34,8 @@ type HTTPRequest struct {
 	Headers map[string]string `yaml:"headers,omitempty"`
 	// Body is an optional parameter which contains the request body for POST methods, etc
 	Body string `yaml:"body,omitempty"`
+	// CookieReuse is an optional setting that makes cookies shared within requests
+	CookieReuse bool `yaml:"cookie-reuse,omitempty"`
 	// Matchers contains the detection mechanism for the request to identify
 	// whether the request was successful
 	Matchers []*matchers.Matcher `yaml:"matchers,omitempty"`
@@ -187,7 +189,7 @@ func (r *HTTPRequest) handleSimpleRaw(raw string, baseURL string, values map[str
 
 	// copy headers
 	for key, value := range compiledRequest.Headers {
-		req.Header.Set(key, value)
+		req.Header[key] = []string{value}
 	}
 
 	request, err := r.fillRequest(req, values)
@@ -242,7 +244,7 @@ func (r *HTTPRequest) handleRawWithPaylods(raw string, baseURL string, values, g
 
 	// copy headers
 	for key, value := range compiledRequest.Headers {
-		req.Header.Set(key, value)
+		req.Header[key] = []string{value}
 	}
 
 	request, err := r.fillRequest(req, values)
@@ -265,7 +267,7 @@ func (r *HTTPRequest) fillRequest(req *http.Request, values map[string]interface
 
 	// Set the header values requested
 	for header, value := range r.Headers {
-		req.Header.Set(header, replacer.Replace(value))
+		req.Header[header] = []string{replacer.Replace(value)}
 	}
 
 	// Set some headers only if the header wasn't supplied by the user

v2/pkg/workflows/var.go

@@ -34,31 +34,29 @@ func (n *NucleiVar) CanCall() bool {
 	return true
 }
 
-// Call logic - actually it doesn't require arguments
+// Call logic - args[0]=headers, args[1]=payloads
 func (n *NucleiVar) Call(args ...tengo.Object) (ret tengo.Object, err error) {
 	n.InternalVars = make(map[string]interface{})
+	headers := make(map[string]string)
 	externalVars := make(map[string]interface{})
 
 	// if external variables are specified and matches the template ones, these gets overwritten
-	if len(args) == 1 {
-		m := args[0]
-		if m.CanIterate() {
-			i := m.Iterate()
-			for i.Next() {
-				key, ok := tengo.ToString(i.Key())
-				if !ok {
-					continue
-				}
-				value := tengo.ToInterface(i.Value())
-				externalVars[key] = value
-			}
-		}
+	if len(args) >= 1 {
+		headers = iterableToMapString(args[0])
+	}
+
+	// if external variables are specified and matches the template ones, these gets overwritten
+	if len(args) >= 2 {
+		externalVars = iterableToMap(args[1])
 	}
 
 	var gotResult bool
 	for _, template := range n.Templates {
 		if template.HTTPOptions != nil {
 			for _, request := range template.HTTPOptions.Template.RequestsHTTP {
+				// apply externally supplied payloads if any
+				request.Headers = generators.MergeMapsWithStrings(request.Headers, headers)
+				// apply externally supplied payloads if any
 				request.Payloads = generators.MergeMaps(request.Payloads, externalVars)
 				template.HTTPOptions.HTTPRequest = request
 				httpExecuter, err := executer.NewHTTPExecuter(template.HTTPOptions)
@@ -159,3 +157,38 @@ func (n *NucleiVar) IndexGet(index tengo.Object) (res tengo.Object, err error) {
 
 	return
 }
+
+func iterableToMap(t tengo.Object) map[string]interface{} {
+	m := make(map[string]interface{})
+	if t.CanIterate() {
+		i := t.Iterate()
+		for i.Next() {
+			key, ok := tengo.ToString(i.Key())
+			if !ok {
+				continue
+			}
+			value := tengo.ToInterface(i.Value())
+			m[key] = value
+		}
+	}
+
+	return m
+}
+
+func iterableToMapString(t tengo.Object) map[string]string {
+	m := make(map[string]string)
+	if t.CanIterate() {
+		i := t.Iterate()
+		for i.Next() {
+			key, ok := tengo.ToString(i.Key())
+			if !ok {
+				continue
+			}
+			if value, ok := tengo.ToString(i.Value()); ok {
+				m[key] = value
+			}
+		}
+	}
+
+	return m
+}

v2/pkg/workflows/workflows.go

@@ -6,6 +6,8 @@ type Workflow struct {
 	ID string `yaml:"id"`
 	// Info contains information about the template
 	Info Info `yaml:"info"`
+	// CookieReuse makes all cookies shared by templates within the workflow
+	CookieReuse bool `yaml:"cookie-reuse,omitempty"`
 	// Variables contains the variables accessible to the pseudo-code
 	Variables map[string]string `yaml:"variables"`
 	// Logic contains the workflow pseudo-code