协程池换成ants框架，最多限度利用发挥内存作用

Author: SummerSec

README.md

@@ -30,6 +30,7 @@
 * [x] 命令执行漏洞式支持交互式执行命令
 * [ ] 验证url是否存活
 * [x] 增加自动更新参数
+* [x] 随机User-Agent请求头
 
 ………
 

cmd/commons/core/runner.go

@@ -5,6 +5,7 @@ import (
 	"github.com/SummerSec/SpringExploit/cmd/commons/attack"
 	"github.com/SummerSec/SpringExploit/cmd/commons/utils"
 	"github.com/fatih/structs"
+	"github.com/panjf2000/ants/v2"
 	log "github.com/sirupsen/logrus"
 	"net/url"
 	"sync"
@@ -57,35 +58,75 @@ func (r *Runner) Run() {
 	k := r.options.Thread
 	var wg sync.WaitGroup
 	hashmap := structs.Map(&r.options)
+	defer ants.Release()
+	// TODO: check if options are valid
+	if k <= 0 {
+		k = 500
+	}
+	log.Info("Running with ", k, " threads")
+	pool, err1 := ants.NewPool(k+1, ants.WithPreAlloc(true))
+	if err1 != nil {
+		log.Error("Error creating pool")
+		return
+	}
+	log.Info("Total URLs: ", len(urls))
 	for i < len(urls) {
+		//TODO 老代码
+		//for t := 0; t < k; t++ {
+		//	if i == len(urls) {
+		//		break
+		//	}
+		//	if urls[i] != "" {
+		//		log.Debugln("Running attack on: ", urls[i])
+		//		// 通道通信 发送url 并且 i++
+		//		c := make(chan int)
+		//		wg.Add(1)
+		//		go func() {
+		//			log.Debugf("Running go func() %d", t)
+		//			Start(urls[i], hashmap, i, c) // Start k goroutines
+		//			wg.Done()
+		//		}()
+		//		i = <-c
+		//	} else {
+		//		wg.Wait()
+		//		i++
+		//		break
+		//	}
+		//}
+
 		for t := 0; t < k; t++ {
-			if i == len(urls) {
+			if i == len(urls)-1 {
 				break
 			}
 			if urls[i] != "" {
-				log.Debugln("Running attack on: ", urls[i])
 				// 通道通信 发送url 并且 i++
-				c := make(chan int)
+				//c := make(chan int)
+				log.Debugf("Now Threads: %d", t)
 				wg.Add(1)
-				go func() {
-					log.Debugf("Running go func() %d", t)
-					Start(urls[i], hashmap, i, c) // Start k goroutines
+				err := pool.Submit(func() {
+					log.Debugf("Running Submit %d url is  %s %d", t, urls[i], i)
+					Start2(urls[i], hashmap, i) // Start k goroutines
 					wg.Done()
-				}()
-				i = <-c
+				})
+				i++
+				if err != nil {
+					log.Error("Error submitting job " + urls[i])
+					log.Error(err)
+				}
+				//i = <-c
 			} else {
-				wg.Wait()
 				i++
-				break
+				wg.Wait()
 			}
 		}
+
 	}
 
 }
 
 func Start(u string, hashmap map[string]interface{}, i int, c chan int) {
 	log.Info("Runner started")
-	log.Infoln("testing URL: ", u)
+	log.Infoln("Pen-testing URL: ", u)
 	//for k, v := range hashmap {
 	//	log.Debugln("key: ", k, " value: ", v)
 	//}
@@ -111,3 +152,32 @@ func Start(u string, hashmap map[string]interface{}, i int, c chan int) {
 	c <- i + 1
 
 }
+
+func Start2(u string, hashmap map[string]interface{}, i int) {
+	log.Infof("%s Runner started", u)
+	//log.Infoln("testing URL: ", u)
+	//for k, v := range hashmap {
+	//	log.Debugln("key: ", k, " value: ", v)
+	//}
+	defer func() {
+		if errs := recover(); errs != nil {
+			log.Debug("Runner panic: ", errs)
+		}
+	}()
+
+	r, err := url.Parse(u)
+	if err != nil {
+		log.Info("URL parse error")
+		log.Errorln(err)
+		return
+	}
+	var target string
+	if r.Path == "" {
+		target = r.Scheme + "://" + r.Host + "/"
+	} else {
+		target = r.Scheme + "://" + r.Host + r.Path
+	}
+	attack.Sevice(target, hashmap)
+	log.Infof("%s Runner finished", u)
+
+}

cmd/test/ants.go

@@ -0,0 +1,37 @@
+package main
+
+import (
+	"fmt"
+	"github.com/panjf2000/ants/v2"
+	"sync"
+	"time"
+)
+
+func Task() {
+	fmt.Println(fmt.Sprintf("%d  %s", time.Now().Nanosecond(), "Hello, World!"))
+	time.Sleep(time.Second * 2)
+
+}
+
+func main() {
+	defer ants.Release()
+	pool, err := ants.NewPool(50)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+	var wg sync.WaitGroup
+
+	task := func() {
+		Task()
+		wg.Done()
+	}
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		fmt.Println(fmt.Sprintf("%d  %s", i, "before submit"))
+		_ = pool.Submit(task)
+	}
+	wg.Wait()
+
+}

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/gosuri/uiprogress v0.0.1 // indirect
 	github.com/hashicorp/errwrap v1.0.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
+	github.com/panjf2000/ants/v2 v2.5.0 // indirect
 	github.com/projectdiscovery/blackrock v0.0.0-20210415162320-b38689ae3a2e // indirect
 	github.com/tj/go-update v2.2.5-0.20200519121640-62b4b798fd68+incompatible
 	golang.org/x/net v0.0.0-20220111093109-d55c255bac03 // indirect

go.sum

@@ -97,6 +97,8 @@ github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+W
 github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
 github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
+github.com/panjf2000/ants/v2 v2.5.0 h1:1rWGWSnxCsQBga+nQbA4/iY6VMeNoOIAM0ZWh9u3q2Q=
+github.com/panjf2000/ants/v2 v2.5.0/go.mod h1:cU93usDlihJZ5CfRGNDYsiBYvoilLvBF5Qp/BT2GNRE=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -129,6 +131,7 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ=
 github.com/tj/assert v0.0.0-20171129193455-018094318fb0/go.mod h1:mZ9/Rh9oLWpLLDRpvE+3b7gP/C2YyLFYxNmcLnPTMe0=
 github.com/tj/assert v0.0.3/go.mod h1:Ne6X72Q+TB1AteidzQncjw9PabbMp4PBMZ1k+vd1Pvk=