From 9582a325861bc97c8f8254c0183c43d59f931c9a Mon Sep 17 00:00:00 2001 From: M09Ic Date: Wed, 26 Oct 2022 18:28:40 +0800 Subject: [PATCH] =?UTF-8?q?=E7=8E=B0=E5=9C=A8=E5=90=8C=E6=97=B6=E6=94=AF?= =?UTF-8?q?=E6=8C=81http.net=E4=B8=8Efasthttp=E4=B8=A4=E4=B8=AA=E5=BA=93,?= =?UTF-8?q?=20=E9=80=82=E7=94=A8=E4=B8=8D=E5=90=8C=E7=9A=84=E5=9C=BA?= =?UTF-8?q?=E6=99=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/baseline.go | 48 ++++++++++++++-------- internal/pool.go | 86 +++++++++++++++++++++++++-------------- internal/runner.go | 77 +++++++++++++++++++---------------- pkg/client.go | 67 ------------------------------ pkg/config.go | 1 + pkg/ihttp/client.go | 94 +++++++++++++++++++++++++++++++++++++++++++ pkg/ihttp/request.go | 31 ++++++++++++++ pkg/ihttp/response.go | 83 ++++++++++++++++++++++++++++++++++++++ pkg/utils.go | 20 +++++++++ 9 files changed, 359 insertions(+), 148 deletions(-) delete mode 100644 pkg/client.go create mode 100644 pkg/ihttp/client.go create mode 100644 pkg/ihttp/request.go create mode 100644 pkg/ihttp/response.go diff --git a/internal/baseline.go b/internal/baseline.go index 55b09b4..1b0be34 100644 --- a/internal/baseline.go +++ b/internal/baseline.go @@ -2,46 +2,48 @@ package internal import ( "encoding/json" - "fmt" "github.com/chainreactors/parsers" "github.com/chainreactors/spray/pkg" - "github.com/valyala/fasthttp" + "github.com/chainreactors/spray/pkg/ihttp" + "strconv" "strings" ) -func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline { +func NewBaseline(u, host string, resp *ihttp.Response) *baseline { bl := &baseline{ - Url: u, - UrlString: u.String(), + //Url: u, + UrlString: u, + Host: host, Status: resp.StatusCode(), IsValid: true, } bl.Body = resp.Body() - bl.BodyLength = resp.Header.ContentLength() - bl.Header = resp.Header.Header() - bl.HeaderLength = resp.Header.Len() - bl.RedirectURL = string(resp.Header.Peek("Location")) + bl.BodyLength = resp.ContentLength() + bl.Header = resp.Header() + bl.HeaderLength = len(bl.Header) + bl.RedirectURL = resp.GetHeader("Location") bl.Raw = append(bl.Header, bl.Body...) return bl } -func NewInvalidBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline { +func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline { bl := &baseline{ - Url: u, - UrlString: u.String(), + //Url: u, + UrlString: u, + Host: host, Status: resp.StatusCode(), IsValid: false, } - bl.RedirectURL = string(resp.Header.Peek("Location")) + bl.RedirectURL = string(resp.GetHeader("Location")) return bl } type baseline struct { - Url *fasthttp.URI `json:"-"` - UrlString string `json:"url_string"` + UrlString string `json:"url"` + Host string `json:"host"` Body []byte `json:"-"` BodyLength int `json:"body_length"` Header []byte `json:"-"` @@ -51,6 +53,7 @@ type baseline struct { Status int `json:"status"` IsDynamicUrl bool `json:"is_dynamic_url"` // 判断是否存在动态的url Spended int `json:"spended"` // 耗时, 毫秒 + Title string `json:"title"` Frameworks pkg.Frameworks `json:"frameworks"` Extracteds pkg.Extracteds `json:"extracts"` Err error `json:"-"` @@ -61,6 +64,7 @@ type baseline struct { // Collect 深度收集信息 func (bl *baseline) Collect() { bl.Hashes = parsers.NewHashes(bl.Raw) + bl.Title = parsers.MatchTitle(string(bl.Body)) // todo extract bl.Extracteds = pkg.Extractors.Extract(string(bl.Raw)) // todo 指纹识别 @@ -95,11 +99,21 @@ func (bl *baseline) String() string { var line strings.Builder //line.WriteString("[+] ") line.WriteString(bl.UrlString) - line.WriteString(fmt.Sprintf(" - %d - %d ", bl.Status, bl.BodyLength)) + line.WriteString(" (" + bl.Host + ")") + line.WriteString(" - ") + line.WriteString(strconv.Itoa(bl.Status)) + line.WriteString(" - ") + line.WriteString(strconv.Itoa(bl.BodyLength)) if bl.RedirectURL != "" { - line.WriteString("-> ") + line.WriteString(" -> ") line.WriteString(bl.RedirectURL) + line.WriteString(" ") } + line.WriteString(" [" + bl.Title + "]") + if bl.Hashes != nil { + line.WriteString(" [" + bl.Hashes.BodyMd5 + "]") + } + line.WriteString(bl.Frameworks.ToString()) //line.WriteString(bl.Extracteds) //line.WriteString("\n") diff --git a/internal/pool.go b/internal/pool.go index ae70be6..ece7b4f 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -4,9 +4,11 @@ import ( "context" "github.com/chainreactors/logs" "github.com/chainreactors/spray/pkg" + "github.com/chainreactors/spray/pkg/ihttp" "github.com/chainreactors/words" "github.com/panjf2000/ants/v2" "github.com/valyala/fasthttp" + "net/http" "sync" "time" ) @@ -25,7 +27,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( Config: config, ctx: pctx, cancel: cancel, - client: pkg.NewClient(config.Thread, 2), + client: ihttp.NewClient(config.Thread, 2, config.ClientType), worder: words.NewWorder(config.Wordlist), outputCh: outputCh, tempCh: make(chan *baseline, config.Thread), @@ -37,13 +39,32 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( switch config.Mod { case pkg.PathSpray: - pool.genReq = func(s string) (*fasthttp.Request, error) { + pool.genReq = func(s string) (*ihttp.Request, error) { return pool.buildPathRequest(s) } + pool.check = func() { + pool.wg.Add(1) + _ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource)) + + if pool.failedCount > breakThreshold { + // 当报错次数超过上限是, 结束任务 + pool.cancel() + } + } case pkg.HostSpray: - pool.genReq = func(s string) (*fasthttp.Request, error) { + pool.genReq = func(s string) (*ihttp.Request, error) { return pool.buildHostRequest(s) } + + pool.check = func() { + pool.wg.Add(1) + _ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource)) + + if pool.failedCount > breakThreshold { + // 当报错次数超过上限是, 结束任务 + pool.cancel() + } + } } p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) { @@ -56,8 +77,11 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( var bl *baseline resp, reqerr := pool.client.Do(pctx, req) - defer fasthttp.ReleaseResponse(resp) - defer fasthttp.ReleaseRequest(req) + if pool.ClientType == ihttp.FAST { + defer fasthttp.ReleaseResponse(resp.FastResponse) + defer fasthttp.ReleaseRequest(req.FastRequest) + } + if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge { pool.failedCount++ bl = &baseline{UrlString: pool.BaseURL + unit.path, Err: reqerr} @@ -65,9 +89,9 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( pool.failedCount = 0 if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource { // 通过预对比跳过一些无用数据, 减少性能消耗 - bl = NewBaseline(req.URI(), resp) + bl = NewBaseline(req.URI(), req.Host(), resp) } else { - bl = NewInvalidBaseline(req.URI(), resp) + bl = NewInvalidBaseline(req.URI(), req.Host(), resp) } bl.Err = reqerr } @@ -108,7 +132,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( type Pool struct { *pkg.Config - client *pkg.Client + client *ihttp.Client pool *ants.PoolWithFunc bar *pkg.Bar ctx context.Context @@ -122,22 +146,13 @@ type Pool struct { checkPeriod int errPeriod int analyzeDone bool - genReq func(s string) (*fasthttp.Request, error) + genReq func(s string) (*ihttp.Request, error) + check func() worder *words.Worder wg sync.WaitGroup initwg sync.WaitGroup // 初始化用, 之后改成锁 } -func (p *Pool) check() { - p.wg.Add(1) - _ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource)) - - if p.failedCount > breakThreshold { - // 当报错次数超过上限是, 结束任务 - p.cancel() - } -} - func (p *Pool) Init() error { p.initwg.Add(1) p.check() @@ -189,12 +204,12 @@ Loop: p.Close() } -func (p *Pool) PreCompare(resp *fasthttp.Response) error { +func (p *Pool) PreCompare(resp *ihttp.Response) error { if !CheckStatusCode(resp.StatusCode()) { return ErrBadStatus } - if CheckRedirect != nil && !CheckRedirect(string(resp.Header.Peek("Location"))) { + if CheckRedirect != nil && !CheckRedirect(string(resp.GetHeader("Location"))) { return ErrRedirect } @@ -238,15 +253,26 @@ func (p *Pool) Close() { } } -func (p *Pool) buildPathRequest(path string) (*fasthttp.Request, error) { - req := fasthttp.AcquireRequest() - req.SetRequestURI(p.BaseURL + path) - return req, nil +func (p *Pool) buildPathRequest(path string) (*ihttp.Request, error) { + if p.Config.ClientType == ihttp.FAST { + req := fasthttp.AcquireRequest() + req.SetRequestURI(p.BaseURL + path) + return &ihttp.Request{FastRequest: req}, nil + } else { + req, err := http.NewRequest("GET", p.BaseURL+path, nil) + return &ihttp.Request{StandardRequest: req}, err + } } -func (p *Pool) buildHostRequest(host string) (*fasthttp.Request, error) { - req := fasthttp.AcquireRequest() - req.SetRequestURI(p.BaseURL) - req.SetHost(host) - return req, nil +func (p *Pool) buildHostRequest(host string) (*ihttp.Request, error) { + if p.Config.ClientType == ihttp.FAST { + req := fasthttp.AcquireRequest() + req.SetRequestURI(p.BaseURL) + req.SetHost(host) + return &ihttp.Request{FastRequest: req}, nil + } else { + req, err := http.NewRequest("GET", p.BaseURL, nil) + req.Host = host + return &ihttp.Request{StandardRequest: req}, err + } } diff --git a/internal/runner.go b/internal/runner.go index cbac025..e85adaa 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -5,7 +5,9 @@ import ( "fmt" "github.com/chainreactors/logs" "github.com/chainreactors/spray/pkg" + "github.com/chainreactors/spray/pkg/ihttp" "github.com/gosuri/uiprogress" + "github.com/panjf2000/ants/v2" "io/ioutil" "net/http" "os" @@ -27,8 +29,9 @@ type Runner struct { Offset int `long:"offset"` Limit int `long:"limit"` Threads int `short:"t" long:"thread" default:"20"` - PoolSize int `short:"p" long:"pool"` - Pools map[string]*Pool + PoolSize int `short:"p" long:"pool" default:"5"` + Pools *ants.PoolWithFunc + poolwg sync.WaitGroup Deadline int `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline Debug bool `long:"debug"` Quiet bool `short:"q" long:"quiet"` @@ -108,46 +111,52 @@ func (r *Runner) Prepare() error { } r.OutputCh = make(chan *baseline, 100) - r.Pools = make(map[string]*Pool) + ctx := context.Background() + + r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { + u := i.(string) + config := &pkg.Config{ + BaseURL: u, + Wordlist: r.Wordlist, + Thread: r.Threads, + Timeout: 2, + Headers: r.Headers, + Mod: pkg.ModMap[r.Mod], + DeadlineTime: r.Deadline, + } + + if config.Mod == pkg.PathSpray { + config.ClientType = ihttp.FAST + } else if config.Mod == pkg.HostSpray { + config.ClientType = ihttp.STANDARD + } + + pool, err := NewPool(ctx, config, r.OutputCh) + if err != nil { + logs.Log.Error(err.Error()) + return + } + pool.bar = pkg.NewBar(u, len(r.Wordlist), r.Progress) + err = pool.Init() + if err != nil { + logs.Log.Error(err.Error()) + return + } + // todo pool 总超时时间 + pool.Run(ctx) + r.poolwg.Done() + }) go r.Outputting() return nil } func (r *Runner) Run() { // todo pool 结束与并发控制 - ctx := context.Background() - var wg sync.WaitGroup for _, u := range r.URLList { - wg.Add(1) - u := u - go func() { - config := &pkg.Config{ - BaseURL: u, - Wordlist: r.Wordlist, - Thread: r.Threads, - Timeout: 2, - Headers: r.Headers, - Mod: pkg.ModMap[r.Mod], - DeadlineTime: r.Deadline, - } - pool, err := NewPool(ctx, config, r.OutputCh) - if err != nil { - logs.Log.Error(err.Error()) - return - } - pool.bar = pkg.NewBar(u, len(r.Wordlist), r.Progress) - err = pool.Init() - if err != nil { - logs.Log.Error(err.Error()) - return - } - r.Pools[u] = pool - // todo pool 总超时时间 - pool.Run(ctx) - wg.Done() - }() + r.poolwg.Add(1) + r.Pools.Invoke(u) } - wg.Wait() + r.poolwg.Wait() for { if len(r.OutputCh) == 0 { close(r.OutputCh) diff --git a/pkg/client.go b/pkg/client.go deleted file mode 100644 index 2bdc742..0000000 --- a/pkg/client.go +++ /dev/null @@ -1,67 +0,0 @@ -package pkg - -import ( - "context" - "crypto/tls" - "github.com/valyala/fasthttp" - "net/http" - "time" -) - -var ( - DefaultMaxBodySize = 1024 * 100 // 100k -) - -func NewClient(thread int, timeout int) *Client { - c := &Client{ - client: &fasthttp.Client{ - TLSConfig: &tls.Config{ - Renegotiation: tls.RenegotiateOnceAsClient, - InsecureSkipVerify: true, - }, - //ReadBufferSize: 20480, - MaxConnsPerHost: thread * 2, - MaxIdleConnDuration: time.Duration(timeout) * time.Second, - MaxConnWaitTimeout: time.Duration(timeout) * time.Second, - ReadTimeout: time.Duration(timeout) * time.Second, - WriteTimeout: time.Duration(timeout) * time.Second, - MaxResponseBodySize: DefaultMaxBodySize, - }, - timeout: time.Duration(timeout) * time.Second, - } - //c := &Client{ - // client: &http.Client{ - // Transport: tr, - // Timeout: time.Second * time.Duration(timeout), - // CheckRedirect: checkRedirect, - // }, - //} - - //c.Method = method - //c.Headers = Opt.Headers - //c.Mod = Opt.Mod - - return c -} - -type Client struct { - client *fasthttp.Client - timeout time.Duration -} - -func (c *Client) Do(ctx context.Context, req *fasthttp.Request) (*fasthttp.Response, error) { - //if req.Header == nil { - // req.Header = c.Headers - //} - resp := fasthttp.AcquireResponse() - return resp, c.client.Do(req, resp) -} - -var MaxRedirects = 0 -var checkRedirect = func(req *http.Request, via []*http.Request) error { - if len(via) > MaxRedirects { - return http.ErrUseLastResponse - } - - return nil -} diff --git a/pkg/config.go b/pkg/config.go index 517d3c8..74fddb3 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -29,4 +29,5 @@ type Config struct { Headers http.Header DeadlineTime int EnableFuzzy bool + ClientType int } diff --git a/pkg/ihttp/client.go b/pkg/ihttp/client.go new file mode 100644 index 0000000..572fa74 --- /dev/null +++ b/pkg/ihttp/client.go @@ -0,0 +1,94 @@ +package ihttp + +import ( + "context" + "crypto/tls" + "fmt" + "github.com/valyala/fasthttp" + "net/http" + "time" +) + +var ( + DefaultMaxBodySize = 1024 * 100 // 100k +) + +const ( + FAST = iota + STANDARD +) + +func NewClient(thread int, timeout int, clientType int) *Client { + if clientType == FAST { + return &Client{ + fastClient: &fasthttp.Client{ + TLSConfig: &tls.Config{ + Renegotiation: tls.RenegotiateOnceAsClient, + InsecureSkipVerify: true, + }, + //ReadBufferSize: 20480, + MaxConnsPerHost: thread * 2, + MaxIdleConnDuration: time.Duration(timeout) * time.Second, + MaxConnWaitTimeout: time.Duration(timeout) * time.Second, + ReadTimeout: time.Duration(timeout) * time.Second, + WriteTimeout: time.Duration(timeout) * time.Second, + MaxResponseBodySize: DefaultMaxBodySize, + }, + timeout: time.Duration(timeout) * time.Second, + } + } else { + return &Client{ + standardClient: &http.Client{ + Transport: &http.Transport{ + //Proxy: Proxy, + //TLSHandshakeTimeout : delay * time.Second, + TLSClientConfig: &tls.Config{ + Renegotiation: tls.RenegotiateOnceAsClient, + InsecureSkipVerify: true, + }, + MaxConnsPerHost: thread, + IdleConnTimeout: time.Duration(timeout) * time.Second, + }, + Timeout: time.Second * time.Duration(timeout), + CheckRedirect: checkRedirect, + }, + timeout: time.Duration(timeout) * time.Second, + } + } +} + +type Client struct { + fastClient *fasthttp.Client + standardClient *http.Client + timeout time.Duration +} + +func (c *Client) FastDo(ctx context.Context, req *fasthttp.Request) (*fasthttp.Response, error) { + resp := fasthttp.AcquireResponse() + return resp, c.fastClient.Do(req, resp) +} + +func (c *Client) StandardDo(ctx context.Context, req *http.Request) (*http.Response, error) { + return c.standardClient.Do(req) +} + +func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) { + if c.fastClient != nil { + resp, err := c.FastDo(ctx, req.FastRequest) + return &Response{FastResponse: resp}, err + } else if c.standardClient != nil { + resp, err := c.StandardDo(ctx, req.StandardRequest) + return &Response{StandardResponse: resp}, err + } else { + return nil, fmt.Errorf("not found client") + } +} + +var MaxRedirects = 0 +var checkRedirect = func(req *http.Request, via []*http.Request) error { + if len(via) > MaxRedirects { + return http.ErrUseLastResponse + } + + return nil +} diff --git a/pkg/ihttp/request.go b/pkg/ihttp/request.go new file mode 100644 index 0000000..cd23c94 --- /dev/null +++ b/pkg/ihttp/request.go @@ -0,0 +1,31 @@ +package ihttp + +import ( + "github.com/valyala/fasthttp" + "net/http" +) + +type Request struct { + StandardRequest *http.Request + FastRequest *fasthttp.Request +} + +func (r *Request) URI() string { + if r.FastRequest != nil { + return r.FastRequest.URI().String() + } else if r.StandardRequest != nil { + return r.StandardRequest.URL.String() + } else { + return "" + } +} + +func (r *Request) Host() string { + if r.FastRequest != nil { + return string(r.FastRequest.Host()) + } else if r.StandardRequest != nil { + return r.StandardRequest.Host + } else { + return "" + } +} diff --git a/pkg/ihttp/response.go b/pkg/ihttp/response.go new file mode 100644 index 0000000..6d251d6 --- /dev/null +++ b/pkg/ihttp/response.go @@ -0,0 +1,83 @@ +package ihttp + +import ( + "bytes" + "github.com/chainreactors/logs" + "github.com/valyala/fasthttp" + "io" + "net/http" +) + +type Response struct { + StandardResponse *http.Response + FastResponse *fasthttp.Response +} + +func (r *Response) StatusCode() int { + if r.FastResponse != nil { + return r.FastResponse.StatusCode() + } else if r.StandardResponse != nil { + return r.StandardResponse.StatusCode + } else { + return 0 + } +} + +func (r *Response) Body() []byte { + if r.FastResponse != nil { + return r.FastResponse.Body() + } else if r.StandardResponse != nil { + body := make([]byte, 20480) + if r.StandardResponse.ContentLength > 0 { + n, err := io.ReadFull(r.StandardResponse.Body, body) + _ = r.StandardResponse.Body.Close() + if err == nil { + return body + } else if err == io.ErrUnexpectedEOF { + return body[:n] + } else { + logs.Log.Error("readfull failed" + err.Error()) + return nil + } + } + return nil + } else { + return nil + } +} + +func (r *Response) ContentLength() int { + if r.FastResponse != nil { + return r.FastResponse.Header.ContentLength() + } else if r.StandardResponse != nil { + return int(r.StandardResponse.ContentLength) + } else { + return 0 + } +} + +func (r *Response) Header() []byte { + if r.FastResponse != nil { + return r.FastResponse.Header.Header() + } else if r.StandardResponse != nil { + var header bytes.Buffer + for k, v := range r.StandardResponse.Header { + for _, i := range v { + header.WriteString(k + ": " + i + "\r\n") + } + } + return header.Bytes() + } else { + return nil + } +} + +func (r *Response) GetHeader(key string) string { + if r.FastResponse != nil { + return string(r.FastResponse.Header.Peek(key)) + } else if r.StandardResponse != nil { + return r.StandardResponse.Header.Get(key) + } else { + return "" + } +} diff --git a/pkg/utils.go b/pkg/utils.go index 1b4f857..f412755 100644 --- a/pkg/utils.go +++ b/pkg/utils.go @@ -56,3 +56,23 @@ func RandPath() string { } return *(*string)(unsafe.Pointer(&b)) } + +func RandHost() string { + n := 8 + b := make([]byte, n) + // A rand.Int63() generates 63 random bits, enough for letterIdMax letters! + for i, cache, remain := n-1, src.Int63(), letterIdMax; i >= 1; { + if remain == 0 { + cache, remain = src.Int63(), letterIdMax + } + if idx := int(cache & letterIdMask); idx < len(letters) { + b[i] = letters[idx] + i-- + } + cache >>= letterIdBits + remain-- + } + + b[5] = byte(0x2e) + return *(*string)(unsafe.Pointer(&b)) +}