现在同时支持http.net与fasthttp两个库, 适用不同的场景

This commit is contained in:
M09Ic 2022-10-26 18:28:40 +08:00
parent 029a83faa8
commit 9582a32586
9 changed files with 359 additions and 148 deletions

View File

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

View File

@ -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) {
func (p *Pool) buildPathRequest(path string) (*ihttp.Request, error) {
if p.Config.ClientType == ihttp.FAST {
req := fasthttp.AcquireRequest()
req.SetRequestURI(p.BaseURL + path)
return req, nil
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) {
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 req, nil
return &ihttp.Request{FastRequest: req}, nil
} else {
req, err := http.NewRequest("GET", p.BaseURL, nil)
req.Host = host
return &ihttp.Request{StandardRequest: req}, err
}
}

View File

@ -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,19 +111,10 @@ func (r *Runner) Prepare() error {
}
r.OutputCh = make(chan *baseline, 100)
r.Pools = make(map[string]*Pool)
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() {
r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
u := i.(string)
config := &pkg.Config{
BaseURL: u,
Wordlist: r.Wordlist,
@ -130,6 +124,13 @@ func (r *Runner) Run() {
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())
@ -141,13 +142,21 @@ func (r *Runner) Run() {
logs.Log.Error(err.Error())
return
}
r.Pools[u] = pool
// todo pool 总超时时间
pool.Run(ctx)
wg.Done()
}()
r.poolwg.Done()
})
go r.Outputting()
return nil
}
wg.Wait()
func (r *Runner) Run() {
// todo pool 结束与并发控制
for _, u := range r.URLList {
r.poolwg.Add(1)
r.Pools.Invoke(u)
}
r.poolwg.Wait()
for {
if len(r.OutputCh) == 0 {
close(r.OutputCh)

View File

@ -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
}

View File

@ -29,4 +29,5 @@ type Config struct {
Headers http.Header
DeadlineTime int
EnableFuzzy bool
ClientType int
}

94
pkg/ihttp/client.go Normal file
View File

@ -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
}

31
pkg/ihttp/request.go Normal file
View File

@ -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 ""
}
}

83
pkg/ihttp/response.go Normal file
View File

@ -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 ""
}
}

View File

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