Added reworked generators package

v2/pkg/protocols/common/generators/generators.go

@@ -0,0 +1,216 @@
+// Inspired from https://github.com/ffuf/ffuf/blob/master/pkg/input/input.go
+
+package generators
+
+// Generator is the generator struct for generating payloads
+type Generator struct {
+	Type     Type
+	payloads map[string][]string
+}
+
+// Type is type of attack
+type Type int
+
+const (
+	// Sniper replaces each variables with values at a time.
+	Sniper Type = iota + 1
+	// PitchFork replaces variables with positional value from multiple wordlists
+	PitchFork
+	// ClusterBomb replaces variables with all possible combinations of values
+	ClusterBomb
+)
+
+// StringToType is an table for conversion of attack type from string.
+var StringToType = map[string]Type{
+	"sniper":      Sniper,
+	"pitchfork":   PitchFork,
+	"clusterbomb": ClusterBomb,
+}
+
+// New creates a new generator structure for payload generation
+func New(payloads map[string]interface{}, Type Type) (*Generator, error) {
+	compiled, err := loadPayloads(payloads)
+	if err != nil {
+		return nil, err
+	}
+	return &Generator{Type: Type, payloads: compiled}, nil
+}
+
+// Iterator is a single instance of an iterator for a generator structure
+type Iterator struct {
+	Type        Type
+	position    int
+	msbIterator int
+	payloads    []*payloadIterator
+}
+
+// NewIterator creates a new iterator for the payloads generator
+func (g *Generator) NewIterator() *Iterator {
+	var payloads []*payloadIterator
+
+	for name, values := range g.payloads {
+		payloads = append(payloads, &payloadIterator{name: name, values: values})
+	}
+	return &Iterator{Type: g.Type, payloads: payloads}
+}
+
+// Next returns true if there are more inputs in iterator
+func (i *Iterator) Next() bool {
+	if i.position >= i.Total() {
+		return false
+	}
+	i.position++
+	return true
+}
+
+//Total returns the amount of input combinations available
+func (i *Iterator) Total() int {
+	count := 0
+	switch i.Type {
+	case Sniper:
+		for _, p := range i.payloads {
+			if p.Total() > count {
+				count = p.Total()
+			}
+		}
+	case PitchFork:
+		for _, p := range i.payloads {
+			if p.Total() > count {
+				count = p.Total()
+			}
+		}
+	case ClusterBomb:
+		count = 1
+		for _, p := range i.payloads {
+			count = count * p.Total()
+		}
+	}
+	return count
+}
+
+// Value returns the next value for an iterator
+func (i *Iterator) Value() map[string]interface{} {
+	switch i.Type {
+	case Sniper:
+		return i.sniperValue()
+	case PitchFork:
+		return i.pitchforkValue()
+	case ClusterBomb:
+		return i.clusterbombValue()
+	default:
+		return i.sniperValue()
+	}
+}
+
+// sniperValue returns a list of all payloads for the iterator
+func (i *Iterator) sniperValue() map[string]interface{} {
+	values := make(map[string]interface{}, len(i.payloads))
+
+	for _, p := range i.payloads {
+		if !p.Next() {
+			p.ResetPosition()
+		}
+		values[p.Keyword()] = p.Value()
+		p.IncrementPosition()
+	}
+	return values
+}
+
+// pitchforkValue returns a map of keyword:value pairs in same index
+func (i *Iterator) pitchforkValue() map[string]interface{} {
+	values := make(map[string]interface{}, len(i.payloads))
+
+	for _, p := range i.payloads {
+		if !p.Next() {
+			p.ResetPosition()
+		}
+		values[p.Keyword()] = p.Value()
+		p.IncrementPosition()
+	}
+	return values
+}
+
+// clusterbombValue returns a combination of all input pairs in key:value format.
+func (i *Iterator) clusterbombValue() map[string]interface{} {
+	values := make(map[string]interface{}, len(i.payloads))
+
+	// Should we signal the next InputProvider in the slice to increment
+	signalNext := false
+	first := true
+	for index, p := range i.payloads {
+		if signalNext {
+			p.IncrementPosition()
+			signalNext = false
+		}
+		if !p.Next() {
+			// No more inputs in this inputprovider
+			if index == i.msbIterator {
+				// Reset all previous wordlists and increment the msb counter
+				i.msbIterator++
+				i.clusterbombIteratorReset()
+				// Start again
+				return i.clusterbombValue()
+			}
+			p.ResetPosition()
+			signalNext = true
+		}
+		values[p.Keyword()] = p.Value()
+		if first {
+			p.IncrementPosition()
+			first = false
+		}
+	}
+	return values
+}
+
+func (i *Iterator) clusterbombIteratorReset() {
+	for index, p := range i.payloads {
+		if index < i.msbIterator {
+			p.ResetPosition()
+		}
+		if index == i.msbIterator {
+			p.IncrementPosition()
+		}
+	}
+}
+
+// payloadIterator is a single instance of an iterator for a single payload list.
+type payloadIterator struct {
+	index  int
+	name   string
+	values []string
+}
+
+// Next returns true if there are more values in payload iterator
+func (i *payloadIterator) Next() bool {
+	if i.index == len(i.values)-1 {
+		return false
+	}
+	return true
+}
+
+// Position returns the position of reader in payload iterator
+func (i *payloadIterator) Position() int {
+	return i.index
+}
+
+func (i *payloadIterator) ResetPosition() {
+	i.index = 0
+}
+
+func (i *payloadIterator) IncrementPosition() {
+	i.index++
+}
+
+func (i *payloadIterator) Value() string {
+	value := i.values[i.index]
+	return value
+}
+
+func (i *payloadIterator) Total() int {
+	return len(i.values)
+}
+
+func (i *payloadIterator) Keyword() string {
+	return i.name
+}

v2/pkg/protocols/common/generators/generators_test.go

@@ -0,0 +1,18 @@
+package generators
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSniperGenerator(t *testing.T) {
+	generator, err := New(map[string]interface{}{"username": []string{"admin", "password", "login", "test"}}, Sniper)
+	require.Nil(t, err, "could not create generator")
+
+	iterator := generator.NewIterator()
+	for iterator.Next() {
+		fmt.Printf("value: %v\n", iterator.Value())
+	}
+}

v2/pkg/protocols/common/generators/load.go

@@ -0,0 +1,60 @@
+package generators
+
+import (
+	"bufio"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/pkg/errors"
+	"github.com/spf13/cast"
+)
+
+// loadPayloads loads the input payloads from a map to a data map
+func loadPayloads(payloads map[string]interface{}) (map[string][]string, error) {
+	loadedPayloads := make(map[string][]string)
+
+	for name, payload := range payloads {
+		switch pt := payload.(type) {
+		case string:
+			elements := strings.Split(pt, "\n")
+			//golint:gomnd // this is not a magic number
+			if len(elements) >= 2 {
+				loadedPayloads[name] = elements
+			} else {
+				payloads, err := loadPayloadsFromFile(pt)
+				if err != nil {
+					return nil, errors.Wrap(err, "could not load payloads")
+				}
+				loadedPayloads[name] = payloads
+			}
+		case interface{}:
+			loadedPayloads[name] = cast.ToStringSlice(pt)
+		}
+	}
+	return loadedPayloads, nil
+}
+
+// loadPayloadsFromFile loads a file to a string slice
+func loadPayloadsFromFile(filepath string) ([]string, error) {
+	var lines []string
+
+	file, err := os.Open(filepath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		text := scanner.Text()
+		if text == "" {
+			continue
+		}
+		lines = append(lines, text)
+	}
+	if err := scanner.Err(); err != nil && err != io.EOF {
+		return lines, scanner.Err()
+	}
+	return lines, nil
+}