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

This commit is contained in:
M09Ic 2022-11-10 04:48:07 +08:00
parent 4981ab326b
commit 78b3dcda75
2 changed files with 74 additions and 49 deletions

View File

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

View File

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