Use RAM-friendly concurrent algorithm (#6)

* Use RAM-friendly concurrent algorithm; refactor

* Add -t to readme

Co-authored-by: Mohammed Diaa <[email protected]>

Author: PolovinaD

README.md

@@ -43,6 +43,8 @@ go install github.com/trickest/mkpath@latest
     	Generate files only, file names are appended to given domains (default false)
   -r string
     	Regex to filter words from wordlist file
+  -t int
+    	Number of threads for every path depth (default 100)
   -w string
     	Wordlist file
 ```

mkpath.go

@@ -2,126 +2,114 @@ package main
 
 import (
 	"bufio"
+	"bytes"
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
 	"regexp"
+	"strconv"
 	"strings"
+	"sync"
+	"syscall"
+
+	roundChan "github.com/trickest/mkpath/round"
 )
 
 const (
 	fileRegex = "[^?*:;{}]+\\.[^/?*:;{}]+"
+
+	bufferSizeMB      = 100
+	maxWorkingThreads = 100000
+	numberOfFiles     = 1
 )
 
-func generatePaths(wordSet map[string]bool, depth int) []string {
-	results := make([]string, 0)
-	for i := 0; i < depth; i += 1 {
-		toMerge := results[0:]
-		if len(toMerge) == 0 {
-			for word := range wordSet {
-				results = append(results, word)
-			}
-		} else {
-			for _, sd := range toMerge {
-				for word := range wordSet {
-					results = append(results, fmt.Sprintf("%s/%s", word, sd))
-				}
-			}
-		}
-	}
-	return results
-}
+var (
+	domain       string
+	inputDomains []string
+	domainFile   string
 
-func generateFiles(paths []string, files map[string]bool) []string {
-	results := make([]string, 0)
-	for _, path := range paths {
-		for file := range files {
-			results = append(results, path+"/"+file)
-		}
-	}
-	return results
-}
+	wordlist    string
+	toLowercase bool
+	regex       string
+	dirWordSet  map[string]bool
+	fileWordSet map[string]bool
 
-func generateAll(paths map[string]bool, files map[string]bool, depth int) []string {
-	var results []string
-	results = generatePaths(paths, depth)
-	results = append(results, generateFiles(generatePaths(paths, depth), files)...)
-	return results
-}
+	depth          int
+	onlyDirs       bool
+	onlyFiles      bool
+	outputFileName string
+	silent         bool
 
-func main() {
-	domain := flag.String("d", "", "Input domain")
-	domainFile := flag.String("df", "", "Input domain file, one domain per line")
-	wordlist := flag.String("w", "", "Wordlist file")
-	toLowercase := flag.Bool("lower", false, "Convert wordlist file content to lowercase (default false)")
-	r := flag.String("r", "", "Regex to filter words from wordlist file")
-	depth := flag.Int("l", 1, "URL path depth to generate (default 1)")
-	output := flag.String("o", "", "Output file (optional)")
-	onlyDirs := flag.Bool("only-dirs", false, "Generate directories only, files are filtered out (default false)")
-	onlyFiles := flag.Bool("only-files", false, "Generate files only, file names are appended to given domains (default false)")
-	flag.Parse()
+	workers         int
+	workerThreadMax = make(chan struct{}, maxWorkingThreads)
+	done            = make(chan struct{})
+	wg              sync.WaitGroup
+	wgWrite         sync.WaitGroup
+	robin           roundChan.RoundRobin
+)
 
-	inputDomains := make([]string, 0)
-	if *domain != "" {
-		inputDomains = append(inputDomains, *domain)
+func readDomainFile() {
+	inputFile, err := os.Open(domainFile)
+	if err != nil {
+		fmt.Println("Could not open file to read domains:", err)
+		os.Exit(1)
 	}
-	if *domainFile != "" {
-		inputFile, err := os.Open(*domainFile)
-		if err != nil {
-			fmt.Println(err.Error())
-			os.Exit(1)
-		}
-		defer inputFile.Close()
-		scanner := bufio.NewScanner(inputFile)
-		for scanner.Scan() {
-			inputDomains = append(inputDomains, scanner.Text())
-		}
+	defer inputFile.Close()
+
+	scanner := bufio.NewScanner(inputFile)
+	for scanner.Scan() {
+		inputDomains = append(inputDomains, strings.TrimSpace(scanner.Text()))
 	}
-	if len(inputDomains) == 0 {
-		fmt.Println("No input provided")
+}
+
+func prepareDomains() {
+	if domain == "" && domainFile == "" {
+		fmt.Println("No domain input provided!")
 		os.Exit(1)
 	}
 
-	wordlistFile, err := os.Open(*wordlist)
-	if err != nil {
-		fmt.Println(err.Error())
-		os.Exit(1)
+	inputDomains = make([]string, 0)
+	if domain != "" {
+		inputDomains = append(inputDomains, domain)
+	} else {
+		if domainFile != "" {
+			readDomainFile()
+		}
 	}
-	defer wordlistFile.Close()
+}
 
+func readWordlistFile() {
 	var reg *regexp.Regexp
-	if *r != "" {
-		reg, err = regexp.Compile(*r)
+	var err error
+	if regex != "" {
+		reg, err = regexp.Compile(regex)
 		if err != nil {
-			fmt.Println(err.Error())
+			fmt.Println(err)
 			os.Exit(1)
 		}
 	}
 
-	var fileReg *regexp.Regexp
-	fileReg, err = regexp.Compile(fileRegex)
+	wordlistFile, err := os.Open(wordlist)
 	if err != nil {
-		fmt.Println(err.Error())
+		fmt.Println("Could not open file to read wordlist:", err)
 		os.Exit(1)
 	}
+	defer wordlistFile.Close()
 
-	var outputFile *os.File
-	if *output != "" {
-		outputFile, err = os.Create(*output)
-		if err != nil {
-			fmt.Println(err.Error())
-			os.Exit(1)
-		}
-		defer outputFile.Close()
+	fileReg, err := regexp.Compile(fileRegex)
+	if err != nil {
+		fmt.Println(err)
+		os.Exit(1)
 	}
 
-	dirWordSet := make(map[string]bool)
-	fileWordSet := make(map[string]bool)
+	dirWordSet = make(map[string]bool)
+	fileWordSet = make(map[string]bool)
 	scanner := bufio.NewScanner(wordlistFile)
 
 	for scanner.Scan() {
 		word := scanner.Text()
-		if *toLowercase {
+		if toLowercase {
 			word = strings.ToLower(word)
 		}
 		word = strings.Trim(word, "/")
@@ -138,24 +126,173 @@ func main() {
 			}
 		}
 	}
+}
+
+func closeWriters(number int) {
+	for i := 0; i < number; i++ {
+		done <- struct{}{}
+	}
+}
+
+func spawnWriters(number int) {
+	for i := 0; i < number; i++ {
+		var bf bytes.Buffer
+		ch := make(chan string, 100000)
+
+		fileName := outputFileName
+		fileSplit := strings.Split(fileName, ".")
+		if len(fileSplit) == 1 {
+			fileName += ".txt"
+		}
+		if number > 1 {
+			fileSplit = strings.Split(fileName, ".")
+			extension := "." + fileSplit[len(fileSplit)-1]
+			fileName = strings.TrimSuffix(fileName, extension) + "-" + strconv.Itoa(i) + extension
+		}
+		file, err := os.Create(fileName)
+		if err != nil {
+			fmt.Println("Couldn't open file to write output:", err)
+			os.Exit(1)
+		}
+
+		wgWrite.Add(1)
+		go write(file, &bf, &ch)
+
+		if robin == nil {
+			robin = roundChan.New(&ch)
+			continue
+		}
+		robin.Add(&ch)
+	}
+}
+
+func write(file *os.File, buffer *bytes.Buffer, ch *chan string) {
+mainLoop:
+	for {
+		select {
+		case <-done:
+			for {
+				if !writeOut(file, buffer, ch) {
+					break
+				}
+			}
+			if buffer.Len() > 0 {
+				if file != nil {
+					_, _ = file.WriteString(buffer.String())
+					buffer.Reset()
+				}
+			}
+			break mainLoop
+		default:
+			writeOut(file, buffer, ch)
+		}
+	}
+	wgWrite.Done()
+}
+
+func writeOut(file *os.File, buffer *bytes.Buffer, outputChannel *chan string) bool {
+	select {
+	case s := <-*outputChannel:
+		buffer.WriteString(s)
+		if buffer.Len() >= bufferSizeMB*1024*1024 {
+			_, _ = file.WriteString(buffer.String())
+			buffer.Reset()
+		}
+		return true
+	default:
+		return false
+	}
+}
+
+func combo(_comb string, level int, wg *sync.WaitGroup, wt *chan struct{}) {
+	defer wg.Done()
+	workerThreadMax <- struct{}{}
+
+	if strings.Count(_comb, "/") > 0 {
+		processOutput(_comb, robin.Next())
+	}
 
-	var results []string
-	if *onlyDirs != *onlyFiles {
-		if *onlyDirs {
-			results = generatePaths(dirWordSet, *depth)
-		} else {
-			results = generateFiles(results, fileWordSet)
+	var nextLevelWaitGroup sync.WaitGroup
+	if level > 1 {
+		nextLevelWt := make(chan struct{}, workers)
+		for dw := range dirWordSet {
+			nextLevelWaitGroup.Add(1)
+			nextLevelWt <- struct{}{}
+			go combo(_comb+"/"+dw, level-1, &nextLevelWaitGroup, &nextLevelWt)
 		}
 	} else {
-		results = generateAll(dirWordSet, fileWordSet, *depth)
+		for dw := range dirWordSet {
+			processOutput(_comb+"/"+dw, robin.Next())
+		}
 	}
 
-	for _, domain := range inputDomains {
-		for _, subpath := range results {
-			fmt.Println(domain + "/" + subpath)
-			if outputFile != nil {
-				_, _ = outputFile.WriteString(domain + "/" + subpath + "\n")
+	nextLevelWaitGroup.Wait()
+	<-workerThreadMax
+	<-*wt
+}
+
+func processOutput(out string, outChan *chan string) {
+	if onlyDirs || onlyFiles == onlyDirs {
+		if !silent {
+			fmt.Print(out + "\n")
+		}
+		*outChan <- out + "\n"
+	}
+
+	if onlyFiles || onlyFiles == onlyDirs {
+		for file := range fileWordSet {
+			if !silent {
+				fmt.Print(out + "/" + file + "\n")
 			}
+			*outChan <- out + "/" + file + "\n"
 		}
 	}
 }
+
+func main() {
+	flag.StringVar(&domain, "d", "", "Input domain")
+	flag.StringVar(&domainFile, "df", "", "Input domain file, one domain per line")
+	flag.StringVar(&wordlist, "w", "", "Wordlist file")
+	flag.BoolVar(&toLowercase, "lower", false, "Convert wordlist file content to lowercase (default false)")
+	flag.StringVar(&regex, "r", "", "Regex to filter words from wordlist file")
+	flag.IntVar(&depth, "l", 1, "URL path depth to generate")
+	flag.StringVar(&outputFileName, "o", "", "Output file (optional)")
+	flag.BoolVar(&onlyDirs, "only-dirs", false, "Generate directories only, files are filtered out (default false)")
+	flag.BoolVar(&onlyFiles, "only-files", false, "Generate files only, file names are appended to given domains (default false)")
+	flag.IntVar(&workers, "t", 100, "Number of threads for every path depth")
+	flag.BoolVar(&silent, "silent", true, "Skip writing generated paths to stdout (faster)")
+	flag.Parse()
+
+	go func() {
+		signalChannel := make(chan os.Signal, 1)
+		signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM, syscall.SIGKILL)
+		<-signalChannel
+
+		fmt.Println("Program interrupted, exiting...")
+		os.Exit(0)
+	}()
+
+	if depth <= 0 || workers <= 0 {
+		fmt.Println("Path depth and number of threads must be positive integers!")
+		os.Exit(0)
+	}
+
+	prepareDomains()
+	readWordlistFile()
+	spawnWriters(numberOfFiles)
+
+	if outputFileName == "" {
+		silent = false
+	}
+
+	for _, d := range inputDomains {
+		wg.Add(1)
+		wt := make(chan struct{}, 1)
+		wt <- struct{}{}
+		go combo(d, depth, &wg, &wt)
+	}
+
+	wg.Wait()
+	closeWriters(numberOfFiles)
+	wgWrite.Wait()
+}