Skip to content

Commit

Permalink
working mxlookup
Browse files Browse the repository at this point in the history
  • Loading branch information
phillip-stephens committed Apr 30, 2024
1 parent 78cc043 commit 5e601f9
Show file tree
Hide file tree
Showing 8 changed files with 302 additions and 135 deletions.
2 changes: 1 addition & 1 deletion pkg/cmd/cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ var mxlookupCmd = &cobra.Command{
correspond with an exchange record.`,
Run: func(cmd *cobra.Command, args []string) {
GC.Module = strings.ToUpper("mxlookup")
//Run(GC, cmd.Flags())
Run(GC, cmd.Flags())
},
}

Expand Down
23 changes: 22 additions & 1 deletion pkg/cmd/modules.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,11 @@
package cmd

import "github.com/zmap/dns"
import (
log "github.com/sirupsen/logrus"
"github.com/spf13/pflag"
"github.com/zmap/dns"
"github.com/zmap/zdns/pkg/modules"
)

var module_to_type map[string]uint16

Expand Down Expand Up @@ -76,3 +81,19 @@ func init() {
"ANY": dns.TypeANY,
}
}

type moduleData struct {
MXLookup *modules.MXLookup
}

func populateModuleData(gc *CLIConf, flags *pflag.FlagSet) *moduleData {
modData := new(moduleData)
switch gc.Module {
case "MXLOOKUP":
modData.MXLookup = modules.Initialize(flags)
default:
log.Debug("nothing to be done for module instantiation")
}
return modData

}
36 changes: 23 additions & 13 deletions pkg/cmd/worker_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -254,9 +254,13 @@ func populateCLIConfig(gc *CLIConf, flags *pflag.FlagSet) *CLIConf {
gc.OutputGroups = append(gc.OutputGroups, gc.ResultVerbosity)
gc.OutputGroups = append(gc.OutputGroups, groups...)

// setup i/o
gc.InputHandler = iohandlers.NewFileInputHandler(gc.InputFilePath)
gc.OutputHandler = iohandlers.NewFileOutputHandler(gc.OutputFilePath)
// setup i/o if not specified
if gc.InputHandler == nil {
gc.InputHandler = iohandlers.NewFileInputHandler(gc.InputFilePath)
}
if gc.OutputHandler == nil {
gc.OutputHandler = iohandlers.NewFileOutputHandler(gc.OutputFilePath)
}
return gc
}

Expand All @@ -282,6 +286,7 @@ func populateResolverConfig(gc *CLIConf) *zdns.ResolverConfig {
func Run(gc CLIConf, flags *pflag.FlagSet) {
gc = *populateCLIConfig(&gc, flags)
resolverConfig := populateResolverConfig(&gc)
modData := populateModuleData(&gc, flags)
// DoLookup:
// - n threads that do processing from in and place results in out
// - process until inChan closes, then wg.done()
Expand Down Expand Up @@ -324,7 +329,7 @@ func Run(gc CLIConf, flags *pflag.FlagSet) {
// create shared cache for all threads to share
for i := 0; i < gc.Threads; i++ {
go func() {
err := doLookupWorker(&gc, resolverConfig, inChan, outChan, metaChan, &lookupWG)
err := doLookupWorker(&gc, resolverConfig, modData, inChan, outChan, metaChan, &lookupWG)
if err != nil {
log.Fatal("could not start lookup worker: %v", err)
}
Expand Down Expand Up @@ -371,7 +376,7 @@ func Run(gc CLIConf, flags *pflag.FlagSet) {
}

// doLookupWorker is a single worker thread that processes lookups from the input channel. It calls wg.Done when it is finished.
func doLookupWorker(gc *CLIConf, rc *zdns.ResolverConfig, input <-chan interface{}, output chan<- string, metaChan chan<- routineMetadata, wg *sync.WaitGroup) error {
func doLookupWorker(gc *CLIConf, rc *zdns.ResolverConfig, moduleData *moduleData, input <-chan interface{}, output chan<- string, metaChan chan<- routineMetadata, wg *sync.WaitGroup) error {
defer wg.Done()
resolver, err := zdns.InitResolver(rc)
if err != nil {
Expand Down Expand Up @@ -409,15 +414,20 @@ func doLookupWorker(gc *CLIConf, rc *zdns.ResolverConfig, input <-chan interface
}
res.Name = rawName
res.Class = dns.Class(gc.Class).String()
dnsType, ok := module_to_type[gc.Module]
if !ok {
log.Fatalf("Module %s not found in module_to_type map", gc.Module)
}

if rc.IsIterative {
innerRes, trace, status, err = resolver.IterativeLookup(&zdns.Question{Name: res.Name, Class: gc.Class, Type: dnsType})
} else if !rc.IsIterative {
innerRes, status, err = resolver.ExternalLookup(&zdns.Question{Name: res.Name, Class: gc.Class, Type: dnsType}, nameServer)
switch gc.Module {
case "MXLOOKUP":
innerRes, trace, status, err = moduleData.MXLookup.DoLookup(resolver, lookupName, nameServer)
default:
dnsType, ok := module_to_type[gc.Module]
if !ok {
log.Fatalf("Module %s not found in module_to_type map", gc.Module)
}
if rc.IsIterative {
innerRes, trace, status, err = resolver.IterativeLookup(&zdns.Question{Name: res.Name, Class: gc.Class, Type: dnsType})
} else if !rc.IsIterative {
innerRes, trace, status, err = resolver.ExternalLookup(&zdns.Question{Name: res.Name, Class: gc.Class, Type: dnsType}, nameServer)
}
}

res.Timestamp = time.Now().Format(gc.TimeFormat)
Expand Down
9 changes: 8 additions & 1 deletion pkg/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package main
import (
"fmt"
"github.com/zmap/dns"
"github.com/zmap/zdns/pkg/cmd"
"github.com/zmap/zdns/pkg/zdns"

log "github.com/sirupsen/logrus"
Expand Down Expand Up @@ -75,6 +76,12 @@ func library() {

}

func cli() {
cmd.Execute()

}

func main() {
library()
//library()
cli()
}
3 changes: 2 additions & 1 deletion pkg/modules/a_lookup.go
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,8 @@ func recursiveIPLookup(r *zdns.Resolver, name string, nameServer string, dnsType
if _, ok := candidateSet[name]; !ok {
var status zdns.Status
var err error
result, trace, status, err = r.Lookup(&zdns.Question{Name: name, Type: dnsType}, nameServer)
result, trace, status, err = r.ExternalLookup(&zdns.Question{Name: name, Type: dnsType, Class: dns.ClassINET}, nameServer)

if status != zdns.STATUS_NOERROR || err != nil {
return nil, trace, status, err
}
Expand Down
127 changes: 127 additions & 0 deletions pkg/modules/mx_lookup.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
package modules

import (
log "github.com/sirupsen/logrus"
"github.com/spf13/pflag"
"github.com/zmap/dns"
"github.com/zmap/zdns/cachehash"
"github.com/zmap/zdns/pkg/zdns"
"strings"
"sync"
)

type CachedAddresses struct {
IPv4Addresses []string
IPv6Addresses []string
}

type MXRecord struct {
Name string `json:"name" groups:"short,normal,long,trace"`
Type string `json:"type" groups:"short,normal,long,trace"`
Class string `json:"class" groups:"normal,long,trace"`
Preference uint16 `json:"preference" groups:"short,normal,long,trace"`
IPv4Addresses []string `json:"ipv4_addresses,omitempty" groups:"short,normal,long,trace"`
IPv6Addresses []string `json:"ipv6_addresses,omitempty" groups:"short,normal,long,trace"`
TTL uint32 `json:"ttl" groups:"ttl,normal,long,trace"`
}

type MXResult struct {
Servers []MXRecord `json:"exchanges" groups:"short,normal,long,trace"`
}

type MXLookup struct {
IPv4Lookup bool
IPv6Lookup bool
MXCacheSize int
CacheHash *cachehash.CacheHash
CHmu sync.Mutex
}

func Initialize(f *pflag.FlagSet) *MXLookup {

ipv4Lookup, err := f.GetBool("ipv4-lookup")
if err != nil {
panic(err)
}
ipv6Lookup, err := f.GetBool("ipv6-lookup")
if err != nil {
panic(err)
}
mxCacheSize, err := f.GetInt("mx-cache-size")
if err != nil {
panic(err)
}

mxLookup := new(MXLookup)
mxLookup.IPv4Lookup = ipv4Lookup
mxLookup.IPv6Lookup = ipv6Lookup
mxLookup.MXCacheSize = mxCacheSize
mxLookup.CacheHash = new(cachehash.CacheHash)
mxLookup.CacheHash.Init(mxCacheSize)
return mxLookup
}

func (mxLookup *MXLookup) LookupIPs(r *zdns.Resolver, name, nameServer string, lookupIpv4 bool, lookupIpv6 bool) (CachedAddresses, zdns.Trace) {
if mxLookup == nil {
log.Fatal("mxLookup is not initialized")
}
mxLookup.CHmu.Lock()
// XXX this should be changed to a miekglookup
res, found := mxLookup.CacheHash.Get(name)
mxLookup.CHmu.Unlock()
if found {
return res.(CachedAddresses), zdns.Trace{}
}
retv := CachedAddresses{}
result, trace, status, _ := DoTargetedLookup(r, name, nameServer, lookupIpv4, lookupIpv6)
if status == zdns.STATUS_NOERROR && result != nil {
retv.IPv4Addresses = result.IPv4Addresses
retv.IPv6Addresses = result.IPv6Addresses
}

mxLookup.CHmu.Lock()
mxLookup.CacheHash.Add(name, retv)
mxLookup.CHmu.Unlock()
return retv, trace
}

func (mxLookup *MXLookup) DoLookup(r *zdns.Resolver, name, nameServer string) (*MXResult, zdns.Trace, zdns.Status, error) {
retv := MXResult{Servers: []MXRecord{}}
res, trace, status, err := r.ExternalLookup(&zdns.Question{Name: name, Type: dns.TypeMX, Class: dns.ClassINET}, nameServer)
if status != zdns.STATUS_NOERROR || err != nil {
return nil, trace, status, err
}

lookupIpv4 := mxLookup.IPv4Lookup || !mxLookup.IPv6Lookup
lookupIpv6 := mxLookup.IPv6Lookup
for _, ans := range res.Answers {
if mxAns, ok := ans.(zdns.PrefAnswer); ok {
name = strings.TrimSuffix(mxAns.Answer.Answer, ".")
rec := MXRecord{TTL: mxAns.Ttl, Type: mxAns.Type, Class: mxAns.Class, Name: name, Preference: mxAns.Preference}
ips, secondTrace := mxLookup.LookupIPs(r, name, nameServer, lookupIpv4, lookupIpv6)
rec.IPv4Addresses = ips.IPv4Addresses
rec.IPv6Addresses = ips.IPv6Addresses
retv.Servers = append(retv.Servers, rec)
trace = append(trace, secondTrace...)
}
}
return &retv, trace, zdns.STATUS_NOERROR, nil
}

// TODO Phillip pull thesse CLI flags over
//func (s *GlobalLookupFactory) SetFlags(f *pflag.FlagSet) {
// // If there's an error, panic is appropriate since we should at least be getting the default here.
// var err error
// s.IPv4Lookup, err = f.GetBool("ipv4-lookup")
// if err != nil {
// panic(err)
// }
// s.IPv6Lookup, err = f.GetBool("ipv6-lookup")
// if err != nil {
// panic(err)
// }
// s.MXCacheSize, err = f.GetInt("mx-cache-size")
// if err != nil {
// panic(err)
// }
//}
Loading

0 comments on commit 5e601f9

Please sign in to comment.