Skip to content

Commit bc4c126

Browse files
caffixcaffix
committed
data sources can now be checked for requests they support
1 parent 60e573c commit bc4c126

File tree

5 files changed

+92
-36
lines changed

5 files changed

+92
-36
lines changed

‎datasrcs/scripting/script.go

Lines changed: 76 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,14 @@ import (
99
"errors"
1010
"fmt"
1111
"regexp"
12+
"sync"
1213

14+
"github.com/caffix/service"
15+
luaurl "github.com/cjoudrey/gluaurl"
1316
"github.com/owasp-amass/amass/v3/config"
1417
"github.com/owasp-amass/amass/v3/net/dns"
1518
"github.com/owasp-amass/amass/v3/requests"
1619
"github.com/owasp-amass/amass/v3/systems"
17-
"github.com/caffix/service"
18-
luaurl "github.com/cjoudrey/gluaurl"
1920
lua "github.com/yuin/gopher-lua"
2021
luajson "layeh.com/gopher-json"
2122
)
@@ -43,6 +44,7 @@ type Script struct {
4344
sys systems.System
4445
luaState *lua.LState
4546
cbs *callbacks
47+
cbsLock sync.Mutex
4648
subre *regexp.Regexp
4749
seconds int
4850
ctx context.Context
@@ -135,8 +137,10 @@ func (s *Script) newLuaState(cfg *config.Config) *lua.LState {
135137

136138
// Save references to the script functions that serve as callbacks for Amass events.
137139
func (s *Script) assignCallbacks() {
138-
L := s.luaState
140+
s.cbsLock.Lock()
141+
defer s.cbsLock.Unlock()
139142

143+
L := s.luaState
140144
s.cbs = &callbacks{
141145
Start: L.GetGlobal("start"),
142146
Stop: L.GetGlobal("stop"),
@@ -193,6 +197,41 @@ func (s *Script) OnStop() error {
193197
return nil
194198
}
195199

200+
// HandlesReq implements the Service interface.
201+
func (s *Script) HandlesReq(req interface{}) bool {
202+
s.cbsLock.Lock()
203+
defer s.cbsLock.Unlock()
204+
205+
var handles bool
206+
switch t := req.(type) {
207+
case *requests.DNSRequest:
208+
if s.cbs.Vertical.Type() != lua.LTNil && t != nil && t.Domain != "" {
209+
handles = true
210+
}
211+
case *requests.ResolvedRequest:
212+
if s.cbs.Resolved.Type() != lua.LTNil && t != nil && t.Name != "" && len(t.Records) > 0 {
213+
handles = true
214+
}
215+
case *requests.SubdomainRequest:
216+
if s.cbs.Subdomain.Type() != lua.LTNil && t != nil && t.Name != "" {
217+
handles = true
218+
}
219+
case *requests.AddrRequest:
220+
if s.cbs.Address.Type() != lua.LTNil && t != nil && t.Address != "" {
221+
handles = true
222+
}
223+
case *requests.ASNRequest:
224+
if s.cbs.Asn.Type() != lua.LTNil && t != nil && (t.Address != "" || t.ASN != 0) {
225+
handles = true
226+
}
227+
case *requests.WhoisRequest:
228+
if s.cbs.Horizontal.Type() != lua.LTNil {
229+
handles = true
230+
}
231+
}
232+
return handles
233+
}
234+
196235
func (s *Script) requests() {
197236
for {
198237
select {
@@ -283,41 +322,57 @@ func (s *Script) stopScript() {
283322
}
284323

285324
func (s *Script) dispatch(in interface{}) {
325+
s.cbsLock.Lock()
326+
286327
switch req := in.(type) {
287328
case *requests.DNSRequest:
288329
if s.cbs.Vertical.Type() != lua.LTNil && req != nil && req.Domain != "" {
330+
callback := s.cbs.Vertical
331+
s.cbsLock.Unlock()
289332
s.CheckRateLimit()
290-
s.dnsRequest(s.ctx, req)
333+
s.dnsRequest(s.ctx, callback, req)
291334
}
292335
case *requests.ResolvedRequest:
293336
if s.cbs.Resolved.Type() != lua.LTNil && req != nil && req.Name != "" && len(req.Records) > 0 {
337+
callback := s.cbs.Resolved
338+
s.cbsLock.Unlock()
294339
s.CheckRateLimit()
295-
s.resolvedRequest(s.ctx, req)
340+
s.resolvedRequest(s.ctx, callback, req)
296341
}
297342
case *requests.SubdomainRequest:
298343
if s.cbs.Subdomain.Type() != lua.LTNil && req != nil && req.Name != "" {
344+
callback := s.cbs.Subdomain
345+
s.cbsLock.Unlock()
299346
s.CheckRateLimit()
300-
s.subdomainRequest(s.ctx, req)
347+
s.subdomainRequest(s.ctx, callback, req)
301348
}
302349
case *requests.AddrRequest:
303350
if s.cbs.Address.Type() != lua.LTNil && req != nil && req.Address != "" {
351+
callback := s.cbs.Address
352+
s.cbsLock.Unlock()
304353
s.CheckRateLimit()
305-
s.addrRequest(s.ctx, req)
354+
s.addrRequest(s.ctx, callback, req)
306355
}
307356
case *requests.ASNRequest:
308357
if s.cbs.Asn.Type() != lua.LTNil && req != nil && (req.Address != "" || req.ASN != 0) {
358+
callback := s.cbs.Asn
359+
s.cbsLock.Unlock()
309360
s.CheckRateLimit()
310-
s.asnRequest(s.ctx, req)
361+
s.asnRequest(s.ctx, callback, req)
311362
}
312363
case *requests.WhoisRequest:
313364
if s.cbs.Horizontal.Type() != lua.LTNil {
365+
callback := s.cbs.Horizontal
366+
s.cbsLock.Unlock()
314367
s.CheckRateLimit()
315-
s.whoisRequest(s.ctx, req)
368+
s.whoisRequest(s.ctx, callback, req)
316369
}
370+
default:
371+
s.cbsLock.Unlock()
317372
}
318373
}
319374

320-
func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
375+
func (s *Script) dnsRequest(ctx context.Context, callback lua.LValue, req *requests.DNSRequest) {
321376
L := s.luaState
322377

323378
if contextExpired(ctx) {
@@ -327,7 +382,7 @@ func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
327382
s.sys.Config().Log.Printf("Querying %s for %s subdomains", s.String(), req.Domain)
328383

329384
err := L.CallByParam(lua.P{
330-
Fn: s.cbs.Vertical,
385+
Fn: callback,
331386
NRet: 0,
332387
Protect: true,
333388
}, s.contextToUserData(ctx), lua.LString(req.Domain))
@@ -336,7 +391,7 @@ func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
336391
}
337392
}
338393

339-
func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequest) {
394+
func (s *Script) resolvedRequest(ctx context.Context, callback lua.LValue, req *requests.ResolvedRequest) {
340395
L := s.luaState
341396

342397
if contextExpired(ctx) {
@@ -354,7 +409,7 @@ func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequ
354409
}
355410

356411
err := L.CallByParam(lua.P{
357-
Fn: s.cbs.Resolved,
412+
Fn: callback,
358413
NRet: 0,
359414
Protect: true,
360415
}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LString(req.Domain), records)
@@ -363,15 +418,15 @@ func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequ
363418
}
364419
}
365420

366-
func (s *Script) subdomainRequest(ctx context.Context, req *requests.SubdomainRequest) {
421+
func (s *Script) subdomainRequest(ctx context.Context, callback lua.LValue, req *requests.SubdomainRequest) {
367422
L := s.luaState
368423

369424
if contextExpired(ctx) {
370425
return
371426
}
372427

373428
err := L.CallByParam(lua.P{
374-
Fn: s.cbs.Subdomain,
429+
Fn: callback,
375430
NRet: 0,
376431
Protect: true,
377432
}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LString(req.Domain), lua.LNumber(req.Times))
@@ -380,15 +435,15 @@ func (s *Script) subdomainRequest(ctx context.Context, req *requests.SubdomainRe
380435
}
381436
}
382437

383-
func (s *Script) addrRequest(ctx context.Context, req *requests.AddrRequest) {
438+
func (s *Script) addrRequest(ctx context.Context, callback lua.LValue, req *requests.AddrRequest) {
384439
L := s.luaState
385440

386441
if contextExpired(ctx) {
387442
return
388443
}
389444

390445
err := L.CallByParam(lua.P{
391-
Fn: s.cbs.Address,
446+
Fn: callback,
392447
NRet: 0,
393448
Protect: true,
394449
}, s.contextToUserData(ctx), lua.LString(req.Address))
@@ -397,15 +452,15 @@ func (s *Script) addrRequest(ctx context.Context, req *requests.AddrRequest) {
397452
}
398453
}
399454

400-
func (s *Script) asnRequest(ctx context.Context, req *requests.ASNRequest) {
455+
func (s *Script) asnRequest(ctx context.Context, callback lua.LValue, req *requests.ASNRequest) {
401456
L := s.luaState
402457

403458
if contextExpired(ctx) {
404459
return
405460
}
406461

407462
err := L.CallByParam(lua.P{
408-
Fn: s.cbs.Asn,
463+
Fn: callback,
409464
NRet: 0,
410465
Protect: true,
411466
}, s.contextToUserData(ctx), lua.LString(req.Address), lua.LNumber(req.ASN))
@@ -414,15 +469,15 @@ func (s *Script) asnRequest(ctx context.Context, req *requests.ASNRequest) {
414469
}
415470
}
416471

417-
func (s *Script) whoisRequest(ctx context.Context, req *requests.WhoisRequest) {
472+
func (s *Script) whoisRequest(ctx context.Context, callback lua.LValue, req *requests.WhoisRequest) {
418473
L := s.luaState
419474

420475
if contextExpired(ctx) {
421476
return
422477
}
423478

424479
err := L.CallByParam(lua.P{
425-
Fn: s.cbs.Horizontal,
480+
Fn: callback,
426481
NRet: 0,
427482
Protect: true,
428483
}, s.contextToUserData(ctx), lua.LString(req.Domain))

‎enum/enum.go

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,14 @@ import (
88
"context"
99
"sync"
1010

11-
"github.com/owasp-amass/amass/v3/config"
12-
"github.com/owasp-amass/amass/v3/datasrcs"
13-
"github.com/owasp-amass/amass/v3/requests"
14-
"github.com/owasp-amass/amass/v3/systems"
1511
"github.com/caffix/netmap"
1612
"github.com/caffix/pipeline"
1713
"github.com/caffix/queue"
1814
"github.com/caffix/service"
15+
"github.com/owasp-amass/amass/v3/config"
16+
"github.com/owasp-amass/amass/v3/datasrcs"
17+
"github.com/owasp-amass/amass/v3/requests"
18+
"github.com/owasp-amass/amass/v3/systems"
1919
)
2020

2121
// Enumeration is the object type used to execute a DNS enumeration.
@@ -145,7 +145,7 @@ func (e *Enumeration) manageDataSrcRequests() {
145145
pending[src.String()] = false
146146
}
147147

148-
finished := make(chan string, len(e.srcs))
148+
finished := make(chan string, len(e.srcs)*2)
149149
requestsMap := make(map[string][]interface{})
150150
loop:
151151
for {
@@ -161,11 +161,13 @@ loop:
161161
}
162162

163163
for name := range nameToSrc {
164-
if len(requestsMap[name]) == 0 && !pending[name] {
165-
go e.fireRequest(nameToSrc[name], element, finished)
166-
pending[name] = true
167-
} else {
168-
requestsMap[name] = append(requestsMap[name], element)
164+
if src := nameToSrc[name]; src != nil && src.HandlesReq(element) {
165+
if len(requestsMap[name]) == 0 && !pending[name] {
166+
go e.fireRequest(src, element, finished)
167+
pending[name] = true
168+
} else {
169+
requestsMap[name] = append(requestsMap[name], element)
170+
}
169171
}
170172
}
171173
case name := <-finished:

‎enum/input.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,14 +189,13 @@ func (r *enumSource) Next(ctx context.Context) bool {
189189
r.markDone()
190190
return false
191191
case <-t.C:
192-
if !r.enum.requestsPending() && r.queue.Len() == 0 {
192+
if !r.enum.requestsPending() && r.pipeline.DataItemCount() <= 0 {
193193
r.markDone()
194194
return false
195195
}
196196
r.fillQueue()
197197
t.Reset(waitForDuration)
198198
case <-r.queue.Signal():
199-
t.Reset(waitForDuration)
200199
return true
201200
}
202201
}

‎go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ require (
77
github.com/caffix/netmap v0.1.0
88
github.com/caffix/pipeline v0.2.1
99
github.com/caffix/queue v0.1.4
10-
github.com/caffix/service v0.2.4
10+
github.com/caffix/service v0.3.0
1111
github.com/caffix/stringset v0.1.1
1212
github.com/cayleygraph/quad v1.2.4
1313
github.com/cjoudrey/gluaurl v0.0.0-20161028222611-31cbb9bef199

‎go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ github.com/caffix/pipeline v0.2.1 h1:sfkBebseEmIh39+wHiaWLwCirx4sfTofUuT7Do0tbM8
100100
github.com/caffix/pipeline v0.2.1/go.mod h1:0q0Dx1s1pIi7peIwz3aTIzQRmTEtH1n45NJyrTkiIHs=
101101
github.com/caffix/queue v0.1.4 h1:sQbFzwGaPM1tRnQHWCgHOwj7hLuhDQ3BhY1/1TFbBiE=
102102
github.com/caffix/queue v0.1.4/go.mod h1:l8Eg7UTUHTRlc5aQ37mRVjzLN6eC7hgwimN0pA4UHe8=
103-
github.com/caffix/service v0.2.4 h1:tTXdKua4dnPLsvO3V7eDZ40SaJ1BRRlR0uC32Xg8XYA=
104-
github.com/caffix/service v0.2.4/go.mod h1:rh+bjCHTYUfxPatYy8xWHypfr0s95UhqUMqW1jC9GxQ=
103+
github.com/caffix/service v0.3.0 h1:Sb0GVFaYnn7mJCWyfcGr4AumdoHKT9+7gn6A96U88eY=
104+
github.com/caffix/service v0.3.0/go.mod h1:rh+bjCHTYUfxPatYy8xWHypfr0s95UhqUMqW1jC9GxQ=
105105
github.com/caffix/stringset v0.1.1 h1:Tm4b7SBFAsRTBbBX90eP8xBv6BxSuU2w+6G/JNXtNpg=
106106
github.com/caffix/stringset v0.1.1/go.mod h1:9Ztc521vlcp8IWdtIowZyWbbddMKR9Rdr+d0pgnjcvk=
107107
github.com/casbin/casbin/v2 v2.37.0/go.mod h1:vByNa/Fchek0KZUgG5wEsl7iFsiviAYKRtgrQfcJqHg=

0 commit comments

Comments
 (0)