新增--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,33 +8,36 @@ import (
"github.com/gosuri/uiprogress" "github.com/gosuri/uiprogress"
"io/ioutil" "io/ioutil"
"os" "os"
"strconv"
"strings" "strings"
) )
type Option struct { type Option struct {
URL string `short:"u" long:"url"` URL string `short:"u" long:"url"`
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"`
Deadline int `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline Suffixes []string `long:"suffix"`
Timeout int `long:"timeout" default:"2"` Replaces map[string]string `long:"replace"`
Headers []string `long:"header"` Deadline int `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
OutputFile string `short:"f"` Timeout int `long:"timeout" default:"2"`
OutputProbe string `long:"probe"` Headers []string `long:"header"`
Offset int `long:"offset"` OutputFile string `short:"f"`
Limit int `long:"limit"` OutputProbe string `long:"probe"`
Threads int `short:"t" long:"thread" default:"20"` Offset int `long:"offset"`
PoolSize int `short:"p" long:"pool" default:"5"` Limit int `long:"limit"`
Debug bool `long:"debug"` Threads int `short:"t" long:"thread" default:"20"`
Quiet bool `short:"q" long:"quiet"` PoolSize int `short:"p" long:"pool" default:"5"`
Mod string `short:"m" long:"mod" default:"path"` Debug bool `long:"debug"`
Client string `short:"c" long:"client" default:"auto"` Quiet bool `short:"q" long:"quiet"`
Mod string `short:"m" long:"mod" default:"path"`
Client string `short:"c" long:"client" default:"auto"`
} }
func (opt *Option) PrepareRunner() (*Runner, error) { func (opt *Option) PrepareRunner() (*Runner, error) {
@ -104,25 +107,32 @@ 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 {
mask.CustomWords = dicts dicts = append(dicts, opt.Prefixes)
if len(exts) > 0 { opt.Word = fmt.Sprintf("{?%d}", len(dicts)-1) + opt.Word
mask.CustomWords = append(mask.CustomWords, exts) }
}
opt.Word += fmt.Sprintf("{?%d}", len(exts)-1) if opt.Extensions != "" {
r.Wordlist, err = mask.Run(opt.Word) dicts = append(dicts, strings.Split(opt.Extensions, ","))
if err != nil { opt.Word += fmt.Sprintf("{?%d}", len(dicts)-1)
return nil, err }
}
mask.CustomWords = dicts
r.Wordlist, err = mask.Run(opt.Word)
if err != nil {
return nil, err
} }
if opt.Uppercase { if opt.Uppercase {
@ -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():