Merge pull request #53 from icann/ft-view_stats_json

Allow "view" command to output JSON.

Author: jschlyter

app/cmd/view.go

@@ -3,7 +3,10 @@
 package cmd
 
 import (
+	"bytes"
 	"dnsmag/internal"
+	"fmt"
+	"os"
 
 	"github.com/spf13/cobra"
 )
@@ -16,38 +19,76 @@ func newViewCmd() *cobra.Command {
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			stderr := cmd.ErrOrStderr()
+			stdout := cmd.OutOrStdout()
 
 			inputFile := args[0]
 
 			var (
 				verbose bool
+				json    bool
 				top     int
+				output  string
 			)
 
 			parseFlags(cmd, map[string]any{
 				"verbose": &verbose,
+				"json":    &json,
 				"top":     &top,
+				"output":  &output,
 			})
 
+			if verbose && json {
+				return fmt.Errorf("--verbose and --json are mutually exclusive")
+			}
+
+			cmd.SilenceUsage = true
+
 			seq := internal.NewDatasetSequence(top, nil)
 
 			if err := loadDatasets(cmd, seq, []string{inputFile}, verbose); err != nil {
-				cmd.SilenceUsage = true
 				return err
 			}
 
 			// Format and print the domain statistics
-			if err := internal.OutputDatasetStats(stderr, seq.Result, verbose); err != nil {
-				cmd.SilenceUsage = true
-				return err
+
+			var buf bytes.Buffer
+			if json {
+				if err := internal.OutputDatasetStatsJSON(&buf, seq.Result); err != nil {
+					return err
+				}
+			} else {
+				if err := internal.OutputDatasetStats(&buf, seq.Result, verbose); err != nil {
+					return err
+				}
+			}
+
+			// Write buffer to stderr (default), stdout or file
+			if output != "" && output != "-" {
+				// Write to file
+				// #nosec G306
+				if err := os.WriteFile(output, buf.Bytes(), 0o644); err != nil {
+					return fmt.Errorf("failed to write to %s: %w", output, err)
+				}
+			} else if output == "-" {
+				// Write to stdout
+				if _, err := stdout.Write(buf.Bytes()); err != nil {
+					return fmt.Errorf("failed to write to stdout: %w", err)
+				}
+			} else {
+				// Default: stderr
+				if _, err := stderr.Write(buf.Bytes()); err != nil {
+					return fmt.Errorf("failed to write to stderr: %w", err)
+				}
 			}
 
 			return nil
 		},
 	}
 
 	viewCmd.Flags().BoolP("verbose", "v", false, "Verbose output")
+	viewCmd.Flags().BoolP("json", "j", false, "JSON output")
 	viewCmd.Flags().IntP("top", "n", internal.DefaultDomainCount, "Number of top domains to display")
+	viewCmd.Flags().StringP("output", "o", "", "Output file (optional, use '-' for stdout, defaults to stderr)")
 
 	return viewCmd
 }

app/cmd/view_test.go

@@ -2,8 +2,12 @@ package cmd
 
 import (
 	"bytes"
+	"dnsmag/internal"
+	"encoding/json"
+	"fmt"
 	"os"
 	"regexp"
+	"strings"
 	"testing"
 )
 
@@ -62,6 +66,194 @@ func TestViewCmd_Integration(t *testing.T) {
 	t.Logf("View command output:\n%s", output)
 }
 
+func TestViewCmd_JSON(t *testing.T) {
+	// Create temporary DNSMAG file
+	tmpDnsmag, err := os.CreateTemp("", "test_view_json_*.dnsmag")
+	if err != nil {
+		t.Fatalf("Failed to create temp DNSMAG file: %v", err)
+	}
+	defer os.Remove(tmpDnsmag.Name())
+	tmpDnsmag.Close()
+
+	// Collect data into temporary DNSMAG file
+	executeCollectAndVerify(t, []string{
+		"../../testdata/test1.pcap.gz",
+		"--output", tmpDnsmag.Name(),
+	}, 100, "PCAP")
+
+	// View the data with JSON output
+	viewCmd := newViewCmd()
+	viewCmd.SetArgs([]string{
+		tmpDnsmag.Name(),
+		"--json",
+	})
+
+	var viewBuf bytes.Buffer
+	viewCmd.SetOut(&viewBuf)
+	viewCmd.SetErr(&viewBuf)
+
+	err = viewCmd.Execute()
+	if err != nil {
+		t.Fatalf("View command with --json failed: %v\nOutput: %s", err, viewBuf.String())
+	}
+
+	output := viewBuf.String()
+
+	// Parse JSON output
+	var stats internal.DatasetStatsJSON
+	if err := json.Unmarshal(viewBuf.Bytes(), &stats); err != nil {
+		t.Fatalf("Failed to parse JSON output: %v\nOutput: %s", err, output)
+	}
+
+	// Verify ID is non-empty before overwriting
+	if stats.DatasetStatistics.ID == "" {
+		t.Error("Expected non-empty ID")
+	}
+
+	// Overwrite random ID field for comparison
+	stats.DatasetStatistics.ID = ""
+
+	expected := internal.DatasetStatsJSON{
+		DatasetStatistics: internal.DatasetStats{
+			ID:                 "",
+			Generator:          fmt.Sprintf("dnsmag %s", internal.Version),
+			Date:               "2000-01-01",
+			TotalUniqueClients: 70,
+			TotalQueryVolume:   100,
+			TotalDomainCount:   4,
+		},
+	}
+
+	if stats != expected {
+		t.Errorf("JSON output mismatch.\nGot:      %+v\nExpected: %+v", stats, expected)
+	}
+
+	t.Logf("View --json output:\n%s", output)
+}
+
+func TestViewCmd_OutputDestination(t *testing.T) {
+	// Create temporary DNSMAG file
+	tmpDnsmag, err := os.CreateTemp("", "test_view_output_*.dnsmag")
+	if err != nil {
+		t.Fatalf("Failed to create temp DNSMAG file: %v", err)
+	}
+	defer os.Remove(tmpDnsmag.Name())
+	tmpDnsmag.Close()
+
+	// Collect data into temporary DNSMAG file
+	executeCollectAndVerify(t, []string{
+		"../../testdata/test1.pcap.gz",
+		"--output", tmpDnsmag.Name(),
+	}, 100, "PCAP")
+
+	tests := []struct {
+		name         string
+		outputFlag   string
+		jsonFlag     bool
+		expectIn     string // "stdout", "stderr", or "file"
+		searchString string
+	}{
+		{
+			name:         "text output to stderr (default)",
+			outputFlag:   "",
+			jsonFlag:     false,
+			expectIn:     "stderr",
+			searchString: "Dataset statistics",
+		},
+		{
+			name:         "text output to file",
+			outputFlag:   "output.txt",
+			jsonFlag:     false,
+			expectIn:     "file",
+			searchString: "Dataset statistics",
+		},
+		{
+			name:         "text output to stdout",
+			outputFlag:   "-",
+			jsonFlag:     false,
+			expectIn:     "stdout",
+			searchString: "Dataset statistics",
+		},
+		{
+			name:         "JSON output to stderr (default)",
+			outputFlag:   "",
+			jsonFlag:     true,
+			expectIn:     "stderr",
+			searchString: "datasetStatistics",
+		},
+		{
+			name:         "JSON output to file",
+			outputFlag:   "output.json",
+			jsonFlag:     true,
+			expectIn:     "file",
+			searchString: "datasetStatistics",
+		},
+		{
+			name:         "JSON output to stdout",
+			outputFlag:   "-",
+			jsonFlag:     true,
+			expectIn:     "stdout",
+			searchString: "datasetStatistics",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			viewCmd := newViewCmd()
+			args := []string{tmpDnsmag.Name()}
+			if tt.jsonFlag {
+				args = append(args, "--json")
+			}
+
+			var outputPath string
+			if tt.outputFlag != "" {
+				if tt.expectIn == "file" {
+					tmpFile, err := os.CreateTemp("", tt.outputFlag)
+					if err != nil {
+						t.Fatalf("Failed to create temp file: %v", err)
+					}
+					outputPath = tmpFile.Name()
+					tmpFile.Close()
+					defer os.Remove(outputPath)
+					args = append(args, "--output", outputPath)
+				} else {
+					args = append(args, "--output", tt.outputFlag)
+				}
+			}
+			viewCmd.SetArgs(args)
+
+			var stdout, stderr bytes.Buffer
+			viewCmd.SetOut(&stdout)
+			viewCmd.SetErr(&stderr)
+
+			err := viewCmd.Execute()
+			if err != nil {
+				t.Fatalf("View command failed: %v\nStdout: %s\nStderr: %s", err, stdout.String(), stderr.String())
+			}
+
+			// Check output appears in expected location
+			switch tt.expectIn {
+			case "stdout":
+				if !strings.Contains(stdout.String(), tt.searchString) {
+					t.Errorf("Expected %q in stdout", tt.searchString)
+				}
+			case "stderr":
+				if !strings.Contains(stderr.String(), tt.searchString) {
+					t.Errorf("Expected %q in stderr", tt.searchString)
+				}
+			case "file":
+				fileContent, err := os.ReadFile(outputPath)
+				if err != nil {
+					t.Fatalf("Failed to read output file: %v", err)
+				}
+				if !strings.Contains(string(fileContent), tt.searchString) {
+					t.Errorf("Expected %q in file", tt.searchString)
+				}
+			}
+		})
+	}
+}
+
 func TestViewCmd_NonExistentFile(t *testing.T) {
 	viewCmd := newViewCmd()
 	viewCmd.SetArgs([]string{"non-existent.dnsmag"})

internal/stats.go

@@ -3,6 +3,7 @@
 package internal
 
 import (
+	"encoding/json"
 	"fmt"
 	"io"
 	"math"
@@ -190,6 +191,44 @@ func OutputDatasetStats(w io.Writer, dataset MagnitudeDataset, verbose bool) err
 	return printTable(w, table)
 }
 
+// DatasetStats represents the nested dataset statistics
+type DatasetStats struct {
+	ID                 string `json:"id"`
+	Generator          string `json:"generator"`
+	Date               string `json:"date"`
+	TotalUniqueClients uint64 `json:"totalUniqueClients"`
+	TotalQueryVolume   uint64 `json:"totalQueryVolume"`
+	TotalDomainCount   uint64 `json:"totalDomainCount"`
+}
+
+// DatasetStatsJSON represents the JSON output format for dataset statistics
+type DatasetStatsJSON struct {
+	DatasetStatistics DatasetStats `json:"datasetStatistics"`
+}
+
+// OutputDatasetStatsJSON formats and prints dataset statistics as JSON
+func OutputDatasetStatsJSON(w io.Writer, dataset MagnitudeDataset) error {
+	dateStr := ""
+	if dataset.Date != nil {
+		dateStr = dataset.Date.Format("2006-01-02")
+	}
+
+	stats := DatasetStatsJSON{
+		DatasetStatistics: DatasetStats{
+			ID:                 dataset.Identifier,
+			Generator:          dataset.Generator,
+			Date:               dateStr,
+			TotalUniqueClients: dataset.AllClientsCount,
+			TotalQueryVolume:   dataset.AllQueriesCount,
+			TotalDomainCount:   uint64(len(dataset.Domains)),
+		},
+	}
+
+	encoder := json.NewEncoder(w)
+	encoder.SetIndent("", "  ")
+	return encoder.Encode(stats)
+}
+
 // OutputCollectorStats formats and prints both dataset and timing statistics for collection operations
 func OutputCollectorStats(w io.Writer, collector *Collector, verbose bool) error {
 	if len(collector.filesLoaded) == 1 {