diff --git a/internal/checkpool.go b/internal/checkpool.go index 3ee5830..c7e8ff7 100644 --- a/internal/checkpool.go +++ b/internal/checkpool.go @@ -194,7 +194,7 @@ func (pool *CheckPool) doRedirect(bl *pkg.Baseline, depth int) { } reURL = bl.RedirectURL } else { - reURL = bl.BaseURL() + FormatURL(bl.BaseURL(), bl.RedirectURL) + reURL = BaseURL(bl.Url) + FormatURL(BaseURL(bl.Url), bl.RedirectURL) } pool.wg.Add(1) diff --git a/internal/option.go b/internal/option.go index d10c091..655223a 100644 --- a/internal/option.go +++ b/internal/option.go @@ -105,6 +105,8 @@ type ModeOptions struct { Scope []string `long:"scope" description:"String, custom scope, e.g.: --scope *.example.com"` Recursive string `long:"recursive" default:"current.IsDir()" description:"String,custom recursive rule, e.g.: --recursive current.IsDir()"` Depth int `long:"depth" default:"0" description:"Int, recursive depth"` + Index string `long:"index" default:"" description:"String, custom index path"` + Random string `long:"random" default:"" description:"String, custom random path"` CheckPeriod int `long:"check-period" default:"200" description:"Int, check period when request"` ErrPeriod int `long:"error-period" default:"10" description:"Int, check period when error"` BreakThreshold int `long:"error-threshold" default:"20" description:"Int, break when the error exceeds the threshold "` @@ -164,6 +166,8 @@ func (opt *Option) PrepareRunner() (*Runner, error) { Common: opt.Common, RetryCount: opt.RetryCount, RandomUserAgent: opt.RandomUserAgent, + Random: opt.Random, + Index: opt.Index, } // log and bar diff --git a/internal/pool.go b/internal/pool.go index b85f326..0d60a4a 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -143,8 +143,20 @@ func (pool *Pool) genReq(mod pkg.SprayMod, s string) (*ihttp.Request, error) { func (pool *Pool) Init() error { // 分成两步是为了避免闭包的线程安全问题 pool.initwg.Add(2) - pool.reqPool.Invoke(newUnit(pool.url.Path, InitIndexSource)) - pool.reqPool.Invoke(newUnit(pool.safePath(pkg.RandPath()), InitRandomSource)) + if pool.Index != "" { + logs.Log.Importantf("custom index url: %s", BaseURL(pool.url)+FormatURL(BaseURL(pool.url), pool.Index)) + pool.reqPool.Invoke(newUnit(pool.Index, InitIndexSource)) + } else { + pool.reqPool.Invoke(newUnit(pool.url.Path, InitIndexSource)) + } + + if pool.Random != "" { + logs.Log.Importantf("custom random url: %s", BaseURL(pool.url)+FormatURL(BaseURL(pool.url), pool.Random)) + pool.reqPool.Invoke(newUnit(pool.Random, InitRandomSource)) + } else { + pool.reqPool.Invoke(newUnit(pool.safePath(pkg.RandPath()), InitRandomSource)) + } + pool.initwg.Wait() if pool.index.ErrString != "" { logs.Log.Error(pool.index.String()) diff --git a/internal/runner.go b/internal/runner.go index b3b0098..ffa5dfc 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -82,6 +82,8 @@ type Runner struct { Common bool RetryCount int RandomUserAgent bool + Random string + Index string } func (r *Runner) PrepareConfig() *pkg.Config { @@ -110,6 +112,8 @@ func (r *Runner) PrepareConfig() *pkg.Config { Retry: r.RetryCount, ClientType: r.ClientType, RandomUserAgent: r.RandomUserAgent, + Random: r.Random, + Index: r.Index, } if config.ClientType == ihttp.Auto { diff --git a/internal/utils.go b/internal/utils.go index 2ff9073..53f4925 100644 --- a/internal/utils.go +++ b/internal/utils.go @@ -213,7 +213,6 @@ func Dir(u string) string { // /a/ /a/ // a/ a/ // aaa / - if strings.HasSuffix(u, "/") { return u } else if i := strings.LastIndex(u, "/"); i == -1 { @@ -251,6 +250,10 @@ func FormatURL(base, u string) string { } } +func BaseURL(u *url.URL) string { + return u.Scheme + "://" + u.Host +} + func RandomUA() string { return randomUserAgent[rand.Intn(uacount)] } diff --git a/pkg/baseline.go b/pkg/baseline.go index c015ba9..ec27c68 100644 --- a/pkg/baseline.go +++ b/pkg/baseline.go @@ -122,10 +122,6 @@ func (bl *Baseline) IsDir() bool { return false } -func (bl *Baseline) BaseURL() string { - return bl.Url.Scheme + "://" + bl.Url.Host -} - // Collect 深度收集信息 func (bl *Baseline) Collect() { if bl.ContentType == "html" || bl.ContentType == "json" || bl.ContentType == "txt" { diff --git a/pkg/config.go b/pkg/config.go index cb6b788..0f34592 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -47,4 +47,6 @@ type Config struct { Common bool Retry int RandomUserAgent bool + Random string + Index string }