From 7df12b222994622b610e649a952a1a09a3f875a8 Mon Sep 17 00:00:00 2001
From: Karan Sharma <hello@mrkaran.dev>
Date: Sun, 13 Dec 2020 09:34:53 +0530
Subject: [PATCH] feat: Add JSON output

---
 cmd/cli.go    |   6 +++
 cmd/hub.go    |   1 +
 cmd/output.go | 122 +++++++++++++++++++++++++++++++++++++++++---------
 3 files changed, 107 insertions(+), 22 deletions(-)

diff --git a/cmd/cli.go b/cmd/cli.go
index c521c04..bbe9037 100644
--- a/cmd/cli.go
+++ b/cmd/cli.go
@@ -88,6 +88,12 @@ func main() {
 			Usage:       "Display how long it took for the response to arrive",
 			Destination: &hub.QueryFlags.DisplayTimeTaken,
 		},
+		&cli.BoolFlag{
+			Name:        "json",
+			Aliases:     []string{"J"},
+			Usage:       "Set the output format as JSON",
+			Destination: &hub.QueryFlags.ShowJSON,
+		},
 		&cli.BoolFlag{
 			Name:        "verbose",
 			Usage:       "Enable verbose logging",
diff --git a/cmd/hub.go b/cmd/hub.go
index 1d3ff48..e584f5b 100644
--- a/cmd/hub.go
+++ b/cmd/hub.go
@@ -29,6 +29,7 @@ type QueryFlags struct {
 	UseIPv4          bool
 	UseIPv6          bool
 	DisplayTimeTaken bool
+	ShowJSON         bool
 }
 
 // NewHub initializes an instance of Hub which holds app wide configuration.
diff --git a/cmd/output.go b/cmd/output.go
index 9f90a6e..ae33024 100644
--- a/cmd/output.go
+++ b/cmd/output.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"strconv"
@@ -11,10 +12,58 @@ import (
 	"github.com/olekukonko/tablewriter"
 )
 
-// Output takes a list of `dns.Answers` and based
-// on the output format specified displays the information.
-func (hub *Hub) Output(responses []resolvers.Response) {
-	// Create SprintXxx functions to mix strings with other non-colorized strings:
+// Output has a list of fields which are produced for the output
+type Output struct {
+	Name      string `json:"name"`
+	Type      string `json:"type"`
+	Class     string `json:"class"`
+	TTL       string `json:"ttl"`
+	Address   string `json:"address"`
+	TimeTaken string `json:"rtt"`
+}
+
+type Query struct {
+	Name  string `json:"name"`
+	Type  string `json:"type"`
+	Class string `json:"class"`
+}
+type Response struct {
+	Output  []Output `json:"answers"`
+	Queries []Query  `json:"queries"`
+}
+
+type JSONResponse struct {
+	Response `json:"responses"`
+}
+
+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)
+		}
+	}
+	resp := JSONResponse{
+		Response{
+			Output:  out,
+			Queries: queries,
+		},
+	}
+	res, err := json.Marshal(resp)
+	if err != nil {
+		hub.Logger.WithError(err).Error("unable to output data in JSON")
+		hub.Logger.Exit(-1)
+	}
+	fmt.Printf("%s", res)
+}
+
+func (hub *Hub) outputTerminal(out []Output) {
 	green := color.New(color.FgGreen).SprintFunc()
 	blue := color.New(color.FgBlue).SprintFunc()
 
@@ -25,25 +74,54 @@ func (hub *Hub) Output(responses []resolvers.Response) {
 	}
 	table.SetHeader(header)
 
-	for _, r := range responses {
-		var res string
-		for _, a := range r.Message.Answer {
-			switch t := a.(type) {
-			case *dns.A:
-				res = t.A.String()
-			}
-			h := a.Header()
-			name := green(h.Name)
-			qclass := dns.Class(h.Class).String()
-			ttl := strconv.FormatInt(int64(h.Ttl), 10) + "s"
-			qtype := blue(dns.Type(h.Rrtype).String())
-			output := []string{name, qtype, qclass, ttl, res}
-			// Print how long it took
-			if hub.QueryFlags.DisplayTimeTaken {
-				output = append(output, fmt.Sprintf("%dms", r.RTT.Milliseconds()))
-			}
-			table.Append(output)
+	for _, o := range out {
+		output := []string{green(o.Name), blue(o.Type), o.Class, o.TTL, o.Address}
+		// Print how long it took
+		if hub.QueryFlags.DisplayTimeTaken {
+			output = append(output, o.TimeTaken)
 		}
+		table.Append(output)
 	}
 	table.Render()
 }
+
+// Output takes a list of `dns.Answers` and based
+// on the output format specified displays the information.
+func (hub *Hub) Output(responses []resolvers.Response) {
+	out := collectOutput(responses)
+	if hub.QueryFlags.ShowJSON {
+		hub.outputJSON(out, responses)
+	} else {
+		hub.outputTerminal(out)
+	}
+}
+
+func collectOutput(responses []resolvers.Response) []Output {
+	out := make([]Output, 0, len(responses))
+	// gather Output from the DNS Messages
+	for _, r := range responses {
+		var addr string
+		for _, a := range r.Message.Answer {
+			switch t := a.(type) {
+			case *dns.A:
+				addr = t.A.String()
+			}
+			h := a.Header()
+			name := h.Name
+			qclass := dns.Class(h.Class).String()
+			ttl := strconv.FormatInt(int64(h.Ttl), 10) + "s"
+			qtype := dns.Type(h.Rrtype).String()
+			rtt := fmt.Sprintf("%dms", r.RTT.Milliseconds())
+			o := Output{
+				Name:      name,
+				Type:      qtype,
+				TTL:       ttl,
+				Class:     qclass,
+				Address:   addr,
+				TimeTaken: rtt,
+			}
+			out = append(out, o)
+		}
+	}
+	return out
+}