diff --git a/internal/option.go b/internal/option.go index 0300b12..7323295 100644 --- a/internal/option.go +++ b/internal/option.go @@ -55,11 +55,12 @@ type RequestOptions struct { } type ModeOptions struct { - Force bool `long:"force"` - CheckOnly bool `long:"check-only"` - CheckPeriod int `long:"check-period" default:"100"` - ErrPeriod int `long:"error-period" default:"10"` - BreakThreshold int `long:"error-threshold" default:"20"` + Force bool `long:"force"` + CheckOnly bool `long:"check-only"` + CheckPeriod int `long:"check-period" default:"100"` + ErrPeriod int `long:"error-period" default:"10"` + BreakThreshold int `long:"error-threshold" default:"20"` + BlackStatus string `long:"black-status" default:"default"` } type MiscOptions struct { @@ -124,6 +125,18 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.ErrPeriod = max } + if opt.BlackStatus != "default" { + for _, s := range strings.Split(opt.BlackStatus, ",") { + si, err := strconv.Atoi(s) + if err != nil { + return nil, err + } + BlackStatus = append(BlackStatus, si) + } + } else { + BlackStatus = []int{400, 404, 410} + } + // prepare url var urls []string var file *os.File diff --git a/internal/pool.go b/internal/pool.go index 78f0a09..d403f97 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -15,10 +15,19 @@ import ( ) var ( - CheckStatusCode func(int) bool - CheckRedirect func(string) bool - CheckWaf func([]byte) bool + CheckBadStatus func(int) bool + CheckRedirect func(string) bool ) + +func CheckWaf(status int) bool { + for _, s := range WAFStatus { + if status == s { + return true + } + } + return false +} + var max = 2147483647 func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { @@ -202,7 +211,7 @@ func (p *Pool) Init() error { p.index.Collect() logs.Log.Important("[baseline.random] " + p.base.String()) - logs.Log.Important("[baseline.index] " + p.base.String()) + logs.Log.Important("[baseline.index] " + p.index.String()) if p.base.RedirectURL != "" { CheckRedirect = func(redirectURL string) bool { @@ -256,23 +265,23 @@ Loop: } func (p *Pool) PreCompare(resp *ihttp.Response) error { - if p.base != nil && p.base.Status != 200 && p.base.Status == resp.StatusCode() { + status := resp.StatusCode() + if p.base != nil && p.base.Status != 200 && p.base.Status == status { return ErrSameStatus } - if !CheckStatusCode(resp.StatusCode()) { + if CheckBadStatus(status) { return ErrBadStatus } + if CheckWaf(status) { + return ErrWaf + } + if CheckRedirect != nil && !CheckRedirect(string(resp.GetHeader("Location"))) { return ErrRedirect } - if CheckWaf != nil && !CheckWaf(nil) { - // todo check waf - return ErrWaf - } - return nil } diff --git a/internal/runner.go b/internal/runner.go index 4f201a2..7ed61bb 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -13,8 +13,11 @@ import ( "time" ) -var BlackStatus = []int{400, 404, 410} -var FuzzyStatus = []int{403, 500, 501, 502, 503} +var ( + BlackStatus = []int{} + FuzzyStatus = []int{403, 500, 501, 502, 503} + WAFStatus = []int{493, 418} +) type Runner struct { URLList chan string @@ -46,13 +49,13 @@ type Runner struct { func (r *Runner) Prepare(ctx context.Context) error { var err error - CheckStatusCode = func(status int) bool { + CheckBadStatus = func(status int) bool { for _, black := range BlackStatus { if black == status { - return false + return true } } - return true + return false } r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {