Skip to content

Commit b18412e

Browse files
caffixcaffix
committed
small improvements
1 parent 9ef5721 commit b18412e

File tree

3 files changed

+91
-106
lines changed

3 files changed

+91
-106
lines changed

‎engine/plugins/horizontals/contact.go‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -32,25 +32,12 @@ func (h *horContact) check(e *et.Event) error {
3232
return nil
3333
}
3434

35-
if ents, err := h.plugin.getContactRecordLocations(e, e.Entity); err == nil && len(ents) > 0 {
36-
for _, ent := range ents {
37-
if assocs := h.lookup(e, ent); len(assocs) > 0 {
38-
for _, assoc := range assocs {
39-
if assoc.ScopeChange {
40-
e.Session.Log().Info(assoc.Rationale)
41-
}
42-
}
43-
}
44-
}
45-
}
46-
47-
if ents, err := h.plugin.getContactRecordOrganizations(e, e.Entity); err == nil && len(ents) > 0 {
48-
for _, ent := range ents {
49-
if assocs := h.lookup(e, ent); len(assocs) > 0 {
50-
for _, assoc := range assocs {
51-
if assoc.ScopeChange {
52-
e.Session.Log().Info(assoc.Rationale)
53-
}
35+
orgs, locs := h.plugin.lookupContactRecordOrgsAndLocations(e.Session, e.Entity)
36+
for _, ent := range append(orgs, locs...) {
37+
if assocs := h.lookup(e, ent); len(assocs) > 0 {
38+
for _, assoc := range assocs {
39+
if assoc.ScopeChange {
40+
e.Session.Log().Info(assoc.Rationale)
5441
}
5542
}
5643
}

‎engine/plugins/horizontals/plugin.go‎

Lines changed: 62 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"log/slog"
1111
"time"
1212

13+
"github.com/caffix/stringset"
1314
"github.com/owasp-amass/amass/v5/engine/plugins/support"
1415
et "github.com/owasp-amass/amass/v5/engine/types"
1516
dbt "github.com/owasp-amass/asset-db/types"
@@ -177,24 +178,70 @@ func (h *horizPlugin) submitIPAddress(e *et.Event, asset *oamnet.IPAddress, src
177178
}
178179
}
179180

180-
func (h *horizPlugin) getContactRecordOrganizations(e *et.Event, cr *dbt.Entity) ([]*dbt.Entity, error) {
181-
since, err := support.TTLStartTime(e.Session.Config(),
181+
func (h *horizPlugin) lookupContactRecordOrgsAndLocations(sess et.Session, cr *dbt.Entity) ([]*dbt.Entity, []*dbt.Entity) {
182+
var orgents []*dbt.Entity
183+
184+
if ents, err := h.getContactRecordOrganizations(sess, cr); err == nil && len(ents) > 0 {
185+
for _, ent := range ents {
186+
if _, valid := ent.Asset.(*oamorg.Organization); valid {
187+
orgents = append(orgents, ent)
188+
}
189+
}
190+
}
191+
192+
set := stringset.New()
193+
defer set.Close()
194+
195+
var locents []*dbt.Entity
196+
for _, o := range orgents {
197+
if ents, err := h.getOrganizationLocations(sess, o); err == nil && len(ents) > 0 {
198+
for _, ent := range ents {
199+
if set.Has(ent.ID) {
200+
continue
201+
}
202+
203+
if _, valid := ent.Asset.(*oamcon.Location); valid {
204+
set.Insert(ent.ID)
205+
locents = append(locents, ent)
206+
}
207+
}
208+
}
209+
}
210+
211+
if ents, err := h.getContactRecordLocations(sess, cr); err == nil && len(ents) > 0 {
212+
for _, ent := range ents {
213+
if set.Has(ent.ID) {
214+
continue
215+
}
216+
217+
if _, valid := ent.Asset.(*oamcon.Location); valid {
218+
set.Insert(ent.ID)
219+
locents = append(locents, ent)
220+
}
221+
}
222+
}
223+
224+
return orgents, locents
225+
}
226+
227+
func (h *horizPlugin) getContactRecordOrganizations(sess et.Session, cr *dbt.Entity) ([]*dbt.Entity, error) {
228+
since, err := support.TTLStartTime(sess.Config(),
182229
string(oam.ContactRecord), string(oam.Organization), h.name)
183230
if err != nil {
184231
return nil, err
185232
}
186233

187-
ctx, cancel := context.WithTimeout(e.Session.Ctx(), 5*time.Second)
234+
ctx, cancel := context.WithTimeout(sess.Ctx(), 5*time.Second)
188235
defer cancel()
189236

190-
edges, err := e.Session.DB().OutgoingEdges(ctx, cr, since, "organization")
237+
edges, err := sess.DB().OutgoingEdges(ctx, cr, since, "organization")
191238
if err != nil || len(edges) == 0 {
192239
return nil, errors.New("zero organizations found")
193240
}
194241

195242
var results []*dbt.Entity
196243
for _, edge := range edges {
197-
to, err := e.Session.DB().FindEntityById(ctx, edge.ToEntity.ID)
244+
to, err := sess.DB().FindEntityById(ctx, edge.ToEntity.ID)
198245
if err != nil {
199246
continue
200247
}
@@ -210,24 +257,24 @@ func (h *horizPlugin) getContactRecordOrganizations(e *et.Event, cr *dbt.Entity)
210257
return results, nil
211258
}
212259

213-
func (h *horizPlugin) getContactRecordLocations(e *et.Event, cr *dbt.Entity) ([]*dbt.Entity, error) {
214-
since, err := support.TTLStartTime(e.Session.Config(),
260+
func (h *horizPlugin) getContactRecordLocations(sess et.Session, cr *dbt.Entity) ([]*dbt.Entity, error) {
261+
since, err := support.TTLStartTime(sess.Config(),
215262
string(oam.ContactRecord), string(oam.Location), h.name)
216263
if err != nil {
217264
return nil, err
218265
}
219266

220-
ctx, cancel := context.WithTimeout(e.Session.Ctx(), 5*time.Second)
267+
ctx, cancel := context.WithTimeout(sess.Ctx(), 5*time.Second)
221268
defer cancel()
222269

223-
edges, err := e.Session.DB().OutgoingEdges(ctx, cr, since, "location")
270+
edges, err := sess.DB().OutgoingEdges(ctx, cr, since, "location")
224271
if err != nil || len(edges) == 0 {
225272
return nil, errors.New("zero locations found")
226273
}
227274

228275
var results []*dbt.Entity
229276
for _, edge := range edges {
230-
to, err := e.Session.DB().FindEntityById(ctx, edge.ToEntity.ID)
277+
to, err := sess.DB().FindEntityById(ctx, edge.ToEntity.ID)
231278
if err != nil {
232279
continue
233280
}
@@ -243,24 +290,24 @@ func (h *horizPlugin) getContactRecordLocations(e *et.Event, cr *dbt.Entity) ([]
243290
return results, nil
244291
}
245292

246-
func (h *horizPlugin) getOrganizationLocations(e *et.Event, o *dbt.Entity) ([]*dbt.Entity, error) {
247-
since, err := support.TTLStartTime(e.Session.Config(),
293+
func (h *horizPlugin) getOrganizationLocations(sess et.Session, o *dbt.Entity) ([]*dbt.Entity, error) {
294+
since, err := support.TTLStartTime(sess.Config(),
248295
string(oam.Organization), string(oam.Location), h.name)
249296
if err != nil {
250297
return nil, err
251298
}
252299

253-
ctx, cancel := context.WithTimeout(e.Session.Ctx(), 5*time.Second)
300+
ctx, cancel := context.WithTimeout(sess.Ctx(), 5*time.Second)
254301
defer cancel()
255302

256-
edges, err := e.Session.DB().OutgoingEdges(ctx, o, since, "hq_address", "location")
303+
edges, err := sess.DB().OutgoingEdges(ctx, o, since, "hq_address", "location")
257304
if err != nil || len(edges) == 0 {
258305
return nil, errors.New("zero locations found")
259306
}
260307

261308
var results []*dbt.Entity
262309
for _, edge := range edges {
263-
to, err := e.Session.DB().FindEntityById(ctx, edge.ToEntity.ID)
310+
to, err := sess.DB().FindEntityById(ctx, edge.ToEntity.ID)
264311
if err != nil {
265312
continue
266313
}

‎engine/plugins/horizontals/reg_records.go‎

Lines changed: 23 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,11 @@ import (
1010
"fmt"
1111
"time"
1212

13-
"github.com/caffix/stringset"
1413
"github.com/owasp-amass/amass/v5/engine/plugins/support"
1514
et "github.com/owasp-amass/amass/v5/engine/types"
1615
dbt "github.com/owasp-amass/asset-db/types"
1716
oam "github.com/owasp-amass/open-asset-model"
1817
oamcon "github.com/owasp-amass/open-asset-model/contact"
19-
oamorg "github.com/owasp-amass/open-asset-model/org"
2018
oamreg "github.com/owasp-amass/open-asset-model/registration"
2119
)
2220

@@ -49,7 +47,12 @@ func (h *horRegRec) check(e *et.Event) error {
4947
return fmt.Errorf("asset type not supported: %s", t)
5048
}
5149

52-
orgs, locs := h.lookupRegistrantOrgsAndLocations(e, rlabel)
50+
cr, err := h.getRegistrantContactRecord(e, rlabel)
51+
if err != nil {
52+
return nil
53+
}
54+
55+
orgs, locs := h.plugin.lookupContactRecordOrgsAndLocations(e.Session, cr)
5356
if len(orgs) == 0 && len(locs) == 0 {
5457
return nil
5558
}
@@ -65,58 +68,6 @@ func (h *horRegRec) check(e *et.Event) error {
6568
return nil
6669
}
6770

68-
func (h *horRegRec) lookupRegistrantOrgsAndLocations(e *et.Event, rlabel string) ([]*oamorg.Organization, []*oamcon.Location) {
69-
cr, err := h.getRegistrantContactRecord(e, rlabel)
70-
if err != nil {
71-
return nil, nil
72-
}
73-
74-
var orgents []*dbt.Entity
75-
var resorgs []*oamorg.Organization
76-
if ents, err := h.plugin.getContactRecordOrganizations(e, cr); err == nil && len(ents) > 0 {
77-
for _, ent := range ents {
78-
if o, valid := ent.Asset.(*oamorg.Organization); valid {
79-
resorgs = append(resorgs, o)
80-
orgents = append(orgents, ent)
81-
}
82-
}
83-
}
84-
85-
set := stringset.New()
86-
defer set.Close()
87-
88-
var reslocs []*oamcon.Location
89-
for _, o := range orgents {
90-
if ents, err := h.plugin.getOrganizationLocations(e, o); err == nil && len(ents) > 0 {
91-
for _, ent := range ents {
92-
if set.Has(ent.ID) {
93-
continue
94-
}
95-
96-
if loc, valid := ent.Asset.(*oamcon.Location); valid {
97-
set.Insert(ent.ID)
98-
reslocs = append(reslocs, loc)
99-
}
100-
}
101-
}
102-
}
103-
104-
if ents, err := h.plugin.getContactRecordLocations(e, cr); err == nil && len(ents) > 0 {
105-
for _, ent := range ents {
106-
if set.Has(ent.ID) {
107-
continue
108-
}
109-
110-
if loc, valid := ent.Asset.(*oamcon.Location); valid {
111-
set.Insert(ent.ID)
112-
reslocs = append(reslocs, loc)
113-
}
114-
}
115-
}
116-
117-
return resorgs, reslocs
118-
}
119-
12071
func (h *horRegRec) getRegistrantContactRecord(e *et.Event, label string) (*dbt.Entity, error) {
12172
since, err := support.TTLStartTime(e.Session.Config(),
12273
string(e.Entity.Asset.AssetType()), string(oam.ContactRecord), h.plugin.name)
@@ -143,29 +94,29 @@ func (h *horRegRec) getRegistrantContactRecord(e *et.Event, label string) (*dbt.
14394
return nil, errors.New("failed to extract the registrant ContactRecord entity")
14495
}
14596

146-
func (h *horRegRec) processAutnumRecord(e *et.Event, orgs []*oamorg.Organization, locs []*oamcon.Location) {
97+
func (h *horRegRec) processAutnumRecord(e *et.Event, orgs []*dbt.Entity, locs []*dbt.Entity) {
14798
// check if the autnum record / registered autonomous system is in scope
14899
if _, conf := e.Session.Scope().IsAssetInScope(e.Entity.Asset, 0); conf > 0 {
149100
for _, o := range orgs {
150-
e.Session.Scope().Add(o)
101+
e.Session.Scope().Add(o.Asset)
151102
}
152103
for _, loc := range locs {
153-
e.Session.Scope().Add(loc)
104+
e.Session.Scope().Add(loc.Asset)
154105
}
155106
return
156107
}
157108

158109
var found bool
159110
for _, o := range orgs {
160-
if _, conf := e.Session.Scope().IsAssetInScope(o, 0); conf > 0 {
111+
if _, conf := e.Session.Scope().IsAssetInScope(o.Asset, 0); conf > 0 {
161112
found = true
162113
break
163114
}
164115
}
165116

166117
if !found {
167118
for _, loc := range locs {
168-
if _, conf := e.Session.Scope().IsAssetInScope(loc, 0); conf > 0 {
119+
if _, conf := e.Session.Scope().IsAssetInScope(loc.Asset, 0); conf > 0 {
169120
found = true
170121
break
171122
}
@@ -178,50 +129,50 @@ func (h *horRegRec) processAutnumRecord(e *et.Event, orgs []*oamorg.Organization
178129
e.Session.Scope().AddASN(an.Number)
179130
}
180131
for _, o := range orgs {
181-
e.Session.Scope().Add(o)
132+
e.Session.Scope().Add(o.Asset)
182133
}
183134
for _, loc := range locs {
184-
e.Session.Scope().Add(loc)
135+
e.Session.Scope().Add(loc.Asset)
185136
}
186137
}
187138
}
188139

189-
func (h *horRegRec) processDomainRecord(e *et.Event, orgs []*oamorg.Organization, locs []*oamcon.Location) {
140+
func (h *horRegRec) processDomainRecord(e *et.Event, orgs []*dbt.Entity, locs []*dbt.Entity) {
190141
// check if the domain record / registered domain name is in scope
191142
if _, conf := e.Session.Scope().IsAssetInScope(e.Entity.Asset, 0); conf > 0 {
192143
for _, o := range orgs {
193-
e.Session.Scope().Add(o)
144+
e.Session.Scope().Add(o.Asset)
194145
}
195146
for _, loc := range locs {
196-
e.Session.Scope().Add(loc)
147+
e.Session.Scope().Add(loc.Asset)
197148
}
198149
return
199150
}
200151
}
201152

202-
func (h *horRegRec) processIPNetRecord(e *et.Event, orgs []*oamorg.Organization, locs []*oamcon.Location) {
153+
func (h *horRegRec) processIPNetRecord(e *et.Event, orgs []*dbt.Entity, locs []*dbt.Entity) {
203154
// check if the ipnet record / registered netblock is in scope
204155
if _, conf := e.Session.Scope().IsAssetInScope(e.Entity.Asset, 0); conf > 0 {
205156
for _, o := range orgs {
206-
e.Session.Scope().Add(o)
157+
e.Session.Scope().Add(o.Asset)
207158
}
208159
for _, loc := range locs {
209-
e.Session.Scope().Add(loc)
160+
e.Session.Scope().Add(loc.Asset)
210161
}
211162
return
212163
}
213164

214165
var found bool
215166
for _, o := range orgs {
216-
if _, conf := e.Session.Scope().IsAssetInScope(o, 0); conf > 0 {
167+
if _, conf := e.Session.Scope().IsAssetInScope(o.Asset, 0); conf > 0 {
217168
found = true
218169
break
219170
}
220171
}
221172

222173
if !found {
223174
for _, loc := range locs {
224-
if _, conf := e.Session.Scope().IsAssetInScope(loc, 0); conf > 0 {
175+
if _, conf := e.Session.Scope().IsAssetInScope(loc.Asset, 0); conf > 0 {
225176
found = true
226177
break
227178
}
@@ -234,10 +185,10 @@ func (h *horRegRec) processIPNetRecord(e *et.Event, orgs []*oamorg.Organization,
234185
e.Session.Scope().AddCIDR(iprec.CIDR.String())
235186
}
236187
for _, o := range orgs {
237-
e.Session.Scope().Add(o)
188+
e.Session.Scope().Add(o.Asset)
238189
}
239190
for _, loc := range locs {
240-
e.Session.Scope().Add(loc)
191+
e.Session.Scope().Add(loc.Asset)
241192
}
242193
}
243194
}

0 commit comments

Comments
 (0)