Skip to content

Commit

Permalink
data sources can now be checked for requests they support
Browse files Browse the repository at this point in the history
  • Loading branch information
@caffix
caffix committed May 9, 2023
1 parent 60e573c commit bc4c126
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 36 deletions.
97 changes: 76 additions & 21 deletions datasrcs/scripting/script.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,14 @@ import (
"errors"
"fmt"
"regexp"
"sync"

"github.com/caffix/service"
luaurl "github.com/cjoudrey/gluaurl"
"github.com/owasp-amass/amass/v3/config"
"github.com/owasp-amass/amass/v3/net/dns"
"github.com/owasp-amass/amass/v3/requests"
"github.com/owasp-amass/amass/v3/systems"
"github.com/caffix/service"
luaurl "github.com/cjoudrey/gluaurl"
lua "github.com/yuin/gopher-lua"
luajson "layeh.com/gopher-json"
)
Expand Down Expand Up @@ -43,6 +44,7 @@ type Script struct {
sys systems.System
luaState *lua.LState
cbs *callbacks
cbsLock sync.Mutex
subre *regexp.Regexp
seconds int
ctx context.Context
Expand Down Expand Up @@ -135,8 +137,10 @@ func (s *Script) newLuaState(cfg *config.Config) *lua.LState {

// Save references to the script functions that serve as callbacks for Amass events.
func (s *Script) assignCallbacks() {
L := s.luaState
s.cbsLock.Lock()
defer s.cbsLock.Unlock()

L := s.luaState
s.cbs = &callbacks{
Start: L.GetGlobal("start"),
Stop: L.GetGlobal("stop"),
Expand Down Expand Up @@ -193,6 +197,41 @@ func (s *Script) OnStop() error {
return nil
}

// HandlesReq implements the Service interface.
func (s *Script) HandlesReq(req interface{}) bool {
s.cbsLock.Lock()
defer s.cbsLock.Unlock()

var handles bool
switch t := req.(type) {
case *requests.DNSRequest:
if s.cbs.Vertical.Type() != lua.LTNil && t != nil && t.Domain != "" {
handles = true
}
case *requests.ResolvedRequest:
if s.cbs.Resolved.Type() != lua.LTNil && t != nil && t.Name != "" && len(t.Records) > 0 {
handles = true
}
case *requests.SubdomainRequest:
if s.cbs.Subdomain.Type() != lua.LTNil && t != nil && t.Name != "" {
handles = true
}
case *requests.AddrRequest:
if s.cbs.Address.Type() != lua.LTNil && t != nil && t.Address != "" {
handles = true
}
case *requests.ASNRequest:
if s.cbs.Asn.Type() != lua.LTNil && t != nil && (t.Address != "" || t.ASN != 0) {
handles = true
}
case *requests.WhoisRequest:
if s.cbs.Horizontal.Type() != lua.LTNil {
handles = true
}
}
return handles
}

func (s *Script) requests() {
for {
select {
Expand Down Expand Up @@ -283,41 +322,57 @@ func (s *Script) stopScript() {
}

func (s *Script) dispatch(in interface{}) {
s.cbsLock.Lock()

switch req := in.(type) {
case *requests.DNSRequest:
if s.cbs.Vertical.Type() != lua.LTNil && req != nil && req.Domain != "" {
callback := s.cbs.Vertical
s.cbsLock.Unlock()
s.CheckRateLimit()
s.dnsRequest(s.ctx, req)
s.dnsRequest(s.ctx, callback, req)
}
case *requests.ResolvedRequest:
if s.cbs.Resolved.Type() != lua.LTNil && req != nil && req.Name != "" && len(req.Records) > 0 {
callback := s.cbs.Resolved
s.cbsLock.Unlock()
s.CheckRateLimit()
s.resolvedRequest(s.ctx, req)
s.resolvedRequest(s.ctx, callback, req)
}
case *requests.SubdomainRequest:
if s.cbs.Subdomain.Type() != lua.LTNil && req != nil && req.Name != "" {
callback := s.cbs.Subdomain
s.cbsLock.Unlock()
s.CheckRateLimit()
s.subdomainRequest(s.ctx, req)
s.subdomainRequest(s.ctx, callback, req)
}
case *requests.AddrRequest:
if s.cbs.Address.Type() != lua.LTNil && req != nil && req.Address != "" {
callback := s.cbs.Address
s.cbsLock.Unlock()
s.CheckRateLimit()
s.addrRequest(s.ctx, req)
s.addrRequest(s.ctx, callback, req)
}
case *requests.ASNRequest:
if s.cbs.Asn.Type() != lua.LTNil && req != nil && (req.Address != "" || req.ASN != 0) {
callback := s.cbs.Asn
s.cbsLock.Unlock()
s.CheckRateLimit()
s.asnRequest(s.ctx, req)
s.asnRequest(s.ctx, callback, req)
}
case *requests.WhoisRequest:
if s.cbs.Horizontal.Type() != lua.LTNil {
callback := s.cbs.Horizontal
s.cbsLock.Unlock()
s.CheckRateLimit()
s.whoisRequest(s.ctx, req)
s.whoisRequest(s.ctx, callback, req)
}
default:
s.cbsLock.Unlock()
}
}

func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
func (s *Script) dnsRequest(ctx context.Context, callback lua.LValue, req *requests.DNSRequest) {
L := s.luaState

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

err := L.CallByParam(lua.P{
Fn: s.cbs.Vertical,
Fn: callback,
NRet: 0,
Protect: true,
}, s.contextToUserData(ctx), lua.LString(req.Domain))
Expand All @@ -336,7 +391,7 @@ func (s *Script) dnsRequest(ctx context.Context, req *requests.DNSRequest) {
}
}

func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequest) {
func (s *Script) resolvedRequest(ctx context.Context, callback lua.LValue, req *requests.ResolvedRequest) {
L := s.luaState

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

err := L.CallByParam(lua.P{
Fn: s.cbs.Resolved,
Fn: callback,
NRet: 0,
Protect: true,
}, s.contextToUserData(ctx), lua.LString(req.Name), lua.LString(req.Domain), records)
Expand All @@ -363,15 +418,15 @@ func (s *Script) resolvedRequest(ctx context.Context, req *requests.ResolvedRequ
}
}

func (s *Script) subdomainRequest(ctx context.Context, req *requests.SubdomainRequest) {
func (s *Script) subdomainRequest(ctx context.Context, callback lua.LValue, req *requests.SubdomainRequest) {
L := s.luaState

if contextExpired(ctx) {
return
}

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

func (s *Script) addrRequest(ctx context.Context, req *requests.AddrRequest) {
func (s *Script) addrRequest(ctx context.Context, callback lua.LValue, req *requests.AddrRequest) {
L := s.luaState

if contextExpired(ctx) {
return
}

err := L.CallByParam(lua.P{
Fn: s.cbs.Address,
Fn: callback,
NRet: 0,
Protect: true,
}, s.contextToUserData(ctx), lua.LString(req.Address))
Expand All @@ -397,15 +452,15 @@ func (s *Script) addrRequest(ctx context.Context, req *requests.AddrRequest) {
}
}

func (s *Script) asnRequest(ctx context.Context, req *requests.ASNRequest) {
func (s *Script) asnRequest(ctx context.Context, callback lua.LValue, req *requests.ASNRequest) {
L := s.luaState

if contextExpired(ctx) {
return
}

err := L.CallByParam(lua.P{
Fn: s.cbs.Asn,
Fn: callback,
NRet: 0,
Protect: true,
}, s.contextToUserData(ctx), lua.LString(req.Address), lua.LNumber(req.ASN))
Expand All @@ -414,15 +469,15 @@ func (s *Script) asnRequest(ctx context.Context, req *requests.ASNRequest) {
}
}

func (s *Script) whoisRequest(ctx context.Context, req *requests.WhoisRequest) {
func (s *Script) whoisRequest(ctx context.Context, callback lua.LValue, req *requests.WhoisRequest) {
L := s.luaState

if contextExpired(ctx) {
return
}

err := L.CallByParam(lua.P{
Fn: s.cbs.Horizontal,
Fn: callback,
NRet: 0,
Protect: true,
}, s.contextToUserData(ctx), lua.LString(req.Domain))
Expand Down
22 changes: 12 additions & 10 deletions enum/enum.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,14 @@ import (
"context"
"sync"

"github.com/owasp-amass/amass/v3/config"
"github.com/owasp-amass/amass/v3/datasrcs"
"github.com/owasp-amass/amass/v3/requests"
"github.com/owasp-amass/amass/v3/systems"
"github.com/caffix/netmap"
"github.com/caffix/pipeline"
"github.com/caffix/queue"
"github.com/caffix/service"
"github.com/owasp-amass/amass/v3/config"
"github.com/owasp-amass/amass/v3/datasrcs"
"github.com/owasp-amass/amass/v3/requests"
"github.com/owasp-amass/amass/v3/systems"
)

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

finished := make(chan string, len(e.srcs))
finished := make(chan string, len(e.srcs)*2)
requestsMap := make(map[string][]interface{})
loop:
for {
Expand All @@ -161,11 +161,13 @@ loop:
}

for name := range nameToSrc {
if len(requestsMap[name]) == 0 && !pending[name] {
go e.fireRequest(nameToSrc[name], element, finished)
pending[name] = true
} else {
requestsMap[name] = append(requestsMap[name], element)
if src := nameToSrc[name]; src != nil && src.HandlesReq(element) {
if len(requestsMap[name]) == 0 && !pending[name] {
go e.fireRequest(src, element, finished)
pending[name] = true
} else {
requestsMap[name] = append(requestsMap[name], element)
}
}
}
case name := <-finished:
Expand Down
3 changes: 1 addition & 2 deletions enum/input.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -189,14 +189,13 @@ func (r *enumSource) Next(ctx context.Context) bool {
r.markDone()
return false
case <-t.C:
if !r.enum.requestsPending() && r.queue.Len() == 0 {
if !r.enum.requestsPending() && r.pipeline.DataItemCount() <= 0 {
r.markDone()
return false
}
r.fillQueue()
t.Reset(waitForDuration)
case <-r.queue.Signal():
t.Reset(waitForDuration)
return true
}
}
Expand Down
2 changes: 1 addition & 1 deletion go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ require (
github.com/caffix/netmap v0.1.0
github.com/caffix/pipeline v0.2.1
github.com/caffix/queue v0.1.4
github.com/caffix/service v0.2.4
github.com/caffix/service v0.3.0
github.com/caffix/stringset v0.1.1
github.com/cayleygraph/quad v1.2.4
github.com/cjoudrey/gluaurl v0.0.0-20161028222611-31cbb9bef199
Expand Down
4 changes: 2 additions & 2 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,8 +100,8 @@ github.com/caffix/pipeline v0.2.1 h1:sfkBebseEmIh39+wHiaWLwCirx4sfTofUuT7Do0tbM8
github.com/caffix/pipeline v0.2.1/go.mod h1:0q0Dx1s1pIi7peIwz3aTIzQRmTEtH1n45NJyrTkiIHs=
github.com/caffix/queue v0.1.4 h1:sQbFzwGaPM1tRnQHWCgHOwj7hLuhDQ3BhY1/1TFbBiE=
github.com/caffix/queue v0.1.4/go.mod h1:l8Eg7UTUHTRlc5aQ37mRVjzLN6eC7hgwimN0pA4UHe8=
github.com/caffix/service v0.2.4 h1:tTXdKua4dnPLsvO3V7eDZ40SaJ1BRRlR0uC32Xg8XYA=
github.com/caffix/service v0.2.4/go.mod h1:rh+bjCHTYUfxPatYy8xWHypfr0s95UhqUMqW1jC9GxQ=
github.com/caffix/service v0.3.0 h1:Sb0GVFaYnn7mJCWyfcGr4AumdoHKT9+7gn6A96U88eY=
github.com/caffix/service v0.3.0/go.mod h1:rh+bjCHTYUfxPatYy8xWHypfr0s95UhqUMqW1jC9GxQ=
github.com/caffix/stringset v0.1.1 h1:Tm4b7SBFAsRTBbBX90eP8xBv6BxSuU2w+6G/JNXtNpg=
github.com/caffix/stringset v0.1.1/go.mod h1:9Ztc521vlcp8IWdtIowZyWbbddMKR9Rdr+d0pgnjcvk=
github.com/casbin/casbin/v2 v2.37.0/go.mod h1:vByNa/Fchek0KZUgG5wEsl7iFsiviAYKRtgrQfcJqHg=
Expand Down

0 comments on commit bc4c126

Please sign in to comment.