diff --git a/cmd/spray.go b/cmd/spray.go index 6fbd653..401ad69 100644 --- a/cmd/spray.go +++ b/cmd/spray.go @@ -8,8 +8,8 @@ import ( ) func main() { - var runner internal.Runner - parser := flags.NewParser(&runner, flags.Default) + var option internal.Option + 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()) diff --git a/internal/baseline.go b/internal/baseline.go index 1b0be34..d5f5901 100644 --- a/internal/baseline.go +++ b/internal/baseline.go @@ -11,11 +11,10 @@ import ( func NewBaseline(u, host string, resp *ihttp.Response) *baseline { bl := &baseline{ - //Url: u, - UrlString: u, - Host: host, - Status: resp.StatusCode(), - IsValid: true, + Url: u, + Host: host, + Status: resp.StatusCode(), + IsValid: true, } bl.Body = resp.Body() @@ -29,11 +28,10 @@ func NewBaseline(u, host string, resp *ihttp.Response) *baseline { func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline { bl := &baseline{ - //Url: u, - UrlString: u, - Host: host, - Status: resp.StatusCode(), - IsValid: false, + Url: u, + Host: host, + Status: resp.StatusCode(), + IsValid: false, } bl.RedirectURL = string(resp.GetHeader("Location")) @@ -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)) diff --git a/internal/option.go b/internal/option.go new file mode 100644 index 0000000..30c10a4 --- /dev/null +++ b/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 +} diff --git a/internal/pool.go b/internal/pool.go index ece7b4f..10f966b 100644 --- a/internal/pool.go +++ b/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(): diff --git a/internal/runner.go b/internal/runner.go index e85adaa..d1926ba 100644 --- a/internal/runner.go +++ b/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"` - OutputFile string `short:"f"` - Offset int `long:"offset"` - 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 - Debug bool `long:"debug"` - Quiet bool `short:"q" long:"quiet"` - Mod string `short:"m" long:"mod" default:"path"` - OutputCh chan *baseline - Progress *uiprogress.Progress + URLList []string + Wordlist []string + Headers http.Header + Threads int + PoolSize int + Pools *ants.PoolWithFunc + poolwg sync.WaitGroup + Timeout int + Mod string + 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 { @@ -116,13 +45,12 @@ func (r *Runner) Prepare() error { r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { u := i.(string) config := &pkg.Config{ - BaseURL: u, - Wordlist: r.Wordlist, - Thread: r.Threads, - Timeout: 2, - Headers: r.Headers, - Mod: pkg.ModMap[r.Mod], - DeadlineTime: r.Deadline, + BaseURL: u, + Wordlist: r.Wordlist, + Thread: r.Threads, + Timeout: r.Timeout, + Headers: r.Headers, + Mod: pkg.ModMap[r.Mod], } 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 } diff --git a/pkg/config.go b/pkg/config.go index 74fddb3..6d0e0be 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -19,15 +19,14 @@ var ModMap = map[string]SprayMod{ } type Config struct { - BaseURL string - Wordlist []string - Thread int - Timeout int - BaseReq *http.Request - Method string - Mod SprayMod - Headers http.Header - DeadlineTime int - EnableFuzzy bool - ClientType int + BaseURL string + Wordlist []string + Thread int + Timeout int + BaseReq *http.Request + Method string + Mod SprayMod + Headers http.Header + EnableFuzzy bool + ClientType int }