burterm/internal/payload/generator.go

package payload

import "strconv"

// Generator — интерфейс источника payload'ов, чтобы легко
// добавлять новые типы (wordlist, range, brute force, mutation).
type Generator interface {
	// Next возвращает следующий payload и false, если генератор исчерпан.
	Next() (string, bool)
	// Reset позволяет переиспользовать генератор (напр. для sniper-режима,
	// где один и тот же генератор прогоняется по каждой точке вставки).
	Reset()
	// Len возвращает общее количество payload'ов, если оно известно заранее.
	Len() int
}

// WordlistGenerator читает payload'ы построчно из слайса (загружен из файла).
type WordlistGenerator struct {
	words []string
	pos   int
}

// NewWordlistGenerator создаёт генератор из уже загруженного списка строк.
func NewWordlistGenerator(words []string) *WordlistGenerator {
	return &WordlistGenerator{words: words}
}

func (g *WordlistGenerator) Next() (string, bool) {
	if g.pos >= len(g.words) {
		return "", false
	}
	w := g.words[g.pos]
	g.pos++
	return w, true
}

func (g *WordlistGenerator) Reset()   { g.pos = 0 }
func (g *WordlistGenerator) Len() int { return len(g.words) }

// NumericRangeGenerator — для перебора числовых ID (IDOR и т.п.).
type NumericRangeGenerator struct {
	current, from, to, step int
}

// NewNumericRangeGenerator создаёт генератор чисел от from до to
// включительно с заданным шагом.
func NewNumericRangeGenerator(from, to, step int) *NumericRangeGenerator {
	return &NumericRangeGenerator{current: from, from: from, to: to, step: step}
}

func (g *NumericRangeGenerator) Next() (string, bool) {
	if g.current > g.to {
		return "", false
	}
	v := g.current
	g.current += g.step
	return strconv.Itoa(v), true
}

func (g *NumericRangeGenerator) Reset() { g.current = g.from }
func (g *NumericRangeGenerator) Len() int {
	if g.step <= 0 {
		return 0
	}
	return (g.to-g.from)/g.step + 1
}

// ExtensionGenerator оборачивает базовый генератор и добавляет к каждому
// его слову варианты с расширениями — аналог -e у ffuf ("-e .php,.html"
// превращает "admin" в "admin", "admin.php", "admin.html"). Голое слово
// без расширения выдаётся первым — тот же порядок, что у ffuf.
type ExtensionGenerator struct {
	base       Generator
	extensions []string

	currentWord string
	// extIdx == -1 означает "нужно новое базовое слово"; иначе — индекс
	// следующего расширения, которое нужно приклеить к currentWord.
	extIdx    int
	exhausted bool
}

// NewExtensionGenerator оборачивает base, добавляя extensions (каждая
// строка — расширение вида ".php"; ведущая точка добавляется автоматически
// вызывающим кодом, сюда передаётся уже готовая строка).
func NewExtensionGenerator(base Generator, extensions []string) *ExtensionGenerator {
	g := &ExtensionGenerator{base: base, extensions: extensions}
	g.Reset()
	return g
}

func (g *ExtensionGenerator) Reset() {
	g.base.Reset()
	g.extIdx = -1
	g.exhausted = false
}

func (g *ExtensionGenerator) Next() (string, bool) {
	if g.exhausted {
		return "", false
	}
	if g.extIdx == -1 {
		w, ok := g.base.Next()
		if !ok {
			g.exhausted = true
			return "", false
		}
		g.currentWord = w
		g.extIdx = 0
		return w, true // сначала голое слово без расширения
	}
	if g.extIdx < len(g.extensions) {
		ext := g.extensions[g.extIdx]
		g.extIdx++
		return g.currentWord + ext, true
	}
	// расширения на это слово кончились — переходим к следующему базовому
	g.extIdx = -1
	return g.Next()
}

func (g *ExtensionGenerator) Len() int {
	return g.base.Len() * (len(g.extensions) + 1) // +1 за голое слово без расширения
}

// BruteForceGenerator перебирает все строки над заданным алфавитом
// длиной от minLen до maxLen включительно — аналог "Brute forcer"
// payload-типа в Burp Intruder. Используется как заготовками "буквы
// алфавита"/"цифры и буквы", так и напрямую через spec "brute:<алфавит>:<min>:<max>".
type BruteForceGenerator struct {
	charset []rune
	minLen  int
	maxLen  int

	length int
	idx    []int // текущая комбинация индексов в charset для текущей длины
	done   bool
}

// NewBruteForceGenerator создаёт генератор по алфавиту charset для длин
// от minLen до maxLen включительно.
func NewBruteForceGenerator(charset string, minLen, maxLen int) *BruteForceGenerator {
	g := &BruteForceGenerator{charset: []rune(charset), minLen: minLen, maxLen: maxLen}
	g.Reset()
	return g
}

func (g *BruteForceGenerator) Reset() {
	g.length = g.minLen
	if g.length < 1 {
		g.length = 1
	}
	g.idx = make([]int, g.length)
	g.done = g.minLen > g.maxLen || len(g.charset) == 0
}

func (g *BruteForceGenerator) Next() (string, bool) {
	if g.done {
		return "", false
	}

	runes := make([]rune, g.length)
	for i, ci := range g.idx {
		runes[i] = g.charset[ci]
	}
	result := string(runes)

	// Одометр: продвигаем индексы для следующего вызова, перенос вправо
	// налево, как в engine.odometer для Cluster Bomb — тот же принцип,
	// просто в одном измерении вместо нескольких точек вставки.
	pos := g.length - 1
	for pos >= 0 {
		g.idx[pos]++
		if g.idx[pos] < len(g.charset) {
			break
		}
		g.idx[pos] = 0
		pos--
	}
	if pos < 0 {
		g.length++
		if g.length > g.maxLen {
			g.done = true
		} else {
			g.idx = make([]int, g.length)
		}
	}

	return result, true
}

// Len — суммарное количество комбинаций по всем длинам от minLen до maxLen.
// Растёт экспоненциально с длиной и размером алфавита — для широких
// алфавитов и maxLen больше 3-4 это может быть очень много; вызывающий
// код (например, TUI) должен относиться к этому числу как к
// предупреждению, а не просто к метрике прогресса.
func (g *BruteForceGenerator) Len() int {
	n := len(g.charset)
	total := 0
	for l := g.minLen; l <= g.maxLen; l++ {
		p := 1
		for i := 0; i < l; i++ {
			p *= n
		}
		total += p
	}
	return total
}