feat. 新增--index --random 支持自定义index与random目录

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)

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

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())

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 {

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)]
 }

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" {

pkg/config.go

@@ -47,4 +47,6 @@ type Config struct {
 	Common          bool
 	Retry           int
 	RandomUserAgent bool
+	Random          string
+	Index           string
 }