Merge pull request #956 from forgedhallpass/master

Re-introducing custom template info attribute support within the new struct

v2/pkg/model/model.go

@@ -17,6 +17,7 @@ type Info struct {
 	Description    string                  `json:"description" yaml:"description"`
 	Reference      StringSlice             `json:"reference" yaml:"reference"`
 	SeverityHolder severity.SeverityHolder `json:"severity" yaml:"severity"`
+	CustomFields   map[string]string       `json:"custom-fields,omitempty" yaml:"custom-fields,omitempty"`
 }
 
 // StringSlice represents a single (in-lined) or multiple string value(s).
@@ -42,13 +43,17 @@ func (stringSlice StringSlice) ToSlice() []string {
 	}
 }
 
+func (stringSlice StringSlice) String() string {
+	return strings.Join(stringSlice.ToSlice(), ", ")
+}
+
 func (stringSlice *StringSlice) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	marshalledSlice, err := marshalStringToSlice(unmarshal)
 	if err != nil {
 		return err
 	}
 
-	result := make([]string, len(marshalledSlice))
+	result := make([]string, 0, len(marshalledSlice))
 	//nolint:gosimple,nolintlint //cannot be replaced with result = append(result, slices...) because the values are being normalized
 	for _, value := range marshalledSlice {
 		result = append(result, strings.ToLower(strings.TrimSpace(value))) // TODO do we need to introduce RawStringSlice and/or NormalizedStringSlices?

v2/pkg/model/model_test.go

@@ -2,6 +2,8 @@ package model
 
 import (
 	"encoding/json"
+	"gopkg.in/yaml.v2"
+	"strings"
 	"testing"
 
 	"github.com/projectdiscovery/nuclei/v2/internal/severity"
@@ -24,3 +26,62 @@ func TestInfoJsonMarshal(t *testing.T) {
 	expected := `{"name":"Test Template Name","author":["forgedhallpass","ice3man"],"tags":["cve","misc"],"description":"Test description","reference":"reference1","severity":"high"}`
 	assert.Equal(t, expected, string(result))
 }
+
+func TestUnmarshal(t *testing.T) {
+	templateName := "Test Template"
+	authors := []string{"forgedhallpass", "ice3man"}
+	tags := []string{"cve", "misc"}
+	references := []string{"http://test.com", "http://domain.com"}
+
+	dynamicKey1 := "customDynamicKey1"
+	dynamicKey2 := "customDynamicKey2"
+
+	dynamicKeysMap := map[string]string{
+		dynamicKey1: "customDynamicValue1",
+		dynamicKey2: "customDynamicValue2",
+	}
+
+	assertUnmarshalledTemplateInfo := func(t *testing.T, yamlPayload string) Info {
+		info := Info{}
+		err := yaml.Unmarshal([]byte(yamlPayload), &info)
+		assert.Nil(t, err)
+		assert.Equal(t, info.Name, templateName)
+		assert.Equal(t, info.Authors.ToSlice(), authors)
+		assert.Equal(t, info.Tags.ToSlice(), tags)
+		assert.Equal(t, info.SeverityHolder.Severity, severity.Critical)
+		assert.Equal(t, info.Reference.ToSlice(), references)
+		assert.Equal(t, info.CustomFields, dynamicKeysMap)
+		return info
+	}
+
+	yamlPayload1 := `
+  name: ` + templateName + `
+  author: ` + strings.Join(authors, ", ") + `
+  tags: ` + strings.Join(tags, ", ") + `
+  severity: critical
+  reference: ` + strings.Join(references, ", ") + `
+  custom-fields:
+     ` + dynamicKey1 + `: ` + dynamicKeysMap[dynamicKey1] + `
+     ` + dynamicKey2 + `: ` + dynamicKeysMap[dynamicKey2] + `
+`
+	yamlPayload2 := `
+  name: ` + templateName + `
+  author:
+    - ` + authors[0] + `
+    - ` + authors[1] + `
+  tags:
+    - ` + tags[0] + `
+    - ` + tags[1] + `
+  severity: critical
+  reference:
+    - ` + references[0] + ` # comments are not unmarshalled
+    - ` + references[1] + `
+  custom-fields:
+     ` + dynamicKey1 + `: ` + dynamicKeysMap[dynamicKey1] + `
+     ` + dynamicKey2 + `: ` + dynamicKeysMap[dynamicKey2] + `
+`
+
+	info1 := assertUnmarshalledTemplateInfo(t, yamlPayload1)
+	info2 := assertUnmarshalledTemplateInfo(t, yamlPayload2)
+	assert.Equal(t, info1, info2)
+}

v2/pkg/reporting/format/format.go

@@ -119,9 +119,8 @@ func MarkdownDescription(event *output.ResultEvent) string { // TODO remove the
 
 	reference := event.Info.Reference
 	if !reference.IsEmpty() {
-		builder.WriteString("\nReference: \n")
+		builder.WriteString("\nReferences: \n")
 
-		/*TODO couldn't the following code replace the logic below?
 		referenceSlice := reference.ToSlice()
 		for i, item := range referenceSlice {
 			builder.WriteString("- ")
@@ -129,23 +128,6 @@ func MarkdownDescription(event *output.ResultEvent) string { // TODO remove the
 			if len(referenceSlice)-1 != i {
 				builder.WriteString("\n")
 			}
-		}*/
-
-		switch value := reference.Value.(type) {
-		case string:
-			if !strings.HasPrefix(value, "-") {
-				builder.WriteString("- ")
-			}
-			builder.WriteString(value)
-		case []interface{}:
-			slice := types.ToStringSlice(value)
-			for i, item := range slice {
-				builder.WriteString("- ")
-				builder.WriteString(item)
-				if len(slice)-1 != i {
-					builder.WriteString("\n")
-				}
-			}
 		}
 	}
 
@@ -171,23 +153,25 @@ func GetMatchedTemplate(event *output.ResultEvent) string {
 }
 
 func ToMarkdownTableString(templateInfo *model.Info) string {
-	fields := map[string]string{
-		"Name":        templateInfo.Name,
-		"Authors":     sliceToString(templateInfo.Authors),
-		"Tags":        sliceToString(templateInfo.Tags),
-		"Description": templateInfo.Description,
-		"Severity":    templateInfo.SeverityHolder.Severity.String(),
-	}
+	fields := utils.NewEmptyInsertionOrderedStringMap(5)
+	fields.Set("Name", templateInfo.Name)
+	fields.Set("Authors", templateInfo.Authors.String())
+	fields.Set("Tags", templateInfo.Tags.String())
+	fields.Set("Severity", templateInfo.SeverityHolder.Severity.String())
+	fields.Set("Description", templateInfo.Description)
 
 	builder := &bytes.Buffer{}
-	for k, v := range fields {
-		if utils.IsNotBlank(v) {
-			builder.WriteString(fmt.Sprintf("| %s | %s |\n", k, v))
+
+	toMarkDownTable := func(insertionOrderedStringMap *utils.InsertionOrderedStringMap) {
+		insertionOrderedStringMap.ForEach(func(key string, value string) {
+			if utils.IsNotBlank(value) {
+				builder.WriteString(fmt.Sprintf("| %s | %s |\n", key, value))
 			}
-	}
-	return builder.String()
+		})
 	}
 
-func sliceToString(stringSlice model.StringSlice) string {
-	return strings.Join(stringSlice.ToSlice(), ", ")
+	toMarkDownTable(fields)
+	toMarkDownTable(utils.NewInsertionOrderedStringMap(templateInfo.CustomFields))
+
+	return builder.String()
 }

v2/pkg/reporting/format/format_test.go

@@ -0,0 +1,45 @@
+package format
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/projectdiscovery/nuclei/v2/internal/severity"
+	"github.com/projectdiscovery/nuclei/v2/pkg/model"
+)
+
+func TestToMarkdownTableString(t *testing.T) {
+	info := model.Info{
+		Name:           "Test Template Name",
+		Authors:        model.StringSlice{Value: []string{"forgedhallpass", "ice3man"}},
+		Description:    "Test description",
+		SeverityHolder: severity.SeverityHolder{Severity: severity.High},
+		Tags:           model.StringSlice{Value: []string{"cve", "misc"}},
+		Reference:      model.StringSlice{Value: "reference1"},
+		CustomFields: map[string]string{
+			"customDynamicKey1": "customDynamicValue1",
+			"customDynamicKey2": "customDynamicValue2",
+		},
+	}
+
+	result := ToMarkdownTableString(&info)
+
+	expectedOrderedAttributes := `| Name | Test Template Name |
+| Authors | forgedhallpass, ice3man |
+| Tags | cve, misc |
+| Severity | high |
+| Description | Test description |`
+
+	expectedDynamicAttributes := []string{
+		"| customDynamicKey1 | customDynamicValue1 |",
+		"| customDynamicKey2 | customDynamicValue2 |",
+		"", // the expected result ends in a new line (\n)
+	}
+
+	actualAttributeSlice := strings.Split(result, "\n")
+	dynamicAttributeIndex := len(actualAttributeSlice) - len(expectedDynamicAttributes)
+	assert.Equal(t, strings.Split(expectedOrderedAttributes, "\n"), actualAttributeSlice[:dynamicAttributeIndex]) // the first part of the result is ordered
+	assert.ElementsMatch(t, expectedDynamicAttributes, actualAttributeSlice[dynamicAttributeIndex:])              // dynamic parameters are not ordered
+}

v2/pkg/reporting/trackers/jira/jira.go

@@ -181,9 +181,8 @@ func jiraFormatDescription(event *output.ResultEvent) string { // TODO remove th
 
 	reference := event.Info.Reference
 	if !reference.IsEmpty() {
-		builder.WriteString("\nReference: \n")
+		builder.WriteString("\nReferences: \n")
 
-		/*TODO couldn't the following code replace the logic below?
 		referenceSlice := reference.ToSlice()
 		for i, item := range referenceSlice {
 			builder.WriteString("- ")
@@ -191,23 +190,6 @@ func jiraFormatDescription(event *output.ResultEvent) string { // TODO remove th
 			if len(referenceSlice)-1 != i {
 				builder.WriteString("\n")
 			}
-		}*/
-
-		switch v := reference.Value.(type) {
-		case string:
-			if !strings.HasPrefix(v, "-") {
-				builder.WriteString("- ")
-			}
-			builder.WriteString(v)
-		case []interface{}:
-			slice := types.ToStringSlice(v)
-			for i, item := range slice {
-				builder.WriteString("- ")
-				builder.WriteString(item)
-				if len(slice)-1 != i {
-					builder.WriteString("\n")
-				}
-			}
 		}
 	}
 	builder.WriteString("\n---\nGenerated by [Nuclei|https://github.com/projectdiscovery/nuclei]")

v2/pkg/utils/insertion_ordered_map.go

@@ -0,0 +1,37 @@
+package utils
+
+type InsertionOrderedStringMap struct {
+	keys   []string `yaml:"-"`
+	values map[string]string
+}
+
+func NewEmptyInsertionOrderedStringMap(size int) *InsertionOrderedStringMap {
+	return &InsertionOrderedStringMap{
+		keys:   make([]string, 0, size),
+		values: make(map[string]string, size),
+	}
+}
+
+func NewInsertionOrderedStringMap(stringMap map[string]string) *InsertionOrderedStringMap {
+	result := NewEmptyInsertionOrderedStringMap(len(stringMap))
+
+	for k, v := range stringMap {
+		result.Set(k, v)
+	}
+
+	return result
+}
+
+func (insertionOrderedStringMap *InsertionOrderedStringMap) ForEach(fn func(key string, data string)) {
+	for _, key := range insertionOrderedStringMap.keys {
+		fn(key, insertionOrderedStringMap.values[key])
+	}
+}
+
+func (insertionOrderedStringMap *InsertionOrderedStringMap) Set(key string, value string) {
+	_, present := insertionOrderedStringMap.values[key]
+	insertionOrderedStringMap.values[key] = value
+	if !present {
+		insertionOrderedStringMap.keys = append(insertionOrderedStringMap.keys, key)
+	}
+}