Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
projectdiscovery
GitHub Repository: projectdiscovery/nuclei
 Path: blob/dev/examples/with_speed_control/main.go
2070 views
1
package main

2


3
import (

4
	"context"

5
	"log"

6
	"sync"

7
	"time"

8


9
	nuclei "github.com/projectdiscovery/nuclei/v3/lib"

10
	"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"

11
)

12


13
func main() {

14
	ne, err := initializeNucleiEngine()

15
	if err != nil {

16
		panic(err)

17
	}

18
	defer ne.Close()

19


20
	ne.LoadTargets([]string{"http://honey.scanme.sh"}, false)

21


22
	var wg sync.WaitGroup

23
	wg.Add(3)

24


25
	go testRateLimit(&wg, ne)

26
	go testThreadsAndBulkSize(&wg, ne)

27
	go testPayloadConcurrency(&wg, ne)

28


29
	err = ne.ExecuteWithCallback(nil)

30
	if err != nil {

31
		panic(err)

32
	}

33


34
	wg.Wait()

35
}

36


37
func initializeNucleiEngine() (*nuclei.NucleiEngine, error) {

38
	return nuclei.NewNucleiEngineCtx(context.TODO(),

39
		nuclei.WithTemplateFilters(nuclei.TemplateFilters{Tags: []string{"oast"}}),

40
		nuclei.EnableStatsWithOpts(nuclei.StatsOptions{MetricServerPort: 6064}),

41
		nuclei.WithGlobalRateLimit(1, time.Second),

42
		nuclei.WithConcurrency(nuclei.Concurrency{

43
			TemplateConcurrency:           1,

44
			HostConcurrency:               1,

45
			HeadlessHostConcurrency:       1,

46
			HeadlessTemplateConcurrency:   1,

47
			JavascriptTemplateConcurrency: 1,

48
			TemplatePayloadConcurrency:    1,

49
			ProbeConcurrency:              1,

50
		}),

51
	)

52
}

53


54
func testRateLimit(wg *sync.WaitGroup, ne *nuclei.NucleiEngine) {

55
	defer wg.Done()

56
	verifyRateLimit(ne, 1, 5000)

57
}

58


59
func testThreadsAndBulkSize(wg *sync.WaitGroup, ne *nuclei.NucleiEngine) {

60
	defer wg.Done()

61
	initialTemplateThreads, initialBulkSize := 1, 1

62
	verifyThreadsAndBulkSize(ne, initialTemplateThreads, initialBulkSize, 25, 25)

63
}

64


65
func testPayloadConcurrency(wg *sync.WaitGroup, ne *nuclei.NucleiEngine) {

66
	defer wg.Done()

67
	verifyPayloadConcurrency(ne, 1, 500)

68
}

69


70
func verifyRateLimit(ne *nuclei.NucleiEngine, initialRate, finalRate int) {

71
	if ne.GetExecuterOptions().RateLimiter.GetLimit() != uint(initialRate) {

72
		panic("wrong initial rate limit")

73
	}

74
	time.Sleep(5 * time.Second)

75
	ne.Options().RateLimit = finalRate

76
	time.Sleep(20 * time.Second)

77
	if ne.GetExecuterOptions().RateLimiter.GetLimit() != uint(finalRate) {

78
		panic("wrong final rate limit")

79
	}

80
}

81


82
func verifyThreadsAndBulkSize(ne *nuclei.NucleiEngine, initialThreads, initialBulk, finalThreads, finalBulk int) {

83
	if ne.Options().TemplateThreads != initialThreads || ne.Options().BulkSize != initialBulk {

84
		panic("wrong initial standard concurrency")

85
	}

86
	time.Sleep(5 * time.Second)

87
	ne.Options().TemplateThreads = finalThreads

88
	ne.Options().BulkSize = finalBulk

89
	time.Sleep(20 * time.Second)

90
	if ne.Engine().GetWorkPool().InputPool(types.HTTPProtocol).Size != finalBulk || ne.Engine().WorkPool().Default.Size != finalThreads {

91
		log.Fatal("wrong final concurrency", ne.Engine().WorkPool().Default.Size, finalThreads, ne.Engine().GetWorkPool().InputPool(types.HTTPProtocol).Size, finalBulk)

92
	}

93
}

94


95
func verifyPayloadConcurrency(ne *nuclei.NucleiEngine, initialPayloadConcurrency, finalPayloadConcurrency int) {

96
	if ne.Options().PayloadConcurrency != initialPayloadConcurrency {

97
		panic("wrong initial payload concurrency")

98
	}

99
	time.Sleep(5 * time.Second)

100
	ne.Options().PayloadConcurrency = finalPayloadConcurrency

101
	time.Sleep(20 * time.Second)

102
	if ne.GetExecuterOptions().GetThreadsForNPayloadRequests(100, 0) != finalPayloadConcurrency {

103
		panic("wrong final payload concurrency")

104
	}

105
}

106


107