Merge pull request #1873 from projectdiscovery/issue-1565-dsl-extractors

adding support for dsl extractors
projectdiscovery · Apr 21, 2022 · 4573c4e · 4573c4e
2 parents 294dcc7 + 72c5c39
commit 4573c4e
Show file tree

Hide file tree

Showing 12 changed files with 68 additions and 7 deletions.
diff --git a/v2/pkg/operators/extractors/compile.go b/v2/pkg/operators/extractors/compile.go
@@ -5,7 +5,9 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/Knetic/govaluate"
 	"github.com/itchyny/gojq"
+	"github.com/projectdiscovery/nuclei/v2/pkg/operators/common/dsl"
 )
 
 // CompileExtractors performs the initial setup operation on an extractor
@@ -40,6 +42,14 @@ func (e *Extractor) CompileExtractors() error {
 		e.jsonCompiled = append(e.jsonCompiled, compiled)
 	}
 
+	for _, dslExp := range e.DSL {
+		compiled, err := govaluate.NewEvaluableExpressionWithFunctions(dslExp, dsl.HelperFunctions())
+		if err != nil {
+			return fmt.Errorf("could not compile dsl: %s", dslExp)
+		}
+		e.dslCompiled = append(e.dslCompiled, compiled)
+	}
+
 	if e.CaseInsensitive {
 		if e.GetType() != KValExtractor {
 			return fmt.Errorf("case-insensitive flag is supported only for 'kval' extractors (not '%s')", e.Type)

diff --git a/v2/pkg/operators/extractors/extract.go b/v2/pkg/operators/extractors/extract.go
@@ -2,6 +2,7 @@ package extractors
 
 import (
 	"encoding/json"
+	"fmt"
 	"strings"
 
 	"github.com/antchfx/htmlquery"
@@ -122,3 +123,24 @@ func (e *Extractor) ExtractJSON(corpus string) map[string]struct{} {
 	}
 	return results
 }
+
+// ExtractDSL execute the expression and returns the results
+func (e *Extractor) ExtractDSL(data map[string]interface{}) map[string]struct{} {
+	results := make(map[string]struct{})
+
+	for _, compiledExpression := range e.dslCompiled {
+		result, err := compiledExpression.Evaluate(data)
+		if err != nil {
+			return results
+		}
+
+		if result != nil {
+			resultString := fmt.Sprint(result)
+			if resultString != "" {
+				results[resultString] = struct{}{}
+			}
+		}
+	}
+
+	return results
+}
diff --git a/v2/pkg/operators/extractors/extractor_types.go b/v2/pkg/operators/extractors/extractor_types.go
@@ -21,6 +21,8 @@ const (
 	XPathExtractor
 	// name:json
 	JSONExtractor
+	// name:dsl
+	DSLExtractor
 	limit
 )
 
@@ -30,6 +32,7 @@ var extractorMappings = map[ExtractorType]string{
 	KValExtractor:  "kval",
 	XPathExtractor: "xpath",
 	JSONExtractor:  "json",
+	DSLExtractor:   "dsl",
 }
 
 // GetType returns the type of the matcher

diff --git a/v2/pkg/operators/extractors/extractors.go b/v2/pkg/operators/extractors/extractors.go
@@ -3,6 +3,7 @@ package extractors
 import (
 	"regexp"
 
+	"github.com/Knetic/govaluate"
 	"github.com/itchyny/gojq"
 )
 
@@ -87,6 +88,11 @@ type Extractor struct {
 	// jsonCompiled is the compiled variant
 	jsonCompiled []*gojq.Code
 
+	// description: |
+	//   Extracts using DSL expressions.
+	DSL         []string
+	dslCompiled []*govaluate.EvaluableExpression
+
 	// description: |
 	//   Part is the part of the request response to extract data from.
 	//

diff --git a/v2/pkg/operators/extractors/util.go b/v2/pkg/operators/extractors/util.go
@@ -0,0 +1,6 @@
+package extractors
+
+// SupportsMap determines if the extractor type requires a map
+func SupportsMap(extractor *Extractor) bool {
+	return extractor.Type.ExtractorType == KValExtractor || extractor.Type.ExtractorType == DSLExtractor
+}
diff --git a/v2/pkg/protocols/dns/operators.go b/v2/pkg/protocols/dns/operators.go
@@ -48,7 +48,7 @@ func (request *Request) Match(data map[string]interface{}, matcher *matchers.Mat
 // Extract performs extracting operation for an extractor on model and returns true or false.
 func (request *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
 	item, ok := request.getMatchPart(extractor.Part, data)
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 
@@ -57,6 +57,8 @@ func (request *Request) Extract(data map[string]interface{}, extractor *extracto
 		return extractor.ExtractRegex(types.ToString(item))
 	case extractors.KValExtractor:
 		return extractor.ExtractKval(data)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }

diff --git a/v2/pkg/protocols/file/operators.go b/v2/pkg/protocols/file/operators.go
@@ -37,7 +37,7 @@ func (request *Request) Match(data map[string]interface{}, matcher *matchers.Mat
 // Extract performs extracting operation for an extractor on model and returns true or false.
 func (request *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
 	itemStr, ok := request.getMatchPart(extractor.Part, data)
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 
@@ -46,6 +46,8 @@ func (request *Request) Extract(data map[string]interface{}, extractor *extracto
 		return extractor.ExtractRegex(itemStr)
 	case extractors.KValExtractor:
 		return extractor.ExtractKval(data)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }

diff --git a/v2/pkg/protocols/headless/operators.go b/v2/pkg/protocols/headless/operators.go
@@ -37,7 +37,7 @@ func (request *Request) Match(data map[string]interface{}, matcher *matchers.Mat
 // Extract performs extracting operation for an extractor on model and returns true or false.
 func (request *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
 	itemStr, ok := request.getMatchPart(extractor.Part, data)
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 
@@ -46,6 +46,8 @@ func (request *Request) Extract(data map[string]interface{}, extractor *extracto
 		return extractor.ExtractRegex(itemStr)
 	case extractors.KValExtractor:
 		return extractor.ExtractKval(data)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }

diff --git a/v2/pkg/protocols/http/operators.go b/v2/pkg/protocols/http/operators.go
@@ -58,7 +58,7 @@ func getStatusCode(data map[string]interface{}) (int, bool) {
 // Extract performs extracting operation for an extractor on model and returns true or false.
 func (request *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
 	item, ok := request.getMatchPart(extractor.Part, data)
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 	switch extractor.GetType() {
@@ -70,6 +70,8 @@ func (request *Request) Extract(data map[string]interface{}, extractor *extracto
 		return extractor.ExtractHTML(item)
 	case extractors.JSONExtractor:
 		return extractor.ExtractJSON(item)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }

diff --git a/v2/pkg/protocols/network/operators.go b/v2/pkg/protocols/network/operators.go
@@ -37,7 +37,7 @@ func (request *Request) Match(data map[string]interface{}, matcher *matchers.Mat
 // Extract performs extracting operation for an extractor on model and returns true or false.
 func (request *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
 	itemStr, ok := request.getMatchPart(extractor.Part, data)
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 
@@ -46,6 +46,8 @@ func (request *Request) Extract(data map[string]interface{}, extractor *extracto
 		return extractor.ExtractRegex(itemStr)
 	case extractors.KValExtractor:
 		return extractor.ExtractKval(data)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }

diff --git a/v2/pkg/protocols/offlinehttp/operators.go b/v2/pkg/protocols/offlinehttp/operators.go
@@ -58,14 +58,16 @@ func getStatusCode(data map[string]interface{}) (int, bool) {
 // Extract performs extracting operation for an extractor on model and returns true or false.
 func (request *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
 	item, ok := getMatchPart(extractor.Part, data)
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 	switch extractor.GetType() {
 	case extractors.RegexExtractor:
 		return extractor.ExtractRegex(item)
 	case extractors.KValExtractor:
 		return extractor.ExtractKval(data)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }

diff --git a/v2/pkg/protocols/protocols.go b/v2/pkg/protocols/protocols.go
@@ -144,7 +144,7 @@ func MakeDefaultExtractFunc(data map[string]interface{}, extractor *extractors.E
 	}
 
 	item, ok := data[part]
-	if !ok && extractor.Type.ExtractorType != extractors.KValExtractor {
+	if !ok && !extractors.SupportsMap(extractor) {
 		return nil
 	}
 	itemStr := types.ToString(item)
@@ -158,6 +158,8 @@ func MakeDefaultExtractFunc(data map[string]interface{}, extractor *extractors.E
 		return extractor.ExtractJSON(itemStr)
 	case extractors.XPathExtractor:
 		return extractor.ExtractHTML(itemStr)
+	case extractors.DSLExtractor:
+		return extractor.ExtractDSL(data)
 	}
 	return nil
 }