refactor: stable sorting (#3)

* refactor: simplify sorting

* stable sort by qtype

Signed-off-by: Bruno Schaatsbergen <[email protected]>

---------

Signed-off-by: Bruno Schaatsbergen <[email protected]>
Co-authored-by: Bruno Schaatsbergen <[email protected]>

Author: mmourick

cmd/printer.go

@@ -12,47 +12,43 @@ import (
 	"github.com/miekg/dns"
 )
 
-// printRecords prints the DNS records to the terminal.
-func printRecords(domainName string, messages []*dns.Msg) {
+// printRecord prints the DNS records to the terminal.
+func printRecord(domainName string, answer dns.RR) {
 	w := ansiterm.NewTabWriter(os.Stdout, 8, 8, 4, ' ', 0)
 	w.SetColorCapable(true)
 
 	domainColored := color.HiBlueString(domainName)
 
-	for _, msg := range messages {
-		for _, answer := range msg.Answer {
-			queryType := dns.TypeToString[answer.Header().Rrtype]
-			formattedTTL := color.HiMagentaString(formatTTL(answer.Header().Ttl))
+	recordType := dns.TypeToString[answer.Header().Rrtype]
+	formattedTTL := color.HiMagentaString(formatTTL(answer.Header().Ttl))
 
-			switch rec := answer.(type) {
-			case *dns.A:
-				printARecord(w, queryType, domainColored, formattedTTL, rec)
+	switch rec := answer.(type) {
+	case *dns.A:
+		printARecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.AAAA:
-				printAAAARecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.AAAA:
+		printAAAARecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.CNAME:
-				printCNAMERecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.CNAME:
+		printCNAMERecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.MX:
-				printMXRecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.MX:
+		printMXRecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.TXT:
-				printTXTRecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.TXT:
+		printTXTRecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.NS:
-				printNSRecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.NS:
+		printNSRecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.SOA:
-				printSOARecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.SOA:
+		printSOARecord(w, recordType, domainColored, formattedTTL, rec)
 
-			case *dns.PTR:
-				printPTRRecord(w, queryType, domainColored, formattedTTL, rec)
+	case *dns.PTR:
+		printPTRRecord(w, recordType, domainColored, formattedTTL, rec)
 
-			default:
-				fmt.Fprintf(os.Stderr, "Unknown record type: %s\n", queryType)
-			}
-		}
+	default:
+		fmt.Fprintf(os.Stderr, "Unknown record type: %s\n", recordType)
 	}
 
 	w.Flush()

cmd/root.go

@@ -86,19 +86,15 @@ var (
 				os.Exit(1)
 			}
 
-			// Sort the messages by resource record type
 			sort.SliceStable(messages, func(i, j int) bool {
-				if len(messages[i].Answer) == 0 {
-					return false
-				}
-				if len(messages[j].Answer) == 0 {
-					return true
-				}
-				return messages[i].Answer[0].Header().Rrtype < messages[j].Answer[0].Header().Rrtype
+				return messages[i].Question[0].Qtype < messages[j].Question[0].Qtype
 			})
 
-			// Print the records
-			printRecords(args[0], messages)
+			for _, m := range messages {
+				for _, record := range m.Answer {
+					printRecord(args[0], record)
+				}
+			}
 		},
 	}
 )