diff --git a/go.mod b/go.mod index 59399c0..6385f03 100644 --- a/go.mod +++ b/go.mod @@ -11,7 +11,7 @@ require ( github.com/chainreactors/logs v0.0.0-20231220102821-19f082ce37c1 github.com/chainreactors/parsers v0.0.0-20231227070753-8cda94b96b6c github.com/chainreactors/utils v0.0.0-20231031063336-9477f1b23886 - github.com/chainreactors/words v0.4.1-0.20231227144711-467c1954322c + github.com/chainreactors/words v0.4.1-0.20240206181137-09f7a0bc8f50 github.com/gosuri/uiprogress v0.0.1 github.com/jessevdk/go-flags v1.5.0 github.com/panjf2000/ants/v2 v2.7.0 diff --git a/go.sum b/go.sum index 55e5baa..e591334 100644 --- a/go.sum +++ b/go.sum @@ -23,6 +23,8 @@ github.com/chainreactors/utils v0.0.0-20231031063336-9477f1b23886 h1:lS2T/uE9tg1 github.com/chainreactors/utils v0.0.0-20231031063336-9477f1b23886/go.mod h1:JA4eiQZm+7AsfjXBcIzIdVKBEhDCb16eNtWFCGTxlvs= github.com/chainreactors/words v0.4.1-0.20231227144711-467c1954322c h1:gkPx9lwkJURAnveWN2PLBBeGcQ6HbMbpAJojq2SBzGc= github.com/chainreactors/words v0.4.1-0.20231227144711-467c1954322c/go.mod h1:DUDx7PdsMEm5PvVhzkFyppzpiUhQb8dOJaWjVc1SMVk= +github.com/chainreactors/words v0.4.1-0.20240206181137-09f7a0bc8f50 h1:5aK4k6ztjvDENa4zuxYPzQH2UAy9VM+tvOhpRSCCXWQ= +github.com/chainreactors/words v0.4.1-0.20240206181137-09f7a0bc8f50/go.mod h1:DUDx7PdsMEm5PvVhzkFyppzpiUhQb8dOJaWjVc1SMVk= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= diff --git a/internal/option.go b/internal/option.go index 50bc210..77aecee 100644 --- a/internal/option.go +++ b/internal/option.go @@ -55,6 +55,7 @@ type InputOptions struct { type FunctionOptions struct { Extensions string `short:"e" long:"extension" description:"String, add extensions (separated by commas), e.g.: -e jsp,jspx"` + ForceExtension bool `long:"force-extension" description:"Bool, force add extensions"` ExcludeExtensions string `long:"exclude-extension" description:"String, exclude extensions (separated by commas), e.g.: --exclude-extension jsp,jspx"` RemoveExtensions string `long:"remove-extension" description:"String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx"` Uppercase bool `short:"U" long:"uppercase" description:"Bool, upper wordlist, e.g.: --uppercase"` @@ -62,6 +63,8 @@ type FunctionOptions struct { Prefixes []string `long:"prefix" description:"Strings, add prefix, e.g.: --prefix aaa --prefix bbb"` Suffixes []string `long:"suffix" description:"Strings, add suffix, e.g.: --suffix aaa --suffix bbb"` Replaces map[string]string `long:"replace" description:"Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"` + Skips []string `long:"skip" description:"String, skip word when generate. rule, e.g.: --skip aaa"` + //SkipEval string `long:"skip-eval" description:"String, skip word when generate. rule, e.g.: --skip-eval 'current.Length < 4'"` } type OutputOptions struct { @@ -298,8 +301,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { opt.Word = "{@prefix}" + opt.Word } - // 类似dirsearch中的 - if opt.Extensions != "" { + if opt.ForceExtension && opt.Extensions != "" { exts := strings.Split(opt.Extensions, ",") for i, e := range exts { if !strings.HasPrefix(e, ".") { @@ -475,41 +477,81 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.Tasks = tasks logs.Log.Importantf("Loaded %d urls from %s", len(tasks), taskfrom) + // 类似dirsearch中的 + if opt.Extensions != "" { + r.AppendFunction(func(s string) []string { + exts := strings.Split(opt.Extensions, ",") + ss := make([]string, len(exts)) + for i, e := range exts { + if strings.Contains(s, "%EXT%") { + ss[i] = strings.Replace(s, "%EXT%", e, -1) + } + } + return ss + }) + } else { + r.AppendFunction(func(s string) []string { + if strings.Contains(s, "%EXT%") { + return nil + } + return []string{s} + }) + } + if opt.Uppercase { - r.Fns = append(r.Fns, strings.ToUpper) + r.AppendFunction(wrapWordsFunc(strings.ToUpper)) } if opt.Lowercase { - r.Fns = append(r.Fns, strings.ToLower) + r.AppendFunction(wrapWordsFunc(strings.ToLower)) } if opt.RemoveExtensions != "" { rexts := strings.Split(opt.ExcludeExtensions, ",") - r.Fns = append(r.Fns, func(s string) string { + r.AppendFunction(func(s string) []string { if ext := parseExtension(s); iutils.StringsContains(rexts, ext) { - return strings.TrimSuffix(s, "."+ext) + return []string{strings.TrimSuffix(s, "."+ext)} } - return s + return []string{s} }) } if opt.ExcludeExtensions != "" { exexts := strings.Split(opt.ExcludeExtensions, ",") - r.Fns = append(r.Fns, func(s string) string { + r.AppendFunction(func(s string) []string { if ext := parseExtension(s); iutils.StringsContains(exexts, ext) { - return "" + return nil } - return s + return []string{s} }) } if len(opt.Replaces) > 0 { - r.Fns = append(r.Fns, func(s string) string { + r.AppendFunction(func(s string) []string { for k, v := range opt.Replaces { s = strings.Replace(s, k, v, -1) } - return s + return []string{s} }) } + + // default skip function, skip %EXT% + r.AppendFunction(func(s string) []string { + if strings.Contains(s, "%EXT%") { + return nil + } + return []string{s} + }) + if len(opt.Skips) > 0 { + r.AppendFunction(func(s string) []string { + for _, skip := range opt.Skips { + if strings.Contains(s, skip) { + return nil + } + } + return []string{s} + }) + } + logs.Log.Importantf("Loaded %d dictionaries and %d decorators", len(opt.Dictionaries), len(r.Fns)) if opt.Match != "" { diff --git a/internal/runner.go b/internal/runner.go index 5c8b5f6..d51ded0 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -42,7 +42,7 @@ type Runner struct { Rules *rule.Program AppendRules *rule.Program Headers map[string]string - Fns []func(string) string + Fns []func(string) []string FilterExpr *vm.Program MatchExpr *vm.Program RecursiveExpr *vm.Program @@ -128,6 +128,10 @@ func (r *Runner) PrepareConfig() *pkg.Config { return config } +func (r *Runner) AppendFunction(fn func(string) []string) { + r.Fns = append(r.Fns, fn) +} + func (r *Runner) Prepare(ctx context.Context) error { var err error if r.CheckOnly { diff --git a/internal/types.go b/internal/types.go index 6ff46a3..f7ebcd9 100644 --- a/internal/types.go +++ b/internal/types.go @@ -55,7 +55,7 @@ type Origin struct { sum int } -func (o *Origin) InitWorder(fns []func(string) string) (*words.Worder, error) { +func (o *Origin) InitWorder(fns []func(string) []string) (*words.Worder, error) { var worder *words.Worder wl, err := loadWordlist(o.Word, o.Dictionaries) if err != nil { diff --git a/internal/utils.go b/internal/utils.go index 53f4925..6ff9051 100644 --- a/internal/utils.go +++ b/internal/utils.go @@ -295,3 +295,9 @@ func (p *bytesPatcher) Visit(node *ast.Node) { }) } } + +func wrapWordsFunc(f func(string) string) func(string) []string { + return func(s string) []string { + return []string{f(s)} + } +} diff --git a/templates b/templates index e73a858..53643d9 160000 --- a/templates +++ b/templates @@ -1 +1 @@ -Subproject commit e73a8581c69bf15245d62276663657095b05855f +Subproject commit 53643d9011f3f932cfad07a1f5d1bc87ae88ea77