Add max file size support

README.md

@@ -63,6 +63,7 @@ INPUT:
 OUTPUT:
    -es, -estimate      estimate permutation count without generating payloads
    -o, -output string  output file to write altered subdomain list
+   -ms, -max-size int  Max export data size (kb, mb, gb, tb) (default mb)
    -v, -verbose        display verbose output
    -silent             display results only
    -version            display alterx version

cmd/alterx/main.go

@@ -19,6 +19,7 @@ func main() {
 		Payloads: cliOpts.Payloads,
 		Limit:    cliOpts.Limit,
 		Enrich:   cliOpts.Enrich, // enrich payloads
+		MaxSize: cliOpts.MaxSize,
 	}
 
 	if cliOpts.PermutationConfig != "" {

internal/runner/runner.go

@@ -3,12 +3,15 @@ package runner
 import (
 	"fmt"
 	"io"
+	"math"
 	"os"
+	"strconv"
 	"strings"
 
 	"github.com/projectdiscovery/goflags"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/gologger/levels"
+	errorutil "github.com/projectdiscovery/utils/errors"
 	fileutil "github.com/projectdiscovery/utils/file"
 	updateutils "github.com/projectdiscovery/utils/update"
 )
@@ -26,11 +29,13 @@ type Options struct {
 	Silent             bool
 	Enrich             bool
 	Limit              int
+	MaxSize            int
 	// internal/unexported fields
 	wordlists goflags.RuntimeMap
 }
 
 func ParseFlags() *Options {
+	var maxFileSize string
 	opts := &Options{}
 	flagSet := goflags.NewFlagSet()
 	flagSet.SetDescription(`Fast and customizable subdomain wordlist generator using DSL.`)
@@ -44,6 +49,7 @@ func ParseFlags() *Options {
 	flagSet.CreateGroup("output", "Output",
 		flagSet.BoolVarP(&opts.Estimate, "estimate", "es", false, "estimate permutation count without generating payloads"),
 		flagSet.StringVarP(&opts.Output, "output", "o", "", "output file to write altered subdomain list"),
+		flagSet.StringVarP(&maxFileSize, "max-size", "ms", "", "Max export data size (kb, mb, gb, tb) (default mb)"),
 		flagSet.BoolVarP(&opts.Verbose, "verbose", "v", false, "display verbose output"),
 		flagSet.BoolVar(&opts.Silent, "silent", false, "display results only"),
 		flagSet.CallbackVar(printVersion, "version", "display alterx version"),
@@ -89,6 +95,15 @@ func ParseFlags() *Options {
 		}
 	}
 
+	opts.MaxSize = math.MaxInt
+	if len(maxFileSize) > 0 {
+		maxSize, err := convertFileSizeToBytes(maxFileSize)
+		if err != nil {
+			gologger.Fatal().Msgf("Could not parse max-size: %s\n", err)
+		}
+		opts.MaxSize = maxSize
+	}
+
 	opts.Payloads = map[string][]string{}
 	for k, v := range opts.wordlists.AsMap() {
 		value, ok := v.(string)
@@ -129,3 +144,32 @@ func printVersion() {
 	gologger.Info().Msgf("Current version: %s", version)
 	os.Exit(0)
 }
+
+func convertFileSizeToBytes(maxFileSize string) (int, error) {
+	maxFileSize = strings.ToLower(maxFileSize)
+	// default to mb
+	if size, err := strconv.Atoi(maxFileSize); err == nil {
+		return size * 1024 * 1024, nil
+	}
+	if len(maxFileSize) < 3 {
+		return 0, errorutil.New("invalid max-size value")
+	}
+	sizeUnit := maxFileSize[len(maxFileSize)-2:]
+	size, err := strconv.Atoi(maxFileSize[:len(maxFileSize)-2])
+	if err != nil {
+		return 0, err
+	}
+	if size < 0 {
+		return 0, errorutil.New("max-size cannot be negative")
+	}
+	if strings.EqualFold(sizeUnit, "kb") {
+		return size * 1024, nil
+	} else if strings.EqualFold(sizeUnit, "mb") {
+		return size * 1024 * 1024, nil
+	} else if strings.EqualFold(sizeUnit, "gb") {
+		return size * 1024 * 1024 * 1024, nil
+	} else if strings.EqualFold(sizeUnit, "tb") {
+		return size * 1024 * 1024 * 1024 * 1024, nil
+	}
+	return 0, errorutil.New("Unsupported max-size unit")
+}

mutator.go

@@ -37,6 +37,8 @@ type Options struct {
 	// Enrich when true alterx extra possible words from input
 	// and adds them to default payloads word,number
 	Enrich bool
+	// MaxSize limits output data size
+	MaxSize int
 }
 
 // Mutator
@@ -138,6 +140,7 @@ func (m *Mutator) ExecuteWithWriter(Writer io.Writer) error {
 	}
 	resChan := m.Execute(context.TODO())
 	m.payloadCount = 0
+	maxFileSize := m.Options.MaxSize
 	for {
 		value, ok := <-resChan
 		if !ok {
@@ -148,11 +151,23 @@ func (m *Mutator) ExecuteWithWriter(Writer io.Writer) error {
 			// we can't early exit, due to abstraction we have to conclude the elaboration to drain all dedupers
 			continue
 		}
-		_, err := Writer.Write([]byte(value + "\n"))
-		m.payloadCount++
+		if maxFileSize <= 0 {
+			// drain all dedupers when max-file size reached
+			continue
+		}
+		outputData := []byte(value + "\n")
+		if len(outputData) > maxFileSize {
+			maxFileSize = 0
+			continue
+		}
+
+		n, err := Writer.Write(outputData)
 		if err != nil {
 			return err
 		}
+		// update maxFileSize limit after each write
+		maxFileSize -= n
+		m.payloadCount++
 	}
 }