From 90924eb27ab3115974e6438e1adfc0fe465cee09 Mon Sep 17 00:00:00 2001
From: Karan Sharma <hello@mrkaran.dev>
Date: Sun, 13 Dec 2020 13:56:38 +0530
Subject: [PATCH] fix: output formatting

---
 TODO.md                 |  2 +-
 cmd/output.go           | 32 +++++++++++++++++++++++---------
 pkg/resolvers/system.go | 11 -----------
 3 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/TODO.md b/TODO.md
index cd0d3f5..208c79d 100644
--- a/TODO.md
+++ b/TODO.md
@@ -14,7 +14,7 @@
 ## CLI Features
 - [ ] `digfile`
 - [x] `ndots` support
-- [x] `search path` support
+- [x] `search list` support
 - [x] JSON output
 - [x] Colorized output
 - [x] Table output
diff --git a/cmd/output.go b/cmd/output.go
index 295807d..18d5d94 100644
--- a/cmd/output.go
+++ b/cmd/output.go
@@ -40,16 +40,15 @@ type JSONResponse struct {
 func (hub *Hub) outputJSON(out []Output, msgs []resolvers.Response) {
 	// get the questions
 	queries := make([]Query, 0, len(msgs))
-	for _, m := range msgs {
-		for _, ques := range m.Message.Question {
-			q := Query{
-				Name:  ques.Name,
-				Type:  dns.ClassToString[ques.Qtype],
-				Class: dns.ClassToString[ques.Qclass],
-			}
-			queries = append(queries, q)
+	for _, ques := range hub.Questions {
+		q := Query{
+			Name:  ques.Name,
+			Type:  dns.ClassToString[ques.Qtype],
+			Class: dns.ClassToString[ques.Qclass],
 		}
+		queries = append(queries, q)
 	}
+
 	resp := JSONResponse{
 		Response{
 			Output:  out,
@@ -73,6 +72,8 @@ func (hub *Hub) outputTerminal(out []Output) {
 	if hub.QueryFlags.DisplayTimeTaken {
 		header = append(header, "Time Taken")
 	}
+	table.SetAutoWrapText(false)
+	table.SetAutoFormatHeaders(true)
 	table.SetHeader(header)
 
 	for _, o := range out {
@@ -90,6 +91,10 @@ func (hub *Hub) outputTerminal(out []Output) {
 // on the output format specified displays the information.
 func (hub *Hub) Output(responses []resolvers.Response) {
 	out := collectOutput(responses)
+	if len(out) == 0 {
+		hub.Logger.Info("No records found")
+		hub.Logger.Exit(0)
+	}
 	if hub.QueryFlags.ShowJSON {
 		hub.outputJSON(out, responses)
 	} else {
@@ -98,7 +103,7 @@ func (hub *Hub) Output(responses []resolvers.Response) {
 }
 
 func collectOutput(responses []resolvers.Response) []Output {
-	out := make([]Output, 0, len(responses))
+	var out []Output
 	// gather Output from the DNS Messages
 	for _, r := range responses {
 		var addr string
@@ -106,7 +111,16 @@ func collectOutput(responses []resolvers.Response) []Output {
 			switch t := a.(type) {
 			case *dns.A:
 				addr = t.A.String()
+			case *dns.AAAA:
+				addr = t.AAAA.String()
+			case *dns.CNAME:
+				addr = t.Target
+			case *dns.MX:
+				addr = strconv.Itoa(int(t.Preference)) + " " + t.Mx
+			case *dns.SOA:
+				addr = t.String()
 			}
+
 			h := a.Header()
 			name := h.Name
 			qclass := dns.Class(h.Class).String()
diff --git a/pkg/resolvers/system.go b/pkg/resolvers/system.go
index 5c9d35e..7e52b5f 100644
--- a/pkg/resolvers/system.go
+++ b/pkg/resolvers/system.go
@@ -49,17 +49,6 @@ func NewSystemResolver(resolvFilePath string) (Resolver, error) {
 // It's possible to send multiple question in one message
 // but some nameservers are not able to
 func (s *SystemResolver) Lookup(questions []dns.Question) ([]Response, error) {
-	for _, q := range questions {
-		domains := s.config.NameList(q.Name)
-		for _, d := range domains {
-			ques := dns.Question{
-				Name:   d,
-				Qtype:  q.Qtype,
-				Qclass: q.Qclass,
-			}
-			questions = append(questions, ques)
-		}
-	}
 	var (
 		messages  = prepareMessages(questions)
 		responses []Response