Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
projectdiscovery
GitHub Repository: projectdiscovery/nuclei
 Path: blob/dev/pkg/protocols/common/generators/load_test.go
2843 views
1
package generators

2


3
import (

4
	"io"

5
	"strings"

6
	"testing"

7


8
	"github.com/pkg/errors"

9
	"github.com/projectdiscovery/nuclei/v3/pkg/catalog"

10
	"github.com/projectdiscovery/nuclei/v3/pkg/types"

11
)

12


13
type fakeCatalog struct{ catalog.Catalog }

14


15
func (f *fakeCatalog) OpenFile(filename string) (io.ReadCloser, error) {

16
	return nil, errors.New("not used")

17
}

18
func (f *fakeCatalog) GetTemplatePath(target string) ([]string, error) { return nil, nil }

19
func (f *fakeCatalog) GetTemplatesPath(definitions []string) ([]string, map[string]error) {

20
	return nil, nil

21
}

22
func (f *fakeCatalog) ResolvePath(templateName, second string) (string, error) {

23
	return templateName, nil

24
}

25


26
func newTestGenerator() *PayloadGenerator {

27
	opts := types.DefaultOptions()

28
	// inject helper loader function

29
	opts.LoadHelperFileFunction = func(path, templatePath string, _ catalog.Catalog) (io.ReadCloser, error) {

30
		switch path {

31
		case "fileA.txt":

32
			return io.NopCloser(strings.NewReader("one\n two\n\nthree\n")), nil

33
		default:

34
			return io.NopCloser(strings.NewReader("x\ny\nz\n")), nil

35
		}

36
	}

37
	return &PayloadGenerator{options: opts, catalog: &fakeCatalog{}}

38
}

39


40
func TestLoadPayloads_FastPathFile(t *testing.T) {

41
	g := newTestGenerator()

42
	out, err := g.loadPayloads(map[string]interface{}{"A": "fileA.txt"}, "")

43
	if err != nil {

44
		t.Fatalf("err: %v", err)

45
	}

46
	got := out["A"]

47
	if len(got) != 3 || got[0] != "one" || got[1] != " two" || got[2] != "three" {

48
		t.Fatalf("unexpected: %#v", got)

49
	}

50
}

51


52
func TestLoadPayloads_InlineMultiline(t *testing.T) {

53
	g := newTestGenerator()

54
	inline := "a\nb\n"

55
	out, err := g.loadPayloads(map[string]interface{}{"B": inline}, "")

56
	if err != nil {

57
		t.Fatalf("err: %v", err)

58
	}

59
	got := out["B"]

60
	if len(got) != 3 || got[0] != "a" || got[1] != "b" || got[2] != "" {

61
		t.Fatalf("unexpected: %#v", got)

62
	}

63
}

64


65
func TestLoadPayloads_SingleLineFallsBackToFile(t *testing.T) {

66
	g := newTestGenerator()

67
	inline := "fileA.txt" // single line, should be treated as file path

68
	out, err := g.loadPayloads(map[string]interface{}{"C": inline}, "")

69
	if err != nil {

70
		t.Fatalf("err: %v", err)

71
	}

72
	got := out["C"]

73
	if len(got) != 3 {

74
		t.Fatalf("unexpected len: %d", len(got))

75
	}

76
}

77


78
func TestLoadPayloads_InterfaceSlice(t *testing.T) {

79
	g := newTestGenerator()

80
	out, err := g.loadPayloads(map[string]interface{}{"D": []interface{}{"p", "q"}}, "")

81
	if err != nil {

82
		t.Fatalf("err: %v", err)

83
	}

84
	got := out["D"]

85
	if len(got) != 2 || got[0] != "p" || got[1] != "q" {

86
		t.Fatalf("unexpected: %#v", got)

87
	}

88
}

89


90
func TestLoadPayloadsFromFile_SkipsEmpty(t *testing.T) {

91
	g := newTestGenerator()

92
	rc := io.NopCloser(strings.NewReader("a\n\n\n b \n"))

93
	lines, err := g.loadPayloadsFromFile(rc)

94
	if err != nil {

95
		t.Fatalf("err: %v", err)

96
	}

97
	if len(lines) != 2 || lines[0] != "a" || lines[1] != " b " {

98
		t.Fatalf("unexpected: %#v", lines)

99
	}

100
}

101


102
func TestValidate_AllowsInlineMultiline(t *testing.T) {

103
	g := newTestGenerator()

104
	inline := "x\ny\n"

105
	if err := g.validate(map[string]interface{}{"E": inline}, ""); err != nil {

106
		t.Fatalf("validate rejected inline multiline: %v", err)

107
	}

108
}

109


110