From 341998b26b3fc1398f5086d78da64dfbfdd2ba59 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Fri, 2 Dec 2022 19:59:15 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E6=96=AD=E7=82=B9=E7=BB=AD?= =?UTF-8?q?=E4=BC=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/cmd.go | 7 ++- internal/option.go | 119 ++++++++++++++++++++++++++++----------------- internal/pool.go | 3 +- internal/runner.go | 45 +++++++++-------- pkg/statistor.go | 24 ++++++++- 5 files changed, 126 insertions(+), 72 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index c453640..f86f604 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -19,8 +19,13 @@ func Spray() { } return } + var runner *internal.Runner + if option.ResumeFrom != "" { + runner, err = option.PrepareRunner() + } else { + runner, err = option.PrepareRunner() + } - runner, err := option.PrepareRunner() if err != nil { logs.Log.Errorf(err.Error()) return diff --git a/internal/option.go b/internal/option.go index e84a3f7..d930b68 100644 --- a/internal/option.go +++ b/internal/option.go @@ -25,6 +25,7 @@ type Option struct { } type InputOptions struct { + ResumeFrom string `short:"r" long:"resume-from"` URL string `short:"u" long:"url" description:"String, input baseurl (separated by commas), e.g.: http://google.com, http://baidu.com"` URLFile string `short:"l" long:"list" description:"File, input filename"` Offset int `long:"offset" description:"Int, wordlist offset"` @@ -76,6 +77,7 @@ type MiscOptions struct { Threads int `short:"t" long:"thread" default:"20" description:"Int, number of threads per pool (seconds)"` Debug bool `long:"debug" description:"Bool, output debug info"` Quiet bool `short:"q" long:"quiet" description:"Bool, Quiet"` + NoBar bool `long:"no-bar"` Mod string `short:"m" long:"mod" default:"path" choice:"path" choice:"host" description:"String, path/host spray"` Client string `short:"c" long:"client" default:"auto" choice:"fast" choice:"standard" choice:"auto" description:"String, Client type"` } @@ -94,8 +96,8 @@ func (opt *Option) PrepareRunner() (*Runner, error) { Timeout: opt.Timeout, Deadline: opt.Deadline, Offset: opt.Offset, - Limit: opt.Limit, - urlCh: make(chan string), + Total: opt.Limit, + taskCh: make(chan *Task), OutputCh: make(chan *pkg.Baseline, 100), FuzzyCh: make(chan *pkg.Baseline, 100), Fuzzy: opt.Fuzzy, @@ -124,13 +126,13 @@ func (opt *Option) PrepareRunner() (*Runner, error) { if opt.Debug { logs.Log.Level = logs.Debug } - if !opt.Quiet { - r.Progress.Start() - logs.Log.Writer = r.Progress.Bypass() - } else { + if opt.Quiet { logs.Log.Quiet = true } - + if opt.Quiet || opt.NoBar { + r.Progress.Start() + logs.Log.Writer = r.Progress.Bypass() + } if opt.SimhashDistance != 0 { pkg.Distance = uint8(opt.SimhashDistance) } @@ -172,34 +174,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) { } } - // prepare url - var urls []string - var file *os.File - urlfrom := opt.URLFile - if opt.URL != "" { - urls = append(urls, 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 - } - urls = strings.Split(strings.TrimSpace(string(content)), "\n") - } - - r.URLList = urls - logs.Log.Importantf("Loaded %d urls from %s", len(urls), urlfrom) - // prepare word dicts := make([][]string, len(opt.Dictionaries)) for i, f := range opt.Dictionaries { @@ -244,20 +218,69 @@ func (opt *Option) PrepareRunner() (*Runner, error) { return nil, err } logs.Log.Importantf("Parsed %d words by %s", len(r.Wordlist), opt.Word) - pkg.DefaultStatistor.Total = len(r.Wordlist) - pkg.DefaultStatistor.Word = opt.Word - pkg.DefaultStatistor.Dictionaries = opt.Dictionaries + pkg.DefaultStatistor = pkg.Statistor{ + Word: opt.Word, + WordCount: len(r.Wordlist), + Dictionaries: opt.Dictionaries, + Offset: opt.Offset, + } - if r.Limit == 0 { - if r.CheckOnly { - r.Limit = len(r.URLList) - } else { - r.Limit = len(r.Wordlist) + r.Total = len(r.Wordlist) + if opt.Limit != 0 { + if total := r.Offset + opt.Limit; total < r.Total { + r.Total = total + } + } + + // prepare task + var tasks []*Task + var taskfrom string + if opt.ResumeFrom != "" { + stats, err := pkg.ReadStatistors(opt.ResumeFrom) + if err != nil { + return nil, err + } + taskfrom = "resume " + opt.ResumeFrom + for _, stat := range stats { + tasks = append(tasks, &Task{baseUrl: stat.BaseUrl, offset: stat.Offset + stat.ReqNumber, total: r.Total}) } } else { - r.Limit = r.Offset + opt.Limit + var file *os.File + var urls []string + if opt.URL != "" { + urls = append(urls, opt.URL) + tasks = append(tasks, &Task{baseUrl: opt.URL, offset: opt.Offset, total: r.Total}) + taskfrom = "cmd" + } else if opt.URLFile != "" { + file, err = os.Open(opt.URLFile) + if err != nil { + return nil, err + } + taskfrom = opt.URLFile + } else if pkg.HasStdin() { + file = os.Stdin + taskfrom = "stdin" + } + + if file != nil { + content, err := ioutil.ReadAll(file) + if err != nil { + return nil, err + } + urls := strings.Split(strings.TrimSpace(string(content)), "\n") + for _, u := range urls { + tasks = append(tasks, &Task{baseUrl: strings.TrimSpace(u), offset: opt.Offset, total: r.Total}) + } + } + if opt.CheckOnly { + r.URLList = urls + r.Total = len(r.URLList) + } } + r.Tasks = tasks + logs.Log.Importantf("Loaded %d urls from %s", len(tasks), taskfrom) + if opt.Uppercase { r.Fns = append(r.Fns, strings.ToUpper) } @@ -361,7 +384,7 @@ func loadFileToSlice(filename string) ([]string, error) { return nil, err } - ss = strings.Split(string(content), "\n") + ss = strings.Split(strings.TrimSpace(string(content)), "\n") // 统一windows与linux的回车换行差异 for i, word := range ss { @@ -395,3 +418,9 @@ func IntsContains(s []int, e int) bool { } return false } + +type Task struct { + baseUrl string + offset int + total int +} diff --git a/internal/pool.go b/internal/pool.go index 4938a14..909f3d9 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -73,7 +73,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { } p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) { - pool.Statistor.ReqNumber++ + pool.Statistor.Total++ unit := i.(*Unit) req, err := pool.genReq(unit.path) if err != nil { @@ -287,7 +287,6 @@ func (p *Pool) addRedirect(bl *pkg.Baseline, reCount int) { } func (p *Pool) Run(ctx context.Context, offset, limit int) { - p.Statistor.Offset = offset Loop: for { select { diff --git a/internal/runner.go b/internal/runner.go index 9822392..5588171 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -11,7 +11,6 @@ import ( "github.com/gosuri/uiprogress" "github.com/panjf2000/ants/v2" "net/http" - "strings" "sync" "time" ) @@ -24,11 +23,12 @@ var ( ) type Runner struct { - urlCh chan string + taskCh chan *Task poolwg sync.WaitGroup bar *uiprogress.Bar finished int + Tasks []*Task URLList []string Wordlist []string Headers http.Header @@ -50,7 +50,7 @@ type Runner struct { Force bool Progress *uiprogress.Progress Offset int - Limit int + Total int Deadline int CheckPeriod int ErrPeriod int @@ -84,6 +84,7 @@ func (r *Runner) PrepareConfig() *pkg.Config { func (r *Runner) Prepare(ctx context.Context) error { var err error if r.CheckOnly { + // 仅check, 类似httpx r.Pools, err = ants.NewPoolWithFunc(1, func(i interface{}) { config := r.PrepareConfig() config.Wordlist = r.URLList @@ -94,31 +95,31 @@ func (r *Runner) Prepare(ctx context.Context) error { r.poolwg.Done() return } - pool.bar = pkg.NewBar("check", r.Limit-r.Offset, r.Progress) - pool.Run(ctx, r.Offset, r.Limit) + pool.bar = pkg.NewBar("check", r.Total-r.Offset, r.Progress) + pool.Run(ctx, r.Offset, r.Total) r.poolwg.Done() }) } else { go func() { - for _, u := range r.URLList { - r.urlCh <- strings.TrimSpace(u) + for _, t := range r.Tasks { + r.taskCh <- t } - close(r.urlCh) + close(r.taskCh) }() - if len(r.URLList) > 0 { - r.bar = r.Progress.AddBar(len(r.URLList)) + if len(r.Tasks) > 0 { + r.bar = r.Progress.AddBar(len(r.Tasks)) r.bar.PrependCompleted() r.bar.PrependFunc(func(b *uiprogress.Bar) string { - return fmt.Sprintf("total progressive: %d/%d ", r.finished, len(r.URLList)) + return fmt.Sprintf("total progressive: %d/%d ", r.finished, len(r.Tasks)) }) r.bar.AppendElapsed() } r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { - u := i.(string) + t := i.(*Task) config := r.PrepareConfig() - config.BaseURL = u + config.BaseURL = t.baseUrl config.Wordlist = r.Wordlist pool, err := NewPool(ctx, config) if err != nil { @@ -128,7 +129,7 @@ func (r *Runner) Prepare(ctx context.Context) error { return } - pool.bar = pkg.NewBar(u, r.Limit-r.Offset, r.Progress) + pool.bar = pkg.NewBar(config.BaseURL, t.total-t.offset, r.Progress) err = pool.Init() if err != nil { logs.Log.Error(err.Error()) @@ -140,12 +141,11 @@ func (r *Runner) Prepare(ctx context.Context) error { } } - pool.Run(ctx, r.Offset, r.Limit) - + pool.Run(ctx, t.offset, t.total) logs.Log.Important(pool.Statistor.String()) logs.Log.Important(pool.Statistor.Detail()) if r.StatFile != nil { - r.StatFile.SafeWrite(pool.Statistor.Json() + "\n") + r.StatFile.SafeWrite(pool.Statistor.Json()) r.StatFile.SafeSync() } r.Done() @@ -167,18 +167,17 @@ Loop: case <-ctx.Done(): logs.Log.Error("cancel with deadline") break Loop - case u, ok := <-r.urlCh: + case t, ok := <-r.taskCh: if !ok { break Loop } r.poolwg.Add(1) - r.Pools.Invoke(u) + r.Pools.Invoke(t) } } r.poolwg.Wait() - - time.Sleep(100) // 延迟100ms, 等所有数据处理完毕 + //time.Sleep(100 * time.Millisecond) // 延迟100ms, 等所有数据处理完毕 for { if len(r.OutputCh) == 0 { close(r.OutputCh) @@ -192,7 +191,7 @@ Loop: break } } - time.Sleep(100) // 延迟100ms, 等所有数据处理完毕 + time.Sleep(100 * time.Millisecond) // 延迟100ms, 等所有数据处理完毕 } func (r *Runner) RunWithCheck(ctx context.Context) { @@ -225,7 +224,7 @@ Loop: } } - time.Sleep(100) // 延迟100ms, 等所有数据处理完毕 + time.Sleep(100 * time.Millisecond) // 延迟100ms, 等所有数据处理完毕 } func (r *Runner) Done() { diff --git a/pkg/statistor.go b/pkg/statistor.go index 61b8a36..fc8b39e 100644 --- a/pkg/statistor.go +++ b/pkg/statistor.go @@ -1,8 +1,10 @@ package pkg import ( + "bytes" "encoding/json" "fmt" + "io/ioutil" "strconv" "strings" "time" @@ -33,6 +35,7 @@ type Statistor struct { Total int `json:"total"` StartTime int64 `json:"start_time"` EndTime int64 `json:"end_time"` + WordCount int `json:"word_count"` Word string `json:"word"` Dictionaries []string `json:"dictionaries"` } @@ -71,5 +74,24 @@ func (stat *Statistor) Json() string { if err != nil { return err.Error() } - return string(content) + return string(content) + "\n" } + +func ReadStatistors(filename string) (Statistors, error) { + content, err := ioutil.ReadFile(filename) + if err != nil { + return nil, err + } + var stats Statistors + for _, line := range bytes.Split(content, []byte("\n")) { + var stat Statistor + err := json.Unmarshal(line, &stat) + if err != nil { + return nil, err + } + stats = append(stats, stat) + } + return stats, nil +} + +type Statistors []Statistor