Skip to content

Commit 77fe168

Browse files
caffixcaffix
committed
started to update the rate limiters
1 parent 2176edd commit 77fe168

34 files changed

+265
-203
lines changed

‎engine/plugins/api/binaryedge.go

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,20 +20,22 @@ import (
2020
dbt "github.com/owasp-amass/asset-db/types"
2121
oam "github.com/owasp-amass/open-asset-model"
2222
oamdns "github.com/owasp-amass/open-asset-model/dns"
23-
"go.uber.org/ratelimit"
23+
"golang.org/x/time/rate"
2424
)
2525

2626
type binaryEdge struct {
2727
name string
2828
log *slog.Logger
29-
rlimit ratelimit.Limiter
29+
rlimit *rate.Limiter
3030
source *et.Source
3131
}
3232

3333
func NewBinaryEdge() et.Plugin {
34+
limit := rate.Every(10 * time.Second)
35+
3436
return &binaryEdge{
3537
name: "BinaryEdge",
36-
rlimit: ratelimit.New(10, ratelimit.WithoutSlack),
38+
rlimit: rate.NewLimiter(limit, 1),
3739
source: &et.Source{
3840
Name: "BinaryEdge",
3941
Confidence: 80,
@@ -101,7 +103,7 @@ func (be *binaryEdge) check(e *et.Event) error {
101103
if support.AssetMonitoredWithinTTL(e.Session, e.Entity, be.source, since) {
102104
names = append(names, be.lookup(e, fqdn.Name, since)...)
103105
} else {
104-
names = append(names, be.query(e, fqdn.Name, be.source, keys)...)
106+
names = append(names, be.query(e, fqdn.Name, keys)...)
105107
support.MarkAssetMonitored(e.Session, e.Entity, be.source)
106108
}
107109

@@ -115,15 +117,15 @@ func (be *binaryEdge) lookup(e *et.Event, name string, since time.Time) []*dbt.E
115117
return support.SourceToAssetsWithinTTL(e.Session, name, string(oam.FQDN), be.source, since)
116118
}
117119

118-
func (be *binaryEdge) query(e *et.Event, name string, src *et.Source, keys []string) []*dbt.Entity {
120+
func (be *binaryEdge) query(e *et.Event, name string, keys []string) []*dbt.Entity {
119121
subs := stringset.New()
120122
defer subs.Close()
121123

122124
pagenum := 1
123125
loop:
124126
for _, key := range keys {
125127
for pagenum <= 500 {
126-
be.rlimit.Take()
128+
_ = be.rlimit.Wait(context.TODO())
127129
resp, err := http.RequestWebPage(context.TODO(), &http.Request{
128130
Header: http.Header{"X-KEY": []string{key}},
129131
URL: "https://api.binaryedge.io/v2/query/domains/subdomain/" + name + "?page=" + strconv.Itoa(pagenum),

‎engine/plugins/api/chaos.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,20 +19,22 @@ import (
1919
dbt "github.com/owasp-amass/asset-db/types"
2020
oam "github.com/owasp-amass/open-asset-model"
2121
oamdns "github.com/owasp-amass/open-asset-model/dns"
22-
"go.uber.org/ratelimit"
22+
"golang.org/x/time/rate"
2323
)
2424

2525
type chaos struct {
2626
name string
2727
log *slog.Logger
28-
rlimit ratelimit.Limiter
28+
rlimit *rate.Limiter
2929
source *et.Source
3030
}
3131

3232
func NewChaos() et.Plugin {
33+
limit := rate.Every(10 * time.Second)
34+
3335
return &chaos{
3436
name: "Chaos",
35-
rlimit: ratelimit.New(10, ratelimit.WithoutSlack),
37+
rlimit: rate.NewLimiter(limit, 1),
3638
source: &et.Source{
3739
Name: "Chaos",
3840
Confidence: 80,
@@ -118,7 +120,7 @@ func (c *chaos) query(e *et.Event, name string, keys []string) []*dbt.Entity {
118120
var names []string
119121

120122
for _, key := range keys {
121-
c.rlimit.Take()
123+
_ = c.rlimit.Wait(context.TODO())
122124
resp, err := http.RequestWebPage(context.TODO(), &http.Request{
123125
URL: "https://dns.projectdiscovery.io/dns/" + name + "/subdomains",
124126
Header: http.Header{"Authorization": []string{key}},

‎engine/plugins/api/crtsh.go

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,20 +19,22 @@ import (
1919
dbt "github.com/owasp-amass/asset-db/types"
2020
oam "github.com/owasp-amass/open-asset-model"
2121
oamdns "github.com/owasp-amass/open-asset-model/dns"
22-
"go.uber.org/ratelimit"
22+
"golang.org/x/time/rate"
2323
)
2424

2525
type crtsh struct {
2626
name string
2727
log *slog.Logger
28-
rlimit ratelimit.Limiter
28+
rlimit *rate.Limiter
2929
source *et.Source
3030
}
3131

3232
func NewCrtsh() et.Plugin {
33+
limit := rate.Every(2 * time.Second)
34+
3335
return &crtsh{
3436
name: "crt.sh",
35-
rlimit: ratelimit.New(2, ratelimit.WithoutSlack),
37+
rlimit: rate.NewLimiter(limit, 1),
3638
source: &et.Source{
3739
Name: "HackerTarget",
3840
Confidence: 100,
@@ -86,9 +88,9 @@ func (c *crtsh) check(e *et.Event) error {
8688

8789
var names []*dbt.Entity
8890
if support.AssetMonitoredWithinTTL(e.Session, e.Entity, c.source, since) {
89-
names = append(names, c.lookup(e, fqdn.Name, c.source, since)...)
91+
names = append(names, c.lookup(e, fqdn.Name, since)...)
9092
} else {
91-
names = append(names, c.query(e, fqdn.Name, c.source)...)
93+
names = append(names, c.query(e, fqdn.Name)...)
9294
support.MarkAssetMonitored(e.Session, e.Entity, c.source)
9395
}
9496

@@ -98,12 +100,12 @@ func (c *crtsh) check(e *et.Event) error {
98100
return nil
99101
}
100102

101-
func (c *crtsh) lookup(e *et.Event, name string, src *et.Source, since time.Time) []*dbt.Entity {
103+
func (c *crtsh) lookup(e *et.Event, name string, since time.Time) []*dbt.Entity {
102104
return support.SourceToAssetsWithinTTL(e.Session, name, string(oam.FQDN), c.source, since)
103105
}
104106

105-
func (c *crtsh) query(e *et.Event, name string, src *et.Source) []*dbt.Entity {
106-
c.rlimit.Take()
107+
func (c *crtsh) query(e *et.Event, name string) []*dbt.Entity {
108+
_ = c.rlimit.Wait(context.TODO())
107109
resp, err := http.RequestWebPage(context.TODO(), &http.Request{
108110
URL: "https://crt.sh/?CN=" + name + "&output=json&exclude=expired",
109111
})

‎engine/plugins/api/dnsrepo.go

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,20 +20,22 @@ import (
2020
dbt "github.com/owasp-amass/asset-db/types"
2121
oam "github.com/owasp-amass/open-asset-model"
2222
oamdns "github.com/owasp-amass/open-asset-model/dns"
23-
"go.uber.org/ratelimit"
23+
"golang.org/x/time/rate"
2424
)
2525

2626
type dnsrepo struct {
2727
name string
2828
log *slog.Logger
29-
rlimit ratelimit.Limiter
29+
rlimit *rate.Limiter
3030
source *et.Source
3131
}
3232

3333
func NewDNSRepo() et.Plugin {
34+
limit := rate.Every(10 * time.Second)
35+
3436
return &dnsrepo{
3537
name: "DNSRepo",
36-
rlimit: ratelimit.New(10, ratelimit.WithoutSlack),
38+
rlimit: rate.NewLimiter(limit, 1),
3739
source: &et.Source{
3840
Name: "DNSRepo",
3941
Confidence: 80,
@@ -97,9 +99,9 @@ func (d *dnsrepo) check(e *et.Event) error {
9799

98100
var names []*dbt.Entity
99101
if support.AssetMonitoredWithinTTL(e.Session, e.Entity, d.source, since) {
100-
names = append(names, d.lookup(e, fqdn.Name, d.source, since)...)
102+
names = append(names, d.lookup(e, fqdn.Name, since)...)
101103
} else {
102-
names = append(names, d.query(e, fqdn.Name, d.source, keys)...)
104+
names = append(names, d.query(e, fqdn.Name, keys)...)
103105
support.MarkAssetMonitored(e.Session, e.Entity, d.source)
104106
}
105107

@@ -109,11 +111,11 @@ func (d *dnsrepo) check(e *et.Event) error {
109111
return nil
110112
}
111113

112-
func (d *dnsrepo) lookup(e *et.Event, name string, src *et.Source, since time.Time) []*dbt.Entity {
114+
func (d *dnsrepo) lookup(e *et.Event, name string, since time.Time) []*dbt.Entity {
113115
return support.SourceToAssetsWithinTTL(e.Session, name, string(oam.FQDN), d.source, since)
114116
}
115117

116-
func (d *dnsrepo) query(e *et.Event, name string, src *et.Source, keys []string) []*dbt.Entity {
118+
func (d *dnsrepo) query(e *et.Event, name string, keys []string) []*dbt.Entity {
117119
var names []string
118120

119121
for _, key := range keys {
@@ -127,7 +129,7 @@ func (d *dnsrepo) query(e *et.Event, name string, src *et.Source, keys []string)
127129
}
128130
}
129131

130-
d.rlimit.Take()
132+
_ = d.rlimit.Wait(context.TODO())
131133
if resp, err := http.RequestWebPage(context.TODO(), req); err == nil {
132134
if key == "" {
133135
names = append(names, d.parseHTML(e, resp.Body)...)

‎engine/plugins/api/gleif/fuzzy.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ func (fc *fuzzyCompletions) query(e *et.Event, orgent *dbt.Entity) *dbt.Entity {
8787
brand := support.ExtractBrandName(o.Name)
8888
u := "https://api.gleif.org/api/v1/fuzzycompletions?field=entity.legalName&q=" + url.QueryEscape(brand)
8989

90-
fc.plugin.rlimit.Take()
90+
_ = fc.plugin.rlimit.Wait(context.TODO())
9191
resp, err := http.RequestWebPage(context.TODO(), &http.Request{URL: u})
9292
if err != nil || resp.Body == "" {
9393
msg := fmt.Sprintf("Failed to obtain the LEI record for %s: %s", o.Name, err)

‎engine/plugins/api/gleif/lei_record.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import (
1818
)
1919

2020
func (g *gleif) getLEIRecord(id *general.Identifier) (*leiRecord, error) {
21-
g.rlimit.Take()
21+
_ = g.rlimit.Wait(context.TODO())
2222

2323
u := "https://api.gleif.org/api/v1/lei-records/" + id.ID
2424
resp, err := http.RequestWebPage(context.TODO(), &http.Request{URL: u})
@@ -36,7 +36,7 @@ func (g *gleif) getLEIRecord(id *general.Identifier) (*leiRecord, error) {
3636
}
3737

3838
func (g *gleif) getDirectParentRecord(id *general.Identifier) (*leiRecord, error) {
39-
g.rlimit.Take()
39+
_ = g.rlimit.Wait(context.TODO())
4040

4141
u := "https://api.gleif.org/api/v1/lei-records/" + id.ID + "/direct-parent"
4242
resp, err := http.RequestWebPage(context.TODO(), &http.Request{URL: u})
@@ -59,7 +59,7 @@ func (g *gleif) getDirectChildrenRecords(id *general.Identifier) ([]*leiRecord,
5959
last := 1
6060
link := "https://api.gleif.org/api/v1/lei-records/" + id.ID + "/direct-children"
6161
for i := 1; i <= last && link != ""; i++ {
62-
g.rlimit.Take()
62+
_ = g.rlimit.Wait(context.TODO())
6363

6464
resp, err := http.RequestWebPage(context.TODO(), &http.Request{URL: link})
6565
if err != nil || resp.StatusCode != 200 || resp.Body == "" {

‎engine/plugins/api/gleif/plugin.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,21 @@ package gleif
66

77
import (
88
"errors"
9+
"time"
910

1011
et "github.com/owasp-amass/amass/v4/engine/types"
1112
dbt "github.com/owasp-amass/asset-db/types"
1213
oam "github.com/owasp-amass/open-asset-model"
1314
"github.com/owasp-amass/open-asset-model/general"
14-
"go.uber.org/ratelimit"
15+
"golang.org/x/time/rate"
1516
)
1617

1718
func NewGLEIF() et.Plugin {
19+
limit := rate.Every(3 * time.Second)
20+
1821
return &gleif{
1922
name: "GLEIF",
20-
rlimit: ratelimit.New(3, ratelimit.WithoutSlack),
23+
rlimit: rate.NewLimiter(limit, 1),
2124
source: &et.Source{
2225
Name: "GLEIF",
2326
Confidence: 100,

‎engine/plugins/api/gleif/types.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,13 @@ import (
88
"log/slog"
99

1010
et "github.com/owasp-amass/amass/v4/engine/types"
11-
"go.uber.org/ratelimit"
11+
"golang.org/x/time/rate"
1212
)
1313

1414
type gleif struct {
1515
name string
1616
log *slog.Logger
17-
rlimit ratelimit.Limiter
17+
rlimit *rate.Limiter
1818
fuzzy *fuzzyCompletions
1919
related *relatedOrgs
2020
source *et.Source

‎engine/plugins/api/grepapp.go

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -23,20 +23,22 @@ import (
2323
dbt "github.com/owasp-amass/asset-db/types"
2424
oam "github.com/owasp-amass/open-asset-model"
2525
oamdns "github.com/owasp-amass/open-asset-model/dns"
26-
"go.uber.org/ratelimit"
26+
"golang.org/x/time/rate"
2727
)
2828

2929
type grepApp struct {
3030
name string
3131
log *slog.Logger
32-
rlimit ratelimit.Limiter
32+
rlimit *rate.Limiter
3333
source *et.Source
3434
}
3535

3636
func NewGrepApp() et.Plugin {
37+
limit := rate.Every(2 * time.Second)
38+
3739
return &grepApp{
3840
name: "Grep.App",
39-
rlimit: ratelimit.New(2, ratelimit.WithoutSlack),
41+
rlimit: rate.NewLimiter(limit, 1),
4042
source: &et.Source{
4143
Name: "Grep.App",
4244
Confidence: 50,
@@ -91,22 +93,22 @@ func (g *grepApp) check(e *et.Event) error {
9193

9294
var names []*dbt.Entity
9395
if support.AssetMonitoredWithinTTL(e.Session, e.Entity, g.source, since) {
94-
names = append(names, g.lookup(e, fqdn.Name, g.source, since)...)
96+
names = append(names, g.lookup(e, fqdn.Name, since)...)
9597
} else {
96-
names = append(names, g.query(e, fqdn.Name, g.source)...)
98+
names = append(names, g.query(e, fqdn.Name)...)
9799
support.MarkAssetMonitored(e.Session, e.Entity, g.source)
98100
}
99101

100102
if len(names) > 0 {
101-
g.process(e, names, g.source)
103+
g.process(e, names)
102104
}
103105
return nil
104106
}
105-
func (g *grepApp) lookup(e *et.Event, name string, src *et.Source, since time.Time) []*dbt.Entity {
107+
func (g *grepApp) lookup(e *et.Event, name string, since time.Time) []*dbt.Entity {
106108
return support.SourceToAssetsWithinTTL(e.Session, name, string(oam.Identifier), g.source, since)
107109
}
108110

109-
func (g *grepApp) query(e *et.Event, name string, src *et.Source) []*dbt.Entity {
111+
func (g *grepApp) query(e *et.Event, name string) []*dbt.Entity {
110112
newdlt := strings.ReplaceAll(name, ".", `\.`)
111113
escapedQuery := url.QueryEscape("([a-zA-Z0-9._-]+)@" + newdlt)
112114
re := regexp.MustCompile(`([a-zA-Z0-9._-]+)@` + newdlt)
@@ -116,7 +118,7 @@ func (g *grepApp) query(e *et.Event, name string, src *et.Source) []*dbt.Entity
116118

117119
cont := true
118120
for page := 1; cont; page++ {
119-
g.rlimit.Take()
121+
_ = g.rlimit.Wait(context.TODO())
120122
resp, err := http.RequestWebPage(context.TODO(), &http.Request{
121123
URL: fmt.Sprintf("https://grep.app/api/search?page=%s&q=%s&regexp=true", strconv.Itoa(page), escapedQuery),
122124
})
@@ -148,13 +150,13 @@ func (g *grepApp) query(e *et.Event, name string, src *et.Source) []*dbt.Entity
148150
}
149151
}
150152

151-
return g.store(e, emails.Slice(), g.source)
153+
return g.store(e, emails.Slice())
152154
}
153155

154-
func (g *grepApp) store(e *et.Event, emails []string, src *et.Source) []*dbt.Entity {
156+
func (g *grepApp) store(e *et.Event, emails []string) []*dbt.Entity {
155157
return support.StoreEmailsWithSource(e.Session, emails, g.source, g.name, g.name+"-Handler")
156158
}
157159

158-
func (g *grepApp) process(e *et.Event, assets []*dbt.Entity, src *et.Source) {
160+
func (g *grepApp) process(e *et.Event, assets []*dbt.Entity) {
159161
support.ProcessEmailsWithSource(e, assets, g.source)
160162
}

0 commit comments

Comments
 (0)