diff --git a/main.go b/main.go index 553961a..6625142 100644 --- a/main.go +++ b/main.go @@ -458,13 +458,22 @@ All long form (--) flags can be toggled with the dig-standard +[no]flag notation } } - entries = append(entries, &output.Entry{ + e := &output.Entry{ Queries: msgs, Replies: replies, Server: server, - Txp: txp, Time: time.Since(startTime), - }) + } + + if opts.ResolveIPs { + e.LoadPTRs(txp) + } + + entries = append(entries, e) + + if err := (*txp).Close(); err != nil { + return fmt.Errorf("closing transport: %s", err) + } } printer := output.Printer{ diff --git a/output/output.go b/output/output.go index b704919..cd97cc2 100644 --- a/output/output.go +++ b/output/output.go @@ -4,11 +4,12 @@ import ( "io" "time" + "github.com/natesales/q/transport" + "github.com/miekg/dns" log "github.com/sirupsen/logrus" "github.com/natesales/q/cli" - "github.com/natesales/q/transport" ) // Printer stores global options across multiple entries @@ -30,15 +31,12 @@ type Entry struct { // Time is the total time it took to query this server Time time.Duration - // Txp is the transport used to resolve IP addresses in A/AAAA records to their PTR records - Txp *transport.Transport `json:"-"` - PTRs map[string]string `json:"-"` // IP -> PTR value existingRRs map[string]bool } // LoadPTRs populates an entry's PTRs map with PTR values for all A/AAAA records -func (e *Entry) LoadPTRs() { +func (e *Entry) LoadPTRs(txp *transport.Transport) { // Initialize PTR cache if it doesn't exist if e.PTRs == nil { e.PTRs = make(map[string]string) @@ -66,7 +64,7 @@ func (e *Entry) LoadPTRs() { msg.SetQuestion(qname, dns.TypePTR) // Resolve qname and cache result - resp, err := (*e.Txp).Exchange(&msg) + resp, err := (*txp).Exchange(&msg) if err != nil { log.Debugf("error resolving PTR record: %s", err) } diff --git a/output/output_test.go b/output/output_test.go index f22fae3..96397c5 100644 --- a/output/output_test.go +++ b/output/output_test.go @@ -38,7 +38,6 @@ var entries = []*Entry{ Replies: replies(), Server: "192.0.2.10", Time: time.Second * 2, - Txp: nil, PTRs: nil, existingRRs: nil, }, diff --git a/output/pretty.go b/output/pretty.go index 13c5135..dd9f06a 100644 --- a/output/pretty.go +++ b/output/pretty.go @@ -47,42 +47,6 @@ func (p Printer) PrettyPrintNSID(entries []*Entry) { } } -// ptr resolves an IP address (not an arpa FQDN) to its PTR record -func (e *Entry) ptr(ip string) (string, error) { - // Initialize ptrCache if it doesn't exist - if e.PTRs == nil { - e.PTRs = make(map[string]string) - } - - // Return the result from cache if we already have it - if ptr, ok := e.PTRs[ip]; ok { - return ptr, nil - } - - // Create PTR query - qname, err := dns.ReverseAddr(ip) - if err != nil { - log.Fatalf("error reversing PTR record: %s", err) - } - msg := dns.Msg{} - msg.SetQuestion(qname, dns.TypePTR) - - // Resolve qname and cache result - resp, err := (*e.Txp).Exchange(&msg) - if err != nil { - return "", err - } - - // Cache and return - if len(resp.Answer) > 0 { - e.PTRs[ip] = resp.Answer[0].(*dns.PTR).Ptr - return e.PTRs[ip], nil - } - - // No value - return "", fmt.Errorf("no PTR record found for %s", ip) -} - // parseRR converts an RR into a pretty string and returns the qname, ttl, type, value, and whether to skip printing it because it's a duplicate func (e *Entry) parseRR(a dns.RR, opts *cli.Flags) *RR { // Initialize existingRRs map if it doesn't exist @@ -128,11 +92,7 @@ func (e *Entry) parseRR(a dns.RR, opts *cli.Flags) *RR { // Handle PTR resolution if opts.ResolveIPs && (a.Header().Rrtype == dns.TypeA || a.Header().Rrtype == dns.TypeAAAA) { - if ptr, err := e.ptr(valCopy); err == nil { - val += util.Color(util.ColorMagenta, fmt.Sprintf(" (%s)", ptr)) - } else { - log.Warnf("PTR resolution: %s", err) - } + val += util.Color(util.ColorMagenta, fmt.Sprintf(" (%s)", e.PTRs[valCopy])) } // Server suffix diff --git a/output/pretty_test.go b/output/pretty_test.go index 6d4e40f..eafcd67 100644 --- a/output/pretty_test.go +++ b/output/pretty_test.go @@ -10,16 +10,6 @@ import ( "github.com/natesales/q/util" ) -func TestOutputPrettyPTRCache(t *testing.T) { - var e Entry - e.PTRs = make(map[string]string) - e.PTRs["192.0.2.1"] = "example.com." - - ptr, err := e.ptr("192.0.2.1") - assert.Nil(t, err) - assert.Equal(t, "example.com.", ptr) -} - func TestOutputPrettyPrintColumn(t *testing.T) { var buf bytes.Buffer util.UseColor = false