新增--suffix,--prefix, --replace

internal/option.go

@@ -8,6 +8,7 @@ import (
 	"github.com/gosuri/uiprogress"
 	"io/ioutil"
 	"os"
+	"strconv"
 	"strings"
 )
 
@@ -16,12 +17,14 @@ type Option struct {
 	URLFile           string            `short:"l" long:"list"`
 	Dictionaries      []string          `short:"d" long:"dict"`
 	Word              string            `short:"w" long:"word"`
-	Extension         string   `short:"e" long:"extensions"`
-	ExcludeExtensions string   `long:"exclude-extensions"`
-	RemoveExtensions  string   `long:"remove-extensions"`
+	Extensions        string            `short:"e" long:"extension"`
+	ExcludeExtensions string            `long:"exclude-extension"`
+	RemoveExtensions  string            `long:"remove-extension"`
 	Uppercase         bool              `short:"U" long:"uppercase"`
 	Lowercase         bool              `short:"L" long:"lowercase"`
-
+	Prefixes          []string          `long:"prefix"`
+	Suffixes          []string          `long:"suffix"`
+	Replaces          map[string]string `long:"replace"`
 	Deadline          int               `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
 	Timeout           int               `long:"timeout" default:"2"`
 	Headers           []string          `long:"header"`
@@ -104,26 +107,33 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		logs.Log.Importantf("load %d word from %s", len(dicts[i]), f)
 	}
 
-	var exts []string
-	if opt.Extension != "" {
-		exts = strings.Split(opt.Extension, ",")
+	if opt.Word == "" {
+		opt.Word = "{?"
+		for i, _ := range dicts {
+			opt.Word += strconv.Itoa(i)
+		}
+		opt.Word = "}"
 	}
 
-	if opt.Word == "" {
-		for _, w := range dicts {
-			r.Wordlist = append(r.Wordlist, w...)
+	if opt.Suffixes == nil {
+		dicts = append(dicts, opt.Suffixes)
+		opt.Word += fmt.Sprintf("{?%d}", len(dicts)-1)
 	}
-	} else {
+	if opt.Prefixes != nil {
+		dicts = append(dicts, opt.Prefixes)
+		opt.Word = fmt.Sprintf("{?%d}", len(dicts)-1) + opt.Word
+	}
+
+	if opt.Extensions != "" {
+		dicts = append(dicts, strings.Split(opt.Extensions, ","))
+		opt.Word += fmt.Sprintf("{?%d}", len(dicts)-1)
+	}
+
 	mask.CustomWords = dicts
-		if len(exts) > 0 {
-			mask.CustomWords = append(mask.CustomWords, exts)
-		}
-		opt.Word += fmt.Sprintf("{?%d}", len(exts)-1)
 	r.Wordlist, err = mask.Run(opt.Word)
 	if err != nil {
 		return nil, err
 	}
-	}
 
 	if opt.Uppercase {
 		r.Fns = append(r.Fns, strings.ToUpper)
@@ -131,15 +141,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 	if opt.Lowercase {
 		r.Fns = append(r.Fns, strings.ToLower)
 	}
-	if opt.ExcludeExtensions != "" {
-		exexts := strings.Split(opt.ExcludeExtensions, ",")
-		r.Fns = append(r.Fns, func(s string) string {
-			if ext := parseExtension(s); SliceContains(exexts, ext) {
-				return ""
-			}
-			return s
-		})
-	}
+
 	if opt.RemoveExtensions != "" {
 		rexts := strings.Split(opt.ExcludeExtensions, ",")
 		r.Fns = append(r.Fns, func(s string) string {
@@ -149,6 +151,26 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 			return s
 		})
 	}
+
+	if opt.ExcludeExtensions != "" {
+		exexts := strings.Split(opt.ExcludeExtensions, ",")
+		r.Fns = append(r.Fns, func(s string) string {
+			if ext := parseExtension(s); SliceContains(exexts, ext) {
+				return ""
+			}
+			return s
+		})
+	}
+
+	if opt.Replaces != nil {
+		r.Fns = append(r.Fns, func(s string) string {
+			for k, v := range opt.Replaces {
+				s = strings.Replace(s, k, v, -1)
+			}
+			return s
+		})
+	}
+
 	// prepare header
 	for _, h := range opt.Headers {
 		i := strings.Index(h, ":")

internal/pool.go

@@ -193,6 +193,9 @@ Loop:
 			for _, fn := range p.Fns {
 				u = fn(u)
 			}
+			if u == "" {
+				continue
+			}
 			p.wg.Add(1)
 			_ = p.pool.Invoke(newUnit(u, WordSource))
 		case <-ctx.Done():