Merge pull request #1062 from projectdiscovery/feature-preserve-input-order

Adding support to preserve input order
This commit is contained in:
Sandeep Singh 2021-10-16 19:17:08 +05:30 committed by GitHub
commit ae55147cfa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 59 additions and 4 deletions

View File

@ -134,6 +134,7 @@ on extensive configurability, massive extensibility and ease of use.`)
flagSet.StringVar(&options.ProjectPath, "project-path", os.TempDir(), "set a specific project path"),
flagSet.BoolVarP(&options.StopAtFirstMatch, "stop-at-first-path", "spm", false, "stop processing HTTP requests after the first match (may break template/workflow logic)"),
flagSet.BoolVar(&options.Stream, "stream", false, "Stream mode - start elaborating without sorting the input"),
)
createGroup(flagSet, "headless", "Headless",

View File

@ -31,6 +31,7 @@ require (
github.com/pkg/errors v0.9.1
github.com/projectdiscovery/clistats v0.0.8
github.com/projectdiscovery/fastdialer v0.0.13-0.20210917073912-cad93d88e69e
github.com/projectdiscovery/filekv v0.0.0-20210915124239-3467ef45dd08
github.com/projectdiscovery/fileutil v0.0.0-20210928100737-cab279c5d4b5
github.com/projectdiscovery/goflags v0.0.8-0.20211007103353-9b9229e8a240
github.com/projectdiscovery/gologger v1.1.4
@ -73,6 +74,8 @@ require (
github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 // indirect
github.com/andybalholm/cascadia v1.1.0 // indirect
github.com/antchfx/xpath v1.1.6 // indirect
github.com/bits-and-blooms/bitset v1.2.0 // indirect
github.com/bits-and-blooms/bloom/v3 v3.0.1 // indirect
github.com/aymerick/douceur v0.2.0 // indirect
github.com/cnf/structhash v0.0.0-20201127153200-e1b16c1ebc08 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect

View File

@ -109,6 +109,10 @@ github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24
github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
github.com/bits-and-blooms/bitset v1.2.0 h1:Kn4yilvwNtMACtf1eYDlG8H77R07mZSPbMjLyS07ChA=
github.com/bits-and-blooms/bitset v1.2.0/go.mod h1:gIdJ4wp64HaoK2YrL1Q5/N7Y16edYb8uY+O0FJTyyDA=
github.com/bits-and-blooms/bloom/v3 v3.0.1 h1:Inlf0YXbgehxVjMPmCGv86iMCKMGPPrPSHtBF5yRHwA=
github.com/bits-and-blooms/bloom/v3 v3.0.1/go.mod h1:MC8muvBzzPOFsrcdND/A7kU7kMhkqb9KI70JlZCP+C8=
github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ=
github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
github.com/bluele/gcache v0.0.2 h1:WcbfdXICg7G/DGBh1PFfcirkWOQV+v077yF1pSy3DGw=
@ -569,7 +573,10 @@ github.com/projectdiscovery/fastdialer v0.0.12/go.mod h1:RkRbxqDCcCFhfNUbkzBIz/i
github.com/projectdiscovery/fastdialer v0.0.13-0.20210824195254-0113c1406542/go.mod h1:TuapmLiqtunJOxpM7g0tpTy/TUF/0S+XFyx0B0Wx0DQ=
github.com/projectdiscovery/fastdialer v0.0.13-0.20210917073912-cad93d88e69e h1:xMAFYJgRxopAwKrj7HDwMBKJGCGDbHqopS8f959xges=
github.com/projectdiscovery/fastdialer v0.0.13-0.20210917073912-cad93d88e69e/go.mod h1:O1l6+vAQy1QRo9FqyuyJ57W3CwpIXXg7oGo14Le6ZYQ=
github.com/projectdiscovery/filekv v0.0.0-20210915124239-3467ef45dd08 h1:NwD1R/du1dqrRKN3SJl9kT6tN3K9puuWFXEvYF2ihew=
github.com/projectdiscovery/filekv v0.0.0-20210915124239-3467ef45dd08/go.mod h1:paLCnwV8sL7ppqIwVQodQrk3F6mnWafwTDwRd7ywZwQ=
github.com/projectdiscovery/fileutil v0.0.0-20210804142714-ebba15fa53ca/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
github.com/projectdiscovery/fileutil v0.0.0-20210914153648-31f843feaad4/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
github.com/projectdiscovery/fileutil v0.0.0-20210928100737-cab279c5d4b5 h1:2dbm7UhrAKnccZttr78CAmG768sSCd+MBn4ayLVDeqA=
github.com/projectdiscovery/fileutil v0.0.0-20210928100737-cab279c5d4b5/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
github.com/projectdiscovery/goflags v0.0.7/go.mod h1:Jjwsf4eEBPXDSQI2Y+6fd3dBumJv/J1U0nmpM+hy2YY=

View File

@ -11,7 +11,7 @@ import (
func (r *Runner) processTemplateWithList(template *templates.Template) bool {
results := &atomic.Bool{}
wg := sizedwaitgroup.New(r.options.BulkSize)
r.hostMap.Scan(func(k, _ []byte) error {
processItem := func(k, _ []byte) error {
URL := string(k)
// Skip if the host has had errors
@ -29,7 +29,13 @@ func (r *Runner) processTemplateWithList(template *templates.Template) bool {
results.CAS(false, match)
}(URL)
return nil
})
}
if r.options.Stream {
_ = r.hostMapStream.Scan(processItem)
} else {
r.hostMap.Scan(processItem)
}
wg.Wait()
return results.Load()
}
@ -39,7 +45,7 @@ func (r *Runner) processWorkflowWithList(template *templates.Template) bool {
results := &atomic.Bool{}
wg := sizedwaitgroup.New(r.options.BulkSize)
r.hostMap.Scan(func(k, _ []byte) error {
processItem := func(k, _ []byte) error {
URL := string(k)
// Skip if the host has had errors
@ -53,7 +59,14 @@ func (r *Runner) processWorkflowWithList(template *templates.Template) bool {
results.CAS(false, match)
}(URL)
return nil
})
}
if r.options.Stream {
_ = r.hostMapStream.Scan(processItem)
} else {
r.hostMap.Scan(processItem)
}
wg.Wait()
return results.Load()
}

View File

@ -16,6 +16,8 @@ import (
"go.uber.org/ratelimit"
"gopkg.in/yaml.v2"
"github.com/projectdiscovery/filekv"
"github.com/projectdiscovery/fileutil"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/hmap/store/hybrid"
"github.com/projectdiscovery/nuclei/v2/internal/colorizer"
@ -45,6 +47,7 @@ import (
// Runner is a client for running the enumeration process.
type Runner struct {
hostMap *hybrid.HybridMap
hostMapStream *filekv.FileDB
output output.Writer
interactsh *interactsh.Client
inputCount int64
@ -119,6 +122,20 @@ func New(options *types.Options) (*Runner, error) {
}
runner.hostMap = hm
if options.Stream {
fkvOptions := filekv.DefaultOptions
if tmpFileName, err := fileutil.GetTempFileName(); err != nil {
return nil, errors.Wrap(err, "could not create temporary input file")
} else {
fkvOptions.Path = tmpFileName
}
fkv, err := filekv.Open(fkvOptions)
if err != nil {
return nil, errors.Wrap(err, "could not create temporary unsorted input file")
}
runner.hostMapStream = fkv
}
runner.inputCount = 0
dupeCount := 0
@ -138,6 +155,9 @@ func New(options *types.Options) (*Runner, error) {
runner.inputCount++
// nolint:errcheck // ignoring error
runner.hostMap.Set(url, nil)
if options.Stream {
_ = runner.hostMapStream.Set([]byte(url), nil)
}
}
}
@ -158,6 +178,9 @@ func New(options *types.Options) (*Runner, error) {
runner.inputCount++
// nolint:errcheck // ignoring error
runner.hostMap.Set(url, nil)
if options.Stream {
_ = runner.hostMapStream.Set([]byte(url), nil)
}
}
}
@ -180,6 +203,9 @@ func New(options *types.Options) (*Runner, error) {
runner.inputCount++
// nolint:errcheck // ignoring error
runner.hostMap.Set(url, nil)
if options.Stream {
_ = runner.hostMapStream.Set([]byte(url), nil)
}
}
input.Close()
}
@ -290,6 +316,9 @@ func (r *Runner) Close() {
if r.projectFile != nil {
r.projectFile.Close()
}
if r.options.Stream {
r.hostMapStream.Close()
}
protocolinit.Close()
}

View File

@ -146,6 +146,8 @@ type Options struct {
Stdin bool
// StopAtFirstMatch stops processing template at first full match (this may break chained requests)
StopAtFirstMatch bool
// Stream the input without sorting
Stream bool
// NoMeta disables display of metadata for the matches
NoMeta bool
// NoTimestamp disables display of timestamp for the matcher