Merge pull request #54 from icann/ft-force_aggregate

add --force-date option to aggregate

Author: jschlyter

app/cmd/aggregate.go

@@ -5,6 +5,7 @@ package cmd
 import (
 	"dnsmag/internal"
 	"fmt"
+	"time"
 
 	"github.com/spf13/cobra"
 )
@@ -22,17 +23,19 @@ func newAggregateCmd() *cobra.Command {
 			timing := internal.NewTimingStats()
 
 			var (
-				top     int
-				verbose bool
-				quiet   bool
-				output  string
+				top       int
+				verbose   bool
+				quiet     bool
+				output    string
+				forceDate string
 			)
 
 			parseFlags(cmd, map[string]any{
-				"top":     &top,
-				"verbose": &verbose,
-				"quiet":   &quiet,
-				"output":  &output,
+				"top":        &top,
+				"verbose":    &verbose,
+				"quiet":      &quiet,
+				"output":     &output,
+				"force-date": &forceDate,
 			})
 
 			// Quiet and verbose flags are mutually exclusive
@@ -41,7 +44,17 @@ func newAggregateCmd() *cobra.Command {
 				return fmt.Errorf("conflicting flags: cannot use both --quiet and --verbose")
 			}
 
-			seq := internal.NewDatasetSequence(top, nil)
+			var forcedDate *time.Time
+			if forceDate != "" {
+				parsedDate, err := time.Parse("2006-01-02", forceDate)
+				if err != nil {
+					cmd.SilenceUsage = true
+					return fmt.Errorf("invalid date format for --force-date: %s (expected YYYY-MM-DD)", forceDate)
+				}
+				forcedDate = &parsedDate
+			}
+
+			seq := internal.NewDatasetSequence(top, forcedDate, forcedDate != nil, stderr)
 
 			// Load all provided DNSMAG files
 			err := loadDatasets(cmd, seq, args, verbose)
@@ -102,6 +115,7 @@ func newAggregateCmd() *cobra.Command {
 	aggregateCmd.Flags().IntP("top", "n", internal.DefaultDomainCount, "Minimum number of domains required in each dataset")
 	aggregateCmd.Flags().BoolP("verbose", "v", false, "Verbose output")
 	aggregateCmd.Flags().BoolP("quiet", "q", false, "Quiet mode")
+	aggregateCmd.Flags().String("force-date", "", "Force a specific date for the aggregated dataset (YYYY-MM-DD format)")
 
 	return aggregateCmd
 }

app/cmd/aggregate_test.go

@@ -209,3 +209,89 @@ func (sr *slowReader) Read(out []byte) (n int, err error) {
 
 	return numBytes, nil
 }
+
+func TestAggregateCmd_DifferentDates_WithoutForceDate(t *testing.T) {
+	file1, file2, cleanup := createDNSMagFilesWithDifferentDates(t, "2024-04-04", "2025-05-05")
+	defer cleanup()
+
+	// Try to aggregate the two DNSMAG files without --force-date
+	aggregateCmd := newAggregateCmd()
+	aggregateCmd.SetArgs([]string{
+		file1,
+		file2,
+	})
+
+	var aggregateBuf bytes.Buffer
+	aggregateCmd.SetOut(&aggregateBuf)
+	aggregateCmd.SetErr(&aggregateBuf)
+
+	err := aggregateCmd.Execute()
+	if err == nil {
+		t.Fatalf("Expected error when aggregating datasets with different dates, but got none. Output: %s", aggregateBuf.String())
+	}
+
+	// Verify the error message mentions date mismatch
+	output := aggregateBuf.String()
+	if !regexp.MustCompile(`date mismatch`).MatchString(err.Error()) {
+		t.Errorf("Expected 'date mismatch' error, got: %v\nOutput: %s", err, output)
+	}
+
+	t.Logf("Expected error occurred: %v", err)
+}
+
+func TestAggregateCmd_DifferentDates_WithForceDate(t *testing.T) {
+	file1, file2, cleanup := createDNSMagFilesWithDifferentDates(t, "2024-04-04", "2025-05-05")
+	defer cleanup()
+
+	// Aggregate the two DNSMAG files WITH --force-date
+	aggregateCmd := newAggregateCmd()
+	aggregateCmd.SetArgs([]string{
+		"--force-date", "2026-01-14",
+		file1,
+		file2,
+	})
+
+	var aggregateBuf bytes.Buffer
+	aggregateCmd.SetOut(&aggregateBuf)
+	aggregateCmd.SetErr(&aggregateBuf)
+
+	err := aggregateCmd.Execute()
+	if err != nil {
+		t.Fatalf("Aggregate command with --force-date failed: %v\nOutput: %s", err, aggregateBuf.String())
+	}
+
+	output := aggregateBuf.String()
+
+	// Verify the output contains expected patterns
+	expectedPatterns := []*regexp.Regexp{
+		regexp.MustCompile(`Aggregated statistics for 2 datasets:`),
+		regexp.MustCompile(`Date\s+:\s+2026-01-14`),                                           // The forced date
+		regexp.MustCompile(`Total queries\s+:\s+40`),                                          // 25 + 15
+		regexp.MustCompile(`Total domains\s+:\s+2`),                                           // com, org
+		regexp.MustCompile(`Warning: Overriding date 2025-05-05 with forced date 2026-01-14`), // Warning message
+	}
+
+	for _, pattern := range expectedPatterns {
+		if !pattern.MatchString(output) {
+			t.Errorf("Expected pattern %q not found in output:\n%s", pattern.String(), output)
+		}
+	}
+
+	t.Logf("Aggregate command with --force-date output:\n%s", output)
+}
+
+// createDNSMagFilesWithDifferentDates creates two temporary DNSMAG files with different dates
+// Returns the two file paths and a cleanup function
+func createDNSMagFilesWithDifferentDates(t *testing.T, date1, date2 string) (string, string, func()) {
+	t.Helper()
+
+	file1 := createDNSMAGFromCSV(t, "192.168.1.1,example.com,25", date1)
+	file2 := createDNSMAGFromCSV(t, "10.0.1.1,example.org,15", date2)
+
+	cleanup := func() {
+		os.Remove(file1)
+		os.Remove(file2)
+	}
+
+	return file1, file2, cleanup
+}

app/cmd/collect.go

@@ -13,7 +13,7 @@ import (
 func newCollectCmd() *cobra.Command {
 	collectCmd := &cobra.Command{
 		Use:   "collect <input-file> [input-file2] [input-file3...]",
-		Short: "Parse PCAP files and generate domain statistics",
+		Short: "Parse PCAP or CSV files and generate domain statistics",
 		Long: `Parse one or more PCAP files containing DNS traffic and generate domain statistics.
 Save them to a DNSMAG file (CBOR format).`,
 		Args: cobra.MinimumNArgs(1),

app/cmd/common_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"dnsmag/internal"
 	"fmt"
+	"os"
 	"regexp"
 	"strings"
 	"testing"
@@ -72,3 +73,30 @@ func loadDatasetFromCSV(csvData string, dateStr string, verbose bool) (*internal
 
 	return collector, nil
 }
+
+// createDNSMAGFromCSV creates a temporary DNSMAG file with the specified date and CSV data
+// Returns the file path. Caller is responsible for deleting the file.
+func createDNSMAGFromCSV(t *testing.T, csvData, date string) string {
+	t.Helper()
+
+	tmpFile, err := os.CreateTemp("", "test_*.dnsmag")
+	if err != nil {
+		t.Fatalf("Failed to create temp DNSMAG file: %v", err)
+	}
+	filePath := tmpFile.Name()
+	tmpFile.Close()
+
+	collector, err := loadDatasetFromCSV(csvData, date, false)
+	if err != nil {
+		os.Remove(filePath)
+		t.Fatalf("Failed to load CSV dataset: %v", err)
+	}
+
+	_, err = internal.WriteDNSMagFile(collector.Result, filePath, os.Stdout)
+	if err != nil {
+		os.Remove(filePath)
+		t.Fatalf("Failed to write DNSMAG file: %v", err)
+	}
+
+	return filePath
+}

app/cmd/report.go

@@ -47,7 +47,7 @@ func newReportCmd() *cobra.Command {
 				"verbose":     &verbose,
 			})
 
-			seq := internal.NewDatasetSequence(0, nil)
+			seq := internal.NewDatasetSequence(0, nil, false, stderr)
 
 			if err := loadDatasets(cmd, seq, []string{filename}, verbose); err != nil {
 				cmd.SilenceUsage = true

app/cmd/view.go

@@ -43,7 +43,7 @@ func newViewCmd() *cobra.Command {
 
 			cmd.SilenceUsage = true
 
-			seq := internal.NewDatasetSequence(top, nil)
+			seq := internal.NewDatasetSequence(top, nil, false, stderr)
 
 			if err := loadDatasets(cmd, seq, []string{inputFile}, verbose); err != nil {
 				return err

internal/store.go

@@ -91,13 +91,17 @@ type DatasetSequence struct {
 	numDomains int
 	Count      int
 	Result     MagnitudeDataset
+	forceDate  bool
+	logger     io.Writer
 }
 
-func NewDatasetSequence(numDomains int, date *time.Time) *DatasetSequence {
+func NewDatasetSequence(numDomains int, date *time.Time, forceDate bool, logger io.Writer) *DatasetSequence {
 	return &DatasetSequence{
 		numDomains: numDomains,
 		Count:      0,
 		Result:     newDataset(date),
+		forceDate:  forceDate,
+		logger:     logger,
 	}
 }
 
@@ -168,6 +172,17 @@ func (seq *DatasetSequence) LoadDNSMagSequenceFromReader(reader io.Reader, filen
 }
 
 func (seq *DatasetSequence) addDataset(dataset MagnitudeDataset) error {
+	// If forceDate is true and the dataset has a different date, log a warning and override it
+	if seq.forceDate && dataset.Date != nil && seq.Result.Date != nil {
+		if dataset.DateString() != seq.Result.DateString() {
+			if seq.logger != nil {
+				fmt.Fprintf(seq.logger, "Warning: Overriding date %s with forced date %s for dataset %s\n",
+					dataset.DateString(), seq.Result.DateString(), dataset.extraSourceFilename)
+			}
+			dataset.SetDate(&seq.Result.Date.Time)
+		}
+	}
+
 	if seq.Count == 0 {
 		seq.Result = dataset
 		seq.Count = 1

internal/store_test.go

@@ -46,7 +46,7 @@ func TestWriteAndLoadDNSMagFile_WriteLoadCycle(t *testing.T) {
 	}
 
 	// Test loading
-	seq := NewDatasetSequence(0, nil)
+	seq := NewDatasetSequence(0, nil, false, nil)
 	err = seq.LoadDNSMagFile(tmpFile.Name())
 	if err != nil {
 		t.Fatalf("LoadDNSMagFile failed: %v", err)
@@ -90,7 +90,7 @@ func TestWriteDNSMagFile_CreateError(t *testing.T) {
 }
 
 func TestLoadDNSMagFile_FileNotFound(t *testing.T) {
-	seq := NewDatasetSequence(0, nil)
+	seq := NewDatasetSequence(0, nil, false, nil)
 	err := seq.LoadDNSMagFile("non-existent.dnsmag")
 	if err == nil {
 		t.Error("Expected error when loading non-existent file, got nil")
@@ -111,7 +111,7 @@ func TestLoadDNSMagFile_InvalidFormat(t *testing.T) {
 	}
 	tmpFile.Close()
 
-	seq := NewDatasetSequence(0, nil)
+	seq := NewDatasetSequence(0, nil, false, nil)
 	err = seq.LoadDNSMagFile(tmpFile.Name())
 	if err == nil {
 		t.Error("Expected error when loading invalid CBOR file, got nil")
@@ -262,7 +262,7 @@ func TestWriteAndLoadDNSMagSequence_MultiDataset(t *testing.T) {
 	tmpFile.Close()
 
 	// Load the two datasets from the single file
-	seq := NewDatasetSequence(100, &dataset1.Date.Time)
+	seq := NewDatasetSequence(100, &dataset1.Date.Time, false, nil)
 	err = seq.LoadDNSMagSequenceFromReaderFile(tmpFile.Name())
 	if err != nil {
 		t.Fatalf("LoadDNSMagSequenceFromReaderFile failed: %v", err)
@@ -318,7 +318,7 @@ func TestLoadDNSMagSequenceFromReader_ExtraBytes(t *testing.T) {
 	extra := []byte("EXTRA BYTES")
 	input := append(cborData, extra...)
 
-	seq := NewDatasetSequence(100, &dataset.Date.Time)
+	seq := NewDatasetSequence(100, &dataset.Date.Time, false, nil)
 	err = seq.LoadDNSMagSequenceFromReader(
 		bytes.NewReader(input),
 		"testfile#%d",
@@ -347,7 +347,7 @@ func TestLoadDNSMagSequenceFromReader_IncompleteCBOR(t *testing.T) {
 	// Truncate the CBOR data to simulate incomplete input
 	truncated := cborData[:len(cborData)-1]
 
-	seq := NewDatasetSequence(100, &dataset.Date.Time)
+	seq := NewDatasetSequence(100, &dataset.Date.Time, false, nil)
 	err = seq.LoadDNSMagSequenceFromReader(
 		bytes.NewReader(truncated),
 		"testfile#%d",