limit goroutine in dirbrute module

README.md

@@ -37,7 +37,7 @@ Flags:
       --host string        singel host
       --host-file string   host file
   -v, --loglevel int       level of your log (default 2)
-  -o, --output string      output file to write log and results (default "result.txt")
+  -o, --output string      output file to write log and results (default "result.log")
       --private-ip         scan private ip
       --proxy string       proxy
       --ua string          set user agent (default "user agent")

cmd/dirbrute.go

@@ -3,16 +3,17 @@ package cmd
 import (
 	"fmt"
 	"os"
-	"sync"
 
 	"github.com/cheggaaa/pb/v3"
 	"github.com/godspeedcurry/godscan/common"
 	"github.com/godspeedcurry/godscan/utils"
 	"github.com/olekukonko/tablewriter"
+	"github.com/spf13/viper"
 )
 
 type DirbruteOptions struct {
 	DirFile string
+	Threads int
 }
 
 var (
@@ -42,27 +43,43 @@ func (o *DirbruteOptions) run() {
 	}
 	utils.Info("Total: %d url(s)", len(targetUrlList))
 	utils.Info("Total: %d payload(s) in dir dict", len(targetDirList))
+	utils.Info("Total: %d threads", viper.GetInt("dirbrute-threads"))
 	utils.Success("🌲🌲🌲 Log at ./dirbrute.csv")
-	var wg sync.WaitGroup
+
 	bar := pb.StartNew(len(targetUrlList) * len(targetDirList))
 
 	table := tablewriter.NewWriter(os.Stdout)
 	table.SetAutoWrapText(false)
 
 	table.SetHeader(common.TableHeader)
 
-	for _, line := range targetUrlList {
-		for _, dir := range targetDirList {
-			wg.Add(1)
-			go func(url string, dir string) {
-				defer wg.Done()
-				ret := utils.DirBrute(url, dir)
-				utils.AddDataToTable(table, ret)
-				bar.Increment()
-			}(line, dir)
+	// 定义最大并发量
+	maxGoroutines := viper.GetInt("dirbrute-threads")
+	sem := make(chan struct{}, maxGoroutines)
+	done := make(chan bool)
+
+	go func() {
+		for _, line := range targetUrlList {
+			for _, dir := range targetDirList {
+				sem <- struct{}{} // 向通道发送信号，表示一个新的协程即将启动
+
+				go func(url string, dir string) {
+					defer func() { <-sem }() // 从通道中取出信号，表示协程结束
+
+					ret := utils.DirBrute(url, dir)
+					utils.AddDataToTable(table, ret)
+					bar.Increment()
+
+					done <- true
+				}(line, dir)
+			}
 		}
+	}()
+
+	// 等待所有任务完成
+	for i := 0; i < len(targetUrlList)*len(targetDirList); i++ {
+		<-done
 	}
-	wg.Wait()
 	bar.Finish()
 	if table.NumLines() >= 1 {
 		table.Render()
@@ -74,5 +91,10 @@ func init() {
 	dirbruteCmd := newCommandWithAliases("dirbrute", "Dirbrute on sensitive file", []string{"dir", "dirb", "dd"}, &dirbruteOptions)
 	dirbruteCmd.PersistentFlags().StringVarP(&dirbruteOptions.DirFile, "dir-file", "", "", "your directory dict")
 
+	dirbruteCmd.PersistentFlags().IntVarP(&dirbruteOptions.Threads, "threads", "t", 30, "Number of goroutine to use")
+
+	viper.BindPFlag("dirbrute-threads", dirbruteCmd.PersistentFlags().Lookup("threads"))
+	viper.SetDefault("dirbrute-threads", 30)
+
 	rootCmd.AddCommand(dirbruteCmd)
 }

cmd/port.go

@@ -53,7 +53,7 @@ func init() {
 	ipCmd.PersistentFlags().BoolVarP(&portOptions.useAllProbes, "all-probe", "a", false, "Use all probes to probe service")
 	ipCmd.PersistentFlags().BoolVarP(&portOptions.nullProbeOnly, "null-probe-only", "n", false, "Use all probes to probe service")
 
-	ipCmd.PersistentFlags().IntVarP(&portOptions.Threads, "threads", "t", 1000, "Number of threads to use, default is 1000")
+	ipCmd.PersistentFlags().IntVarP(&portOptions.Threads, "threads", "t", 1000, "Number of threads to use")
 
 	viper.BindPFlag("host", ipCmd.PersistentFlags().Lookup("host"))
 	viper.SetDefault("host", "")

cmd/root.go

@@ -111,7 +111,7 @@ func Execute() {
 	rootCmd.PersistentFlags().StringVarP(&GlobalOption.HostFile, "host-file", "", "", "host file")
 
 	rootCmd.PersistentFlags().IntVarP(&GlobalOption.LogLevel, "loglevel", "v", 2, "level of your log")
-	rootCmd.PersistentFlags().StringVarP(&GlobalOption.OutputFile, "output", "o", "result.txt", "output file to write log and results")
+	rootCmd.PersistentFlags().StringVarP(&GlobalOption.OutputFile, "output", "o", "result.log", "output file to write log and results")
 
 	rootCmd.PersistentFlags().StringVarP(&GlobalOption.DefaultUA, "ua", "", "user agent", "set user agent")
 
@@ -126,6 +126,9 @@ func Execute() {
 	viper.BindPFlag("proxy", rootCmd.PersistentFlags().Lookup("proxy"))
 	viper.SetDefault("proxy", SetProxyFromEnv())
 
+	viper.BindPFlag("output", rootCmd.PersistentFlags().Lookup("output"))
+	viper.SetDefault("output", "result.log")
+
 	if viper.GetString("proxy") != "" {
 		fmt.Fprintf(os.Stderr, "[%s] [%s] Proxy is %s\n", utils.GetCurrentTime(), color.New(color.FgCyan).Sprintf("%s", "INFO"), viper.GetString("proxy"))
 	} else {

utils/logs.go

@@ -26,7 +26,7 @@ func log_record(level int, detail string) {
 	if level > viper.GetInt("loglevel") {
 		return
 	}
-	FileWrite("result.log", detail+"\n")
+	FileWrite(viper.GetString("output"), detail+"\n")
 }
 
 func log_print(level int, detail string) {