Skip to content

Commit 4c58a63

Browse files
caffixcaffix
committed
finished implementing the HTTP client selection and session network semaphore
1 parent aed9f4e commit 4c58a63

30 files changed

+198
-146
lines changed

‎engine/plugins/api/aviato/company_enrich.go‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,15 @@ import (
1616

1717
"github.com/owasp-amass/amass/v5/engine/plugins/support"
1818
et "github.com/owasp-amass/amass/v5/engine/types"
19-
"github.com/owasp-amass/amass/v5/internal/net/http"
19+
amasshttp "github.com/owasp-amass/amass/v5/internal/net/http"
2020
dbt "github.com/owasp-amass/asset-db/types"
2121
oam "github.com/owasp-amass/open-asset-model"
22-
"github.com/owasp-amass/open-asset-model/general"
23-
"github.com/owasp-amass/open-asset-model/org"
22+
oamgen "github.com/owasp-amass/open-asset-model/general"
23+
oamorg "github.com/owasp-amass/open-asset-model/org"
2424
)
2525

2626
func (ce *companyEnrich) check(e *et.Event) error {
27-
oamid, ok := e.Entity.Asset.(*general.Identifier)
27+
oamid, ok := e.Entity.Asset.(*oamgen.Identifier)
2828
if !ok {
2929
return errors.New("failed to extract the Identifier asset")
3030
} else if oamid.Type != AviatoCompanyID {
@@ -81,7 +81,7 @@ func (ce *companyEnrich) lookup(e *et.Event, ident *dbt.Entity, since time.Time)
8181
continue
8282
}
8383
if a, err := e.Session.DB().FindEntityById(ctx, edge.FromEntity.ID); err == nil && a != nil {
84-
if _, ok := a.Asset.(*org.Organization); ok {
84+
if _, ok := a.Asset.(*oamorg.Organization); ok {
8585
return a
8686
}
8787
}
@@ -91,7 +91,7 @@ func (ce *companyEnrich) lookup(e *et.Event, ident *dbt.Entity, since time.Time)
9191
}
9292

9393
func (ce *companyEnrich) query(e *et.Event, ident *dbt.Entity, apikey []string) (*dbt.Entity, *companyEnrichResult) {
94-
oamid := e.Entity.Asset.(*general.Identifier)
94+
oamid := e.Entity.Asset.(*oamgen.Identifier)
9595

9696
orgent := ce.lookup(e, ident, time.Time{})
9797
if orgent == nil {
@@ -102,15 +102,19 @@ func (ce *companyEnrich) query(e *et.Event, ident *dbt.Entity, apikey []string)
102102

103103
var enrich *companyEnrichResult
104104
for _, key := range apikey {
105-
headers := http.Header{"Content-Type": []string{"application/json"}}
105+
headers := amasshttp.Header{"Content-Type": []string{"application/json"}}
106106
headers["Authorization"] = []string{"Bearer " + key}
107107

108108
_ = ce.plugin.rlimit.Wait(e.Session.Ctx())
109+
e.Session.NetSem().Acquire()
110+
109111
ctx, cancel := context.WithTimeout(e.Session.Ctx(), 20*time.Second)
110112
defer cancel()
111113

112114
u := fmt.Sprintf("https://data.api.aviato.co/company/enrich?id=%s", url.QueryEscape(oamid.ID))
113-
resp, err := http.RequestWebPage(ctx, &http.Request{URL: u, Header: headers})
115+
resp, err := amasshttp.RequestWebPage(ctx,
116+
e.Session.Clients().General, &amasshttp.Request{URL: u, Header: headers})
117+
e.Session.NetSem().Release()
114118
if err != nil {
115119
msg := fmt.Sprintf("failed to obtain the company enrich result for %s: %s", oamid.ID, err)
116120
e.Session.Log().Error(msg, slog.Group("plugin", "name", ce.plugin.name, "handler", ce.name))
@@ -146,7 +150,7 @@ func (ce *companyEnrich) query(e *et.Event, ident *dbt.Entity, apikey []string)
146150
}
147151

148152
func (ce *companyEnrich) store(e *et.Event, orgent *dbt.Entity, data *companyEnrichResult) {
149-
o := orgent.Asset.(*org.Organization)
153+
o := orgent.Asset.(*oamorg.Organization)
150154

151155
o.Active = false
152156
if strings.EqualFold(data.Status, "active") {
@@ -162,10 +166,10 @@ func (ce *companyEnrich) store(e *et.Event, orgent *dbt.Entity, data *companyEnr
162166
if o.LegalName == "" && data.LegalName != "" {
163167
o.LegalName = data.LegalName
164168

165-
oamid := &general.Identifier{
166-
UniqueID: fmt.Sprintf("%s:%s", general.LegalName, o.LegalName),
169+
oamid := &oamgen.Identifier{
170+
UniqueID: fmt.Sprintf("%s:%s", oamgen.LegalName, o.LegalName),
167171
ID: o.LegalName,
168-
Type: general.LegalName,
172+
Type: oamgen.LegalName,
169173
}
170174

171175
ident, err := e.Session.DB().CreateAsset(ctx, oamid)
@@ -175,7 +179,7 @@ func (ce *companyEnrich) store(e *et.Event, orgent *dbt.Entity, data *companyEnr
175179
return
176180
}
177181

178-
_, err = e.Session.DB().CreateEntityProperty(ctx, ident, &general.SourceProperty{
182+
_, err = e.Session.DB().CreateEntityProperty(ctx, ident, &oamgen.SourceProperty{
179183
Source: ce.name,
180184
Confidence: ce.plugin.source.Confidence,
181185
})
@@ -186,7 +190,7 @@ func (ce *companyEnrich) store(e *et.Event, orgent *dbt.Entity, data *companyEnr
186190
return
187191
}
188192

189-
err = ce.plugin.createRelation(ctx, e.Session, orgent, general.SimpleRelation{Name: "id"}, ident, ce.plugin.source.Confidence)
193+
err = ce.plugin.createRelation(ctx, e.Session, orgent, oamgen.SimpleRelation{Name: "id"}, ident, ce.plugin.source.Confidence)
190194
if err != nil {
191195
msg := fmt.Sprintf("failed to create the relation for %s: %s", o.LegalName, err)
192196
e.Session.Log().Error(msg, slog.Group("plugin", "name", ce.plugin.name, "handler", ce.name))
@@ -204,15 +208,15 @@ func (ce *companyEnrich) store(e *et.Event, orgent *dbt.Entity, data *companyEnr
204208
}
205209

206210
func (ce *companyEnrich) process(e *et.Event, ident, orgent *dbt.Entity) {
207-
o := orgent.Asset.(*org.Organization)
211+
o := orgent.Asset.(*oamorg.Organization)
208212

209213
_ = e.Dispatcher.DispatchEvent(&et.Event{
210214
Name: fmt.Sprintf("%s:%s", o.Name, o.ID),
211215
Entity: orgent,
212216
Session: e.Session,
213217
})
214218

215-
id := ident.Asset.(*general.Identifier)
219+
id := ident.Asset.(*oamgen.Identifier)
216220
e.Session.Log().Info("relationship discovered", "from", id.UniqueID, "relation", "id",
217221
"to", o.Name, slog.Group("plugin", "name", ce.plugin.name, "handler", ce.name))
218222
}

‎engine/plugins/api/aviato/company_rounds.go‎

Lines changed: 35 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -18,18 +18,18 @@ import (
1818
"github.com/owasp-amass/amass/v5/engine/plugins/support"
1919
"github.com/owasp-amass/amass/v5/engine/plugins/support/org"
2020
et "github.com/owasp-amass/amass/v5/engine/types"
21-
"github.com/owasp-amass/amass/v5/internal/net/http"
21+
amasshttp "github.com/owasp-amass/amass/v5/internal/net/http"
2222
dbt "github.com/owasp-amass/asset-db/types"
2323
oam "github.com/owasp-amass/open-asset-model"
24-
"github.com/owasp-amass/open-asset-model/account"
25-
"github.com/owasp-amass/open-asset-model/financial"
26-
"github.com/owasp-amass/open-asset-model/general"
24+
oamacct "github.com/owasp-amass/open-asset-model/account"
25+
oamfin "github.com/owasp-amass/open-asset-model/financial"
26+
oamgen "github.com/owasp-amass/open-asset-model/general"
2727
oamorg "github.com/owasp-amass/open-asset-model/org"
28-
"github.com/owasp-amass/open-asset-model/people"
28+
oampeop "github.com/owasp-amass/open-asset-model/people"
2929
)
3030

3131
func (cr *companyRounds) check(e *et.Event) error {
32-
oamid, ok := e.Entity.Asset.(*general.Identifier)
32+
oamid, ok := e.Entity.Asset.(*oamgen.Identifier)
3333
if !ok {
3434
return errors.New("failed to extract the Identifier asset")
3535
} else if oamid.Type != AviatoCompanyID {
@@ -107,7 +107,7 @@ func (cr *companyRounds) lookup(e *et.Event, ident *dbt.Entity, since time.Time)
107107
continue
108108
}
109109
if a, err := e.Session.DB().FindEntityById(ctx, edge.ToEntity.ID); err == nil && a != nil {
110-
if acc, ok := a.Asset.(*account.Account); ok && acc.Type == account.Checking && acc.Number == "default" {
110+
if acc, ok := a.Asset.(*oamacct.Account); ok && acc.Type == oamacct.Checking && acc.Number == "default" {
111111
accountents = append(accountents, a)
112112
}
113113
}
@@ -121,7 +121,7 @@ func (cr *companyRounds) lookup(e *et.Event, ident *dbt.Entity, since time.Time)
121121
continue
122122
}
123123
if a, err := e.Session.DB().FindEntityById(ctx, edge.FromEntity.ID); err == nil && a != nil {
124-
if _, ok := a.Asset.(*financial.FundsTransfer); ok {
124+
if _, ok := a.Asset.(*oamfin.FundsTransfer); ok {
125125
fundents = append(fundents, a)
126126
}
127127
}
@@ -133,7 +133,7 @@ func (cr *companyRounds) lookup(e *et.Event, ident *dbt.Entity, since time.Time)
133133
}
134134

135135
func (cr *companyRounds) query(e *et.Event, ident *dbt.Entity, apikey []string) []*dbt.Entity {
136-
oamid := e.Entity.Asset.(*general.Identifier)
136+
oamid := e.Entity.Asset.(*oamgen.Identifier)
137137

138138
orgent := cr.getAssociatedOrg(e, ident)
139139
if orgent == nil {
@@ -149,15 +149,18 @@ func (cr *companyRounds) query(e *et.Event, ident *dbt.Entity, apikey []string)
149149
loop:
150150
for _, key := range apikey {
151151
for ; page < total; page++ {
152-
headers := http.Header{"Content-Type": []string{"application/json"}}
152+
headers := amasshttp.Header{"Content-Type": []string{"application/json"}}
153153
headers["Authorization"] = []string{"Bearer " + key}
154154

155155
_ = cr.plugin.rlimit.Wait(e.Session.Ctx())
156+
e.Session.NetSem().Acquire()
157+
156158
ctx, cancel := context.WithTimeout(e.Session.Ctx(), 20*time.Second)
157159
defer cancel()
158160

159161
u := fmt.Sprintf("https://data.api.aviato.co/company/%s/funding-rounds?perPage=%d&page=%d", url.QueryEscape(oamid.ID), perPage, page)
160-
resp, err := http.RequestWebPage(ctx, &http.Request{URL: u, Header: headers})
162+
resp, err := amasshttp.RequestWebPage(ctx, e.Session.Clients().General, &amasshttp.Request{URL: u, Header: headers})
163+
e.Session.NetSem().Release()
161164
if err != nil {
162165
msg := fmt.Sprintf("failed to obtain the funding rounds for %s: %s", oamid.ID, err)
163166
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
@@ -247,7 +250,7 @@ func (cr *companyRounds) store(e *et.Event, ident, orgent *dbt.Entity, funds *co
247250
continue
248251
}
249252

250-
f := &financial.FundsTransfer{
253+
f := &oamfin.FundsTransfer{
251254
ID: fmt.Sprintf("%s:%s", round.Name, round.Stage),
252255
Amount: float64(round.MoneyRaised),
253256
Currency: "USD",
@@ -261,7 +264,7 @@ func (cr *companyRounds) store(e *et.Event, ident, orgent *dbt.Entity, funds *co
261264
continue
262265
}
263266

264-
_, err = e.Session.DB().CreateEntityProperty(ctx, fundent, &general.SourceProperty{
267+
_, err = e.Session.DB().CreateEntityProperty(ctx, fundent, &oamgen.SourceProperty{
265268
Source: cr.plugin.source.Name,
266269
Confidence: cr.plugin.source.Confidence,
267270
})
@@ -272,14 +275,14 @@ func (cr *companyRounds) store(e *et.Event, ident, orgent *dbt.Entity, funds *co
272275
}
273276

274277
if err := cr.plugin.createRelation(ctx, e.Session, fundent,
275-
general.SimpleRelation{Name: "recipient"}, orgacc, cr.plugin.source.Confidence); err != nil {
278+
oamgen.SimpleRelation{Name: "recipient"}, orgacc, cr.plugin.source.Confidence); err != nil {
276279
msg := fmt.Sprintf("failed to create the recipient relation for %s: %s", f.ID, err)
277280
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
278281
continue
279282
}
280283

281284
if err := cr.plugin.createRelation(ctx, e.Session, fundent,
282-
general.SimpleRelation{Name: "sender"}, seedacc, cr.plugin.source.Confidence); err != nil {
285+
oamgen.SimpleRelation{Name: "sender"}, seedacc, cr.plugin.source.Confidence); err != nil {
283286
msg := fmt.Sprintf("failed to create the sender relation for %s: %s", f.ID, err)
284287
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
285288
continue
@@ -301,7 +304,7 @@ func (cr *companyRounds) orgCheckingAccount(e *et.Event, orgent *dbt.Entity) *db
301304
if edges, err := e.Session.DB().OutgoingEdges(ctx, orgent, time.Time{}, "account"); err == nil {
302305
for _, edge := range edges {
303306
if a, err := e.Session.DB().FindEntityById(ctx, edge.ToEntity.ID); err == nil && a != nil {
304-
if acc, ok := a.Asset.(*account.Account); ok && acc.Type == account.Checking && acc.Number == "default" {
307+
if acc, ok := a.Asset.(*oamacct.Account); ok && acc.Type == oamacct.Checking && acc.Number == "default" {
305308
accountent = a
306309
break
307310
}
@@ -310,9 +313,9 @@ func (cr *companyRounds) orgCheckingAccount(e *et.Event, orgent *dbt.Entity) *db
310313
}
311314

312315
if accountent == nil {
313-
ent, err := e.Session.DB().CreateAsset(ctx, &account.Account{
316+
ent, err := e.Session.DB().CreateAsset(ctx, &oamacct.Account{
314317
ID: uuid.New().String(),
315-
Type: account.Checking,
318+
Type: oamacct.Checking,
316319
Number: "default",
317320
Balance: 0,
318321
Active: o.Active,
@@ -326,7 +329,7 @@ func (cr *companyRounds) orgCheckingAccount(e *et.Event, orgent *dbt.Entity) *db
326329
accountent = ent
327330
}
328331

329-
_, err := e.Session.DB().CreateEntityProperty(ctx, accountent, &general.SourceProperty{
332+
_, err := e.Session.DB().CreateEntityProperty(ctx, accountent, &oamgen.SourceProperty{
330333
Source: cr.plugin.source.Name,
331334
Confidence: cr.plugin.source.Confidence,
332335
})
@@ -337,7 +340,7 @@ func (cr *companyRounds) orgCheckingAccount(e *et.Event, orgent *dbt.Entity) *db
337340
}
338341

339342
if err := cr.plugin.createRelation(ctx, e.Session, orgent,
340-
general.SimpleRelation{Name: "account"}, accountent, cr.plugin.source.Confidence); err != nil {
343+
oamgen.SimpleRelation{Name: "account"}, accountent, cr.plugin.source.Confidence); err != nil {
341344
msg := fmt.Sprintf("failed to create the account relation for %s: %s", o.Name, err)
342345
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
343346
}
@@ -350,9 +353,9 @@ func (cr *companyRounds) createSeedAccount(e *et.Event, round *companyFundingRou
350353
defer cancel()
351354

352355
name := fmt.Sprintf("%s:%s", round.Name, round.Stage)
353-
accountent, err := e.Session.DB().CreateAsset(ctx, &account.Account{
356+
accountent, err := e.Session.DB().CreateAsset(ctx, &oamacct.Account{
354357
ID: name,
355-
Type: account.Checking,
358+
Type: oamacct.Checking,
356359
Number: "default",
357360
Balance: float64(round.MoneyRaised),
358361
Active: false,
@@ -363,7 +366,7 @@ func (cr *companyRounds) createSeedAccount(e *et.Event, round *companyFundingRou
363366
return nil
364367
}
365368

366-
_, err = e.Session.DB().CreateEntityProperty(ctx, accountent, &general.SourceProperty{
369+
_, err = e.Session.DB().CreateEntityProperty(ctx, accountent, &oamgen.SourceProperty{
367370
Source: cr.plugin.source.Name,
368371
Confidence: cr.plugin.source.Confidence,
369372
})
@@ -380,7 +383,7 @@ func (cr *companyRounds) createSeedAccount(e *et.Event, round *companyFundingRou
380383

381384
for _, investor := range investors {
382385
if err := cr.plugin.createRelation(ctx, e.Session, investor,
383-
general.SimpleRelation{Name: "account"}, accountent, cr.plugin.source.Confidence); err != nil {
386+
oamgen.SimpleRelation{Name: "account"}, accountent, cr.plugin.source.Confidence); err != nil {
384387
msg := fmt.Sprintf("failed to create the account relation for %s: %s", investor.ID, err)
385388
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
386389
}
@@ -400,7 +403,7 @@ func (cr *companyRounds) createOrgInvestors(e *et.Event, round *companyFundingRo
400403
defer cancel()
401404

402405
for _, investor := range round.CompanyInvestors {
403-
oamid := &general.Identifier{
406+
oamid := &oamgen.Identifier{
404407
UniqueID: fmt.Sprintf("%s:%s", AviatoCompanyID, investor.CompanyID),
405408
ID: investor.CompanyID,
406409
Type: AviatoCompanyID,
@@ -442,7 +445,7 @@ func (cr *companyRounds) createOrgInvestors(e *et.Event, round *companyFundingRo
442445
continue
443446
}
444447

445-
_, err = e.Session.DB().CreateEntityProperty(ctx, orgent, &general.SourceProperty{
448+
_, err = e.Session.DB().CreateEntityProperty(ctx, orgent, &oamgen.SourceProperty{
446449
Source: cr.plugin.source.Name,
447450
Confidence: cr.plugin.source.Confidence,
448451
})
@@ -453,7 +456,7 @@ func (cr *companyRounds) createOrgInvestors(e *et.Event, round *companyFundingRo
453456
}
454457

455458
if err := cr.plugin.createRelation(ctx, e.Session, orgent,
456-
general.SimpleRelation{Name: "id"}, ident, cr.plugin.source.Confidence); err != nil {
459+
oamgen.SimpleRelation{Name: "id"}, ident, cr.plugin.source.Confidence); err != nil {
457460
msg := fmt.Sprintf("failed to create the id relation for %s: %s", oamid.UniqueID, err)
458461
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
459462
}
@@ -480,7 +483,7 @@ func (cr *companyRounds) createPersonInvestors(e *et.Event, round *companyFundin
480483
defer cancel()
481484

482485
for _, investor := range round.PersonInvestors {
483-
oamid := &general.Identifier{
486+
oamid := &oamgen.Identifier{
484487
UniqueID: fmt.Sprintf("%s:%s", AviatoPersonID, investor.PersonID),
485488
ID: investor.PersonID,
486489
Type: AviatoPersonID,
@@ -501,7 +504,7 @@ func (cr *companyRounds) createPersonInvestors(e *et.Event, round *companyFundin
501504
continue
502505
}
503506
if a, err := e.Session.DB().FindEntityById(ctx, edge.FromEntity.ID); err == nil && a != nil {
504-
if _, ok := a.Asset.(*people.Person); ok {
507+
if _, ok := a.Asset.(*oampeop.Person); ok {
505508
personent = a
506509
break
507510
}
@@ -526,7 +529,7 @@ func (cr *companyRounds) createPersonInvestors(e *et.Event, round *companyFundin
526529
continue
527530
}
528531

529-
_, err = e.Session.DB().CreateEntityProperty(ctx, personent, &general.SourceProperty{
532+
_, err = e.Session.DB().CreateEntityProperty(ctx, personent, &oamgen.SourceProperty{
530533
Source: cr.plugin.source.Name,
531534
Confidence: cr.plugin.source.Confidence,
532535
})
@@ -537,7 +540,7 @@ func (cr *companyRounds) createPersonInvestors(e *et.Event, round *companyFundin
537540
}
538541

539542
if err := cr.plugin.createRelation(ctx, e.Session, personent,
540-
general.SimpleRelation{Name: "id"}, ident, cr.plugin.source.Confidence); err != nil {
543+
oamgen.SimpleRelation{Name: "id"}, ident, cr.plugin.source.Confidence); err != nil {
541544
msg := fmt.Sprintf("failed to create the id relation for %s: %s", oamid.UniqueID, err)
542545
e.Session.Log().Error(msg, slog.Group("plugin", "name", cr.plugin.name, "handler", cr.name))
543546
}
@@ -555,7 +558,7 @@ func (cr *companyRounds) createPersonInvestors(e *et.Event, round *companyFundin
555558

556559
func (cr *companyRounds) process(e *et.Event, fundents []*dbt.Entity) {
557560
for _, fund := range fundents {
558-
f := fund.Asset.(*financial.FundsTransfer)
561+
f := fund.Asset.(*oamfin.FundsTransfer)
559562

560563
_ = e.Dispatcher.DispatchEvent(&et.Event{
561564
Name: fmt.Sprintf("%f:%s", f.Amount, f.ID),

0 commit comments

Comments
 (0)