From 4981ab326b32dea764f85117131b2b1a2374cd32 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Thu, 10 Nov 2022 04:22:42 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E8=A3=85-U,-L,=20--remove-extensions,?=
 =?UTF-8?q?=20-extensions,=20--exclude-extension?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/option.go | 66 ++++++++++++++++++++++++++++++++++++++++++++--
 internal/pool.go   |  3 +++
 internal/runner.go |  2 ++
 pkg/config.go      |  1 +
 4 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/internal/option.go b/internal/option.go
index 64c944c..e2f8ebb 100644
--- a/internal/option.go
+++ b/internal/option.go
@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"fmt"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/spray/pkg"
 	"github.com/chainreactors/words/mask"
@@ -16,8 +17,8 @@ type Option struct {
 	Dictionaries      []string `short:"d" long:"dict"`
 	Word              string   `short:"w" long:"word"`
 	Extension         string   `short:"e" long:"extensions"`
-	ExcludeExtensions bool     `long:"exclude-extensions"`
-	RemoveExtensions  bool     `long:"remove-extensions"`
+	ExcludeExtensions string   `long:"exclude-extensions"`
+	RemoveExtensions  string   `long:"remove-extensions"`
 	Uppercase         bool     `short:"U" long:"uppercase"`
 	Lowercase         bool     `short:"L" long:"lowercase"`
 
@@ -37,6 +38,10 @@ type Option struct {
 }
 
 func (opt *Option) PrepareRunner() (*Runner, error) {
+	ok := opt.Validate()
+	if !ok {
+		return nil, fmt.Errorf("validate failed")
+	}
 	var err error
 	r := &Runner{
 		Progress: uiprogress.New(),
@@ -99,18 +104,51 @@ 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 == "" {
 		for _, w := range dicts {
 			r.Wordlist = append(r.Wordlist, w...)
 		}
 	} else {
 		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)
+	}
+	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 {
+			if ext := parseExtension(s); SliceContains(rexts, ext) {
+				return strings.TrimSuffix(s, "."+ext)
+			}
+			return s
+		})
+	}
 	// prepare header
 	for _, h := range opt.Headers {
 		i := strings.Index(h, ":")
@@ -124,6 +162,14 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 	return r, nil
 }
 
+func (opt *Option) Validate() bool {
+	if opt.Uppercase && opt.Lowercase {
+		logs.Log.Error("Cannot set -U and -L at the same time")
+		return false
+	}
+	return true
+}
+
 func loadFileToSlice(filename string) ([]string, error) {
 	var ss []string
 	content, err := ioutil.ReadFile(filename)
@@ -140,3 +186,19 @@ func loadFileToSlice(filename string) ([]string, error) {
 
 	return ss, nil
 }
+
+func parseExtension(s string) string {
+	if i := strings.Index(s, "."); i != -1 {
+		return s[i+1:]
+	}
+	return ""
+}
+
+func SliceContains(s []string, e string) bool {
+	for _, v := range s {
+		if v == e {
+			return true
+		}
+	}
+	return false
+}
diff --git a/internal/pool.go b/internal/pool.go
index ef97e33..53e9720 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -190,6 +190,9 @@ Loop:
 			if !ok {
 				break Loop
 			}
+			for _, fn := range p.Fns {
+				u = fn(u)
+			}
 			p.wg.Add(1)
 			_ = p.pool.Invoke(newUnit(u, WordSource))
 		case <-ctx.Done():
diff --git a/internal/runner.go b/internal/runner.go
index d1926ba..731f260 100644
--- a/internal/runner.go
+++ b/internal/runner.go
@@ -18,6 +18,7 @@ type Runner struct {
 	URLList  []string
 	Wordlist []string
 	Headers  http.Header
+	Fns      []func(string) string
 	Threads  int
 	PoolSize int
 	Pools    *ants.PoolWithFunc
@@ -51,6 +52,7 @@ func (r *Runner) Prepare() error {
 			Timeout:  r.Timeout,
 			Headers:  r.Headers,
 			Mod:      pkg.ModMap[r.Mod],
+			Fns:      r.Fns,
 		}
 
 		if config.Mod == pkg.PathSpray {
diff --git a/pkg/config.go b/pkg/config.go
index 6d0e0be..025a023 100644
--- a/pkg/config.go
+++ b/pkg/config.go
@@ -29,4 +29,5 @@ type Config struct {
 	Headers     http.Header
 	EnableFuzzy bool
 	ClientType  int
+	Fns         []func(string) string
 }