添加option文件, 从runner中分离不必要的功能

cmd/spray.go

@@ -8,8 +8,8 @@ import (
 )
 
 func main() {
-	var runner internal.Runner
+	var option internal.Option
-	parser := flags.NewParser(&runner, flags.Default)
+	parser := flags.NewParser(&option, flags.Default)
 	_, err := parser.Parse()
 	if err != nil {
 		if err.(*flags.Error).Type != flags.ErrHelp {
@@ -18,6 +18,12 @@ func main() {
 		return
 	}
 
+	runner, err := option.PrepareRunner()
+	if err != nil {
+		logs.Log.Errorf(err.Error())
+		return
+	}
+
 	err = runner.Prepare()
 	if err != nil {
 		logs.Log.Errorf(err.Error())

internal/baseline.go

@@ -11,8 +11,7 @@ import (
 
 func NewBaseline(u, host string, resp *ihttp.Response) *baseline {
 	bl := &baseline{
-		//Url:       u,
+		Url:     u,
-		UrlString: u,
 		Host:    host,
 		Status:  resp.StatusCode(),
 		IsValid: true,
@@ -29,8 +28,7 @@ func NewBaseline(u, host string, resp *ihttp.Response) *baseline {
 
 func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
 	bl := &baseline{
-		//Url:       u,
+		Url:     u,
-		UrlString: u,
 		Host:    host,
 		Status:  resp.StatusCode(),
 		IsValid: false,
@@ -42,7 +40,7 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
 }
 
 type baseline struct {
-	UrlString    string         `json:"url"`
+	Url          string         `json:"url"`
 	Host         string         `json:"host"`
 	Body         []byte         `json:"-"`
 	BodyLength   int            `json:"body_length"`
@@ -98,7 +96,7 @@ func (bl *baseline) FuzzyEqual(other *baseline) bool {
 func (bl *baseline) String() string {
 	var line strings.Builder
 	//line.WriteString("[+] ")
-	line.WriteString(bl.UrlString)
+	line.WriteString(bl.Url)
 	line.WriteString(" (" + bl.Host + ")")
 	line.WriteString(" - ")
 	line.WriteString(strconv.Itoa(bl.Status))

internal/option.go

@@ -0,0 +1,127 @@
+package internal
+
+import (
+	"github.com/chainreactors/logs"
+	"github.com/chainreactors/spray/pkg"
+	"github.com/gosuri/uiprogress"
+	"io/ioutil"
+	"os"
+	"strings"
+)
+
+type Option struct {
+	URL               string   `short:"u" long:"url"`
+	URLFile           string   `short:"l" long:"list"`
+	WordLists         []string `short:"w" long:"word"`
+	Extension         string   `short:"e" long:"extensions"`
+	ExcludeExtensions bool     `long:"exclude-extensions"`
+	RemoveExtensions  bool     `long:"remove-extensions"`
+	Uppercase         bool     `short:"U" long:"uppercase"`
+	Lowercase         bool     `short:"L" long:"lowercase"`
+
+	Deadline    int      `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
+	Timeout     int      `long:"timeout" default:"2"`
+	Headers     []string `long:"header"`
+	OutputFile  string   `short:"f"`
+	OutputProbe string   `long:"probe"`
+	Offset      int      `long:"offset"`
+	Limit       int      `long:"limit"`
+	Threads     int      `short:"t" long:"thread" default:"20"`
+	PoolSize    int      `short:"p" long:"pool" default:"5"`
+	Debug       bool     `long:"debug"`
+	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) {
+	r := &Runner{
+		Progress: uiprogress.New(),
+		Threads:  opt.Threads,
+		PoolSize: opt.PoolSize,
+		Mod:      opt.Mod,
+		Timeout:  opt.Timeout,
+	}
+
+	if opt.Debug {
+		logs.Log.Level = logs.Debug
+	}
+	if !opt.Quiet {
+		r.Progress.Start()
+		logs.Log.Writer = r.Progress.Bypass()
+	}
+
+	// prepare url
+	var file *os.File
+	var err error
+	urlfrom := opt.URLFile
+	if opt.URL != "" {
+		r.URLList = append(r.URLList, opt.URL)
+		urlfrom = "cmd"
+	} else if opt.URLFile != "" {
+		file, err = os.Open(opt.URLFile)
+		if err != nil {
+			return nil, err
+		}
+	} else if pkg.HasStdin() {
+		file = os.Stdin
+		urlfrom = "stdin"
+	}
+
+	if file != nil {
+		content, err := ioutil.ReadAll(file)
+		if err != nil {
+			return nil, err
+		}
+		r.URLList = strings.Split(string(content), "\n")
+	}
+
+	for i, u := range r.URLList {
+		r.URLList[i] = strings.TrimSpace(u)
+	}
+	logs.Log.Importantf("load %d urls from %s", len(r.URLList), urlfrom)
+
+	// prepare word
+	words := make([][]string, len(opt.WordLists))
+	for i, f := range opt.WordLists {
+		words[i], err = loadFileToSlice(f)
+		if err != nil {
+			return nil, err
+		}
+		logs.Log.Importantf("load %d word from %s", len(r.Wordlist), f)
+	}
+
+	for _, w := range words {
+		r.Wordlist = append(r.Wordlist, w...)
+	}
+	// todo mask
+
+	// prepare header
+	for _, h := range opt.Headers {
+		i := strings.Index(h, ":")
+		if i == -1 {
+			logs.Log.Warn("invalid header")
+		} else {
+			r.Headers.Add(h[:i], h[i+2:])
+		}
+	}
+
+	return r, nil
+}
+
+func loadFileToSlice(filename string) ([]string, error) {
+	var ss []string
+	content, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return nil, err
+	}
+
+	ss = strings.Split(string(content), "\n")
+
+	// 统一windows与linux的回车换行差异
+	for i, word := range ss {
+		ss[i] = strings.TrimSpace(word)
+	}
+
+	return ss, nil
+}

internal/pool.go

@@ -84,7 +84,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 		if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge {
 			pool.failedCount++
-			bl = &baseline{UrlString: pool.BaseURL + unit.path, Err: reqerr}
+			bl = &baseline{Url: pool.BaseURL + unit.path, Err: reqerr}
 		} else {
 			pool.failedCount = 0
 			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
@@ -192,8 +192,6 @@ Loop:
 			}
 			p.wg.Add(1)
 			_ = p.pool.Invoke(newUnit(u, WordSource))
-		case <-time.NewTimer(time.Duration(p.DeadlineTime) * time.Second).C:
-			break Loop
 		case <-ctx.Done():
 			break Loop
 		case <-p.ctx.Done():

internal/runner.go

@@ -2,16 +2,12 @@ package internal
 
 import (
 	"context"
-	"fmt"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/spray/pkg"
 	"github.com/chainreactors/spray/pkg/ihttp"
 	"github.com/gosuri/uiprogress"
 	"github.com/panjf2000/ants/v2"
-	"io/ioutil"
 	"net/http"
-	"os"
-	"strings"
 	"sync"
 )
 
@@ -19,88 +15,21 @@ var BlackStatus = []int{400, 404, 410}
 var FuzzyStatus = []int{403, 500, 501, 502, 503}
 
 type Runner struct {
-	URL        string `short:"u" long:"url"`
-	URLFile    string `short:"l" long:"list"`
 	URLList  []string
-	WordFile   string `short:"w" long:"work"`
 	Wordlist []string
-	Headers    http.Header `long:"header"`
+	Headers  http.Header
-	OutputFile string      `short:"f"`
+	Threads  int
-	Offset     int         `long:"offset"`
+	PoolSize int
-	Limit      int         `long:"limit"`
-	Threads    int         `short:"t" long:"thread" default:"20"`
-	PoolSize   int         `short:"p" long:"pool" default:"5"`
 	Pools    *ants.PoolWithFunc
 	poolwg   sync.WaitGroup
-	Deadline   int    `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
+	Timeout  int
-	Debug      bool   `long:"debug"`
+	Mod      string
-	Quiet      bool   `short:"q" long:"quiet"`
-	Mod        string `short:"m" long:"mod" default:"path"`
 	OutputCh chan *baseline
 	Progress *uiprogress.Progress
 }
 
 func (r *Runner) Prepare() error {
-	r.Progress = uiprogress.New()
-
-	if r.Debug {
-		logs.Log.Level = logs.Debug
-	}
-	if !r.Quiet {
-		r.Progress.Start()
-		logs.Log.Writer = r.Progress.Bypass()
-	}
-
-	var file *os.File
 	var err error
-	urlfrom := r.URLFile
-	if r.URL != "" {
-		r.URLList = append(r.URLList, r.URL)
-		urlfrom = "cmd"
-	} else if r.URLFile != "" {
-		file, err = os.Open(r.URLFile)
-		if err != nil {
-			return err
-		}
-	} else if pkg.HasStdin() {
-		file = os.Stdin
-		urlfrom = "stdin"
-	}
-
-	if file != nil {
-		content, err := ioutil.ReadAll(file)
-		if err != nil {
-			return err
-		}
-		r.URLList = strings.Split(string(content), "\n")
-	}
-
-	// todo url formatter
-	for i, u := range r.URLList {
-		r.URLList[i] = strings.TrimSpace(u)
-	}
-	logs.Log.Importantf("load %d urls from %s", len(r.URLList), urlfrom)
-
-	if r.WordFile != "" {
-		content, err := ioutil.ReadFile(r.WordFile)
-		if err != nil {
-			return err
-		}
-		r.Wordlist = strings.Split(string(content), "\n")
-	} else {
-		return fmt.Errorf("not special wordlist")
-	}
-
-	if r.Wordlist != nil && len(r.Wordlist) > 0 {
-		// todo  suffix/prefix/trim/generator
-		for i, word := range r.Wordlist {
-			r.Wordlist[i] = strings.TrimSpace(word)
-		}
-		logs.Log.Importantf("load %d word from %s", len(r.Wordlist), r.WordFile)
-	} else {
-		return fmt.Errorf("no wordlist")
-	}
-
 	CheckStatusCode = func(status int) bool {
 		for _, black := range BlackStatus {
 			if black == status {
@@ -119,10 +48,9 @@ func (r *Runner) Prepare() error {
 			BaseURL:  u,
 			Wordlist: r.Wordlist,
 			Thread:   r.Threads,
-			Timeout:      2,
+			Timeout:  r.Timeout,
 			Headers:  r.Headers,
 			Mod:      pkg.ModMap[r.Mod],
-			DeadlineTime: r.Deadline,
 		}
 
 		if config.Mod == pkg.PathSpray {
@@ -146,6 +74,10 @@ func (r *Runner) Prepare() error {
 		pool.Run(ctx)
 		r.poolwg.Done()
 	})
+
+	if err != nil {
+		return err
+	}
 	go r.Outputting()
 	return nil
 }

pkg/config.go

@@ -27,7 +27,6 @@ type Config struct {
 	Method      string
 	Mod         SprayMod
 	Headers     http.Header
-	DeadlineTime int
 	EnableFuzzy bool
 	ClientType  int
 }